Integrate Gemini API for embedding and fact extraction, replacing local model dependencies

app/core/embedding.py

@@ -1,23 +1,25 @@
 import os
-from app.core.device_setup import device
-from sentence_transformers import SentenceTransformer
 
-cache_dir = os.getenv("HF_HOME", "/app/hf_cache")
-model_name = 'sentence-transformers/all-MiniLM-L6-v2'
+# Replace local embedding model with Gemini API integration
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 
-# Load the model (this model will be downloaded the first time if it's not cached)
-embedding_model = SentenceTransformer(model_name, cache_folder=cache_dir)
-
-# Move the model to the appropriate device (GPU if available, otherwise CPU)
-embedding_model = embedding_model.to(device)
+if not GEMINI_API_KEY:
+    raise ValueError("❌ Gemini API Key (GEMINI_API_KEY) has not been set yet")
 
 def generate_embedding(texts: list):
+    import requests
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "contents": [{"parts": [{"text": text}]} for text in texts]
+    }
+
     try:
-        # Use the model's encode method to get embeddings for the input text
-        embedding = embedding_model.encode(texts, convert_to_tensor=True)  # Automatically moves tensor to the correct device
-        embedding_cpu = embedding.cpu()  # Move to CPU if needed
-        # Convert to list for easier handling if needed
-        return embedding_cpu.numpy().tolist()
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        # Parse embeddings from the response (adjust based on Gemini API's actual response structure)
+        return [item.get("embedding", []) for item in data.get("contents", [])]
     except Exception as e:
-        print(f"Error generating embedding: {e}")
+        print(f"Error generating embedding via Gemini API: {e}")
         return None

app/core/fact_extraction.py

@@ -1,35 +1,44 @@
 import os
 from app.core.device_setup import device
-from transformers import pipeline
 from app.core.fact_management import save_user_fact
 from app.core.logging_setup import logger
 
-model_name = 'dslim/distilbert-NER'
-cache_dir = os.getenv("HF_HOME", "/app/hf_cache")
+# Replace local model loading with Gemini API integration
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 
-# Ensure the device is used for the pipeline
-pipeline_device = 0 if device == "cuda" else -1
-nlp = pipeline("token-classification", model=model_name, device=pipeline_device, model_kwargs={"cache_dir": cache_dir})
+if not GEMINI_API_KEY:
+    raise ValueError("❌ Gemini API Key (GEMINI_API_KEY) has not been set yet")
 
-def extract_name(text):
-    """Extract name from text using Transformers."""
-    entities = nlp(text)
-    names = [entity['word'] for entity in entities if entity['entity'].startswith('B-PER')]
-    return names
+def query_gemini_for_entities(text: str):
+    import requests
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "contents": [{"parts": [{"text": text}]}]
+    }
 
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        # Parse entities from the response (adjust based on Gemini API's actual response structure)
+        return data.get("entities", [])
+    except Exception as e:
+        logger.error(f"🚨 Error querying Gemini API for entities: {e}")
+        return []
+
+# Replace the NLP pipeline with Gemini API calls
 def extract_and_store_facts(message):
-    """Extract personal facts like name, location, and interests."""
-    entities = nlp(message)
-    
+    entities = query_gemini_for_entities(message)
+
     # Extract name
-    name = next((entity['word'] for entity in entities if entity['entity'].startswith('B-PER')), None)
+    name = next((entity['name'] for entity in entities if entity.get('type') == 'PERSON'), None)
     if name:
-        clean_name = name.split(".")[0]  # Store only the first sentence
-        save_user_fact("name", clean_name)
+        save_user_fact("name", name)
         logger.info(f"User name '{name}' stored in memory.")
-    
+
     # Extract location
-    location = next((entity['word'] for entity in entities if entity['entity'].startswith('B-LOC')), None)
+    location = next((entity['name'] for entity in entities if entity.get('type') == 'LOCATION'), None)
     if location:
         save_user_fact("location", location)
         logger.info(f"User location '{location}' stored in memory.")

app/routes/chat_hf.py

@@ -15,7 +15,6 @@ from app.core.logging_setup import logger
 from app.core.prompts import SYSTEM_PROMPT
 from app.core.interaction_trends import get_time_of_day
 from app.core.search_utils import needs_web_search, search_duckduckgo
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 import asyncio
 
@@ -30,113 +29,36 @@ headers = {
     "Authorization": f"Bearer {HUGGINGFACE_TOKEN}"
 }
 
