feat: Update environment configuration and integrate Google Generative AI for document processing

.env

@@ -6,3 +6,4 @@ CHUNK_SIZE=1000
 CHUNK_OVERLAP=100
 RATE_LIMIT=20/minute
 EMBEDDING_MODEL=all-MiniLM-L6-v2
+HUGGINGFACEHUB_API_TOKEN=your_hf_token_here

.env.example

@@ -1,5 +1,4 @@
-HUGGINGFACEHUB_API_TOKEN=your_hf_token_here
-HF_MODEL=mistralai/Mistral-7B-Instruct-v0.2
+GOOGLE_API_KEY=your_gemini_api_key_here
 CHUNK_SIZE=800
 CHUNK_OVERLAP=120
 TOP_K_DEFAULT=4

app/config.py

@@ -13,5 +13,4 @@ CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "800"))
 CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "120"))
 TOP_K_DEFAULT = int(os.getenv("TOP_K_DEFAULT", "4"))
 
-HF_MODEL = os.getenv("HF_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
-HF_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN", "")
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")

app/rag.py

@@ -7,7 +7,7 @@ from typing import List, Dict
 
 import faiss
 import numpy as np
-import requests
+import google.generativeai as genai
 from sentence_transformers import SentenceTransformer
 from PyPDF2 import PdfReader
 
@@ -18,8 +18,7 @@ from .config import (
     DATA_DIR,
     CHUNK_SIZE,
     CHUNK_OVERLAP,
-    HF_MODEL,
-    HF_API_TOKEN,
+    GOOGLE_API_KEY,
 )
 
 _index_lock = Lock()
@@ -166,9 +165,9 @@ def generate_answer(question: str, contexts: List[Dict]) -> str:
     context_text = "\n\n".join(
         f"[Source: {c['source']}] {c['chunk']}" for c in contexts
     )
-    if not HF_API_TOKEN:
+    if not GOOGLE_API_KEY:
         return (
-            "No HF token configured. Here are the most relevant excerpts:\n\n"
+            "No Gemini API key configured. Here are the most relevant excerpts:\n\n"
             + context_text
         )
 
@@ -178,28 +177,16 @@ def generate_answer(question: str, contexts: List[Dict]) -> str:
         f"Context:\n{context_text}\n\nQuestion: {question}\nAnswer:"
     )
 
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 256,
-            "temperature": 0.2,
-            "return_full_text": False,
-        },
-    }
-    url = f"https://api-inference.huggingface.co/models/{HF_MODEL}"
-    response = requests.post(url, headers=headers, json=payload, timeout=60)
-    if response.status_code != 200:
+    try:
+        genai.configure(api_key=GOOGLE_API_KEY)
+        model = genai.GenerativeModel("gemini-1.5-flash")
+        response = model.generate_content(prompt)
+        return response.text.strip() or context_text
+    except Exception as e:
         return (
-            "LLM generation failed. Here are the most relevant excerpts:\n\n"
+            f"LLM generation failed: {str(e)}. Here are the most relevant excerpts:\n\n"
             + context_text
         )
-    data = response.json()
-    if isinstance(data, list) and data:
-        return data[0].get("generated_text", "").strip() or context_text
-    if isinstance(data, dict) and "generated_text" in data:
-        return data["generated_text"].strip() or context_text
-    return context_text
 
 
 def ingest_document(path: Path, doc_id: str, source: str) -> Dict:

requirements.txt

@@ -5,6 +5,6 @@ pydantic==2.8.2
 sentence-transformers==3.0.1
 faiss-cpu==1.8.0.post1
 PyPDF2==3.0.1
-requests==2.32.3
+google-generativeai>=0.3.0
 streamlit==1.37.1
 python-multipart==0.0.9