-# Load the model and tokenizer locally
-cache_dir = os.getenv("HF_HOME", "/app/hf_cache")
-model_name = "google/gemma-3-1b-it"
-tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir).to(device)
-model.config.pad_token_id = tokenizer.eos_token_id
-
-# Check model and tokenizer types
-logger.info(f"Model type: {type(model)}")
-logger.info(f"Tokenizer type: {type(tokenizer)}")
-logger.info("Model and tokenizer loaded successfully.")
-
-def build_clean_prompt(messages):
-    role_map = {
-        "system": "System",
-        "user": "User",
-        "assistant": "Arina"
-    }
-    prompt = ""
-    for msg in messages:
-        role = role_map.get(msg["role"], "User")
-        prompt += f"{role}: {msg['content'].strip()}\n"
-    prompt += "Arina:"
-    return prompt
-
-def generate_response(prompt_text):
-    try:
-        logger.info("Starting to generate response.")
-        logger.info(f"Original prompt text: {prompt_text}")
-
-        # Sanity check for prompt structure
-        if "User:" in prompt_text[-80:] or prompt_text.count("User:") > prompt_text.count("Arina:"):
-            logger.warning("⚠️ Possible misalignment in role markers. Last prompt may confuse model.")
-
-        # Tokenize the prompt
-        logger.info("Tokenizing the prompt...")
-        logger.info(f"Tokenizer: {tokenizer}")
-        model_inputs = tokenizer(
-            prompt_text.strip(),
-            return_tensors="pt",
-            truncation=True,
-            max_length=1024  # Can be increased based on your model's context window
-        )
-        logger.info("Prompt tokenized.")
-
-        # Log input token length
-        input_len = model_inputs["input_ids"].shape[-1]
-        logger.info(f"🧾 Prompt token length: {input_len}")
-
-        assert prompt_text.count("User:") == prompt_text.count("Arina:"), "⚠️ Prompt imbalance may confuse model"
-
-        # Generate response
-        logger.info("Generating model response...")
-        logger.info(f"Model: {model}")
-        model_outputs = model.generate(
-            **model_inputs,
-            max_new_tokens=512,  # Can be adjusted, output token limit
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.7,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
-        logger.info("Model response generated.")
-
-        # Decode output
-        logger.info("Decoding model output...")
-        full_output = tokenizer.decode(model_outputs[0], skip_special_tokens=True).strip()
-        logger.info("Model output decoded.")
+# Replace local model loading with Gemini API integration
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
 
-        # Extract only the part after "Arina:" (if present)
-        if "Arina:" in full_output:
-            response = full_output.split("Arina:", 1)[-1].strip()
-        else:
-            response = full_output.strip()
-
-        # Prevent output starting with "User:" as it is hallucination
-        if response.startswith("User:"):
-            logger.warning("⚠️ Model hallucinated user input.")
-            response = response.split("Arina:")[-1].strip() if "Arina:" in response else response
-
-        # Clean echo if present
-        if response.startswith(prompt_text):
-            response = response[len(prompt_text):].strip()
-
-        # Fallback if empty
-        if not response:
-            logger.warning("⚠️ Empty response generated. Returning fallback.")
-            response = "I'm not sure how to respond to that, but I'm here to help."
+if not GEMINI_API_KEY:
+    raise ValueError("❌ Gemini API Key (GEMINI_API_KEY) has not been set yet")
 
-        logger.info(f"✅ Final Arina response: {response}")
-        return response
+def query_gemini_api(prompt: str) -> str:
+    import requests
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "contents": [{"parts": [{"text": prompt}]}]
+    }
 
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        return data.get("contents", [{}])[0].get("parts", [{}])[0].get("text", "")
     except Exception as e:
-        logger.error(f"🚨 Unexpected error in generate_response: {e}")
-        return "❌ An unexpected error occurred while generating a response."
+        logger.error(f"🚨 Error querying Gemini API: {e}")
+        return "⚠️ An error occurred while generating a response."
 
-# Test the generate_response function
-logger.info(f"generate_response is: {type(generate_response)}")
+# Replace generate_response and query_huggingface with query_gemini_api
+def generate_response(prompt_text):
+    return query_gemini_api(prompt_text)
 
-# Ensure query_huggingface is not shadowing generate_response
+# Ensure query_huggingface uses the Gemini API
 def query_huggingface(prompt: str) -> str:
-    logger.debug(f"Calling generate_response with prompt: {prompt}")
-    response = generate_response(prompt)
-    return response
+    return query_gemini_api(prompt)
 
 router = APIRouter()