Update app.py

app.py

@@ -5,7 +5,7 @@ import chromadb
 import tiktoken
 from groq import Groq
 from googleapiclient.discovery import build
-from nomic import AtlasClient  # fixed import
+from sentence_transformers import SentenceTransformer
 
 # --- INITIALIZATION ---
 app = FastAPI()
@@ -14,95 +14,44 @@ groq_client = Groq()
 # Environment variables
 google_api_key = os.getenv("GOOGLE_API_KEY")
 search_engine_id = os.getenv("SEARCH_ENGINE_ID")
-NOMIC_API_KEY = os.getenv("NOMIC_API_KEY")
 
 # Google Custom Search setup
 google_search_service = build("customsearch", "v1", developerKey=google_api_key)
 
-# Nomic Atlas client initialization
-atlas = AtlasClient(api_key=NOMIC_API_KEY)
-print("Connected to Nomic Atlas API.")
+# SentenceTransformer model for embeddings (runs locally in HF Space)
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+print("Loaded SentenceTransformer embeddings.")
 
 # Connect to local ChromaDB
 client = chromadb.PersistentClient(path="./chroma_db")
 collection = client.get_collection(name="legal_docs")
-print(f"Connected to ChromaDB. Documents in collection: {collection.count()}")
 
-# --- DATA MODELS ---
+# --- REQUEST MODEL ---
 class QueryRequest(BaseModel):
     query: str
 
-# --- TOKENIZER FUNCTION ---
-def count_tokens(text, model="gpt-4"):
-    encoding = tiktoken.encoding_for_model(model)
-    return len(encoding.encode(text))
-
-# --- API ENDPOINT ---
-@app.post("/api/ai/research/query")
-async def research_query(request: QueryRequest):
-    try:
-        context = ""
-        citations = []
-
-        # 1. Search local ChromaDB first
-        query_embedding = atlas.embed_text(request.query)  # Nomic embedding
-        results = collection.query(
-            query_embeddings=[query_embedding],
-            n_results=3,
-            include=['documents', 'metadatas', 'distances']
-        )
-
-        # 2. Check if the best result is relevant enough
-        SIMILARITY_THRESHOLD = 0.7
-        if results['distances'] and results['distances'][0] and results['distances'][0][0] < SIMILARITY_THRESHOLD:
-            print("INFO: Found relevant documents in local ChromaDB.")
-            context_chunks = results['documents'][0]
-            citations = [meta['source'] for meta in results['metadatas'][0]]
-        else:
-            # 3. Fallback to Google Custom Search
-            print("INFO: No relevant results locally. Using Google Search.")
-            search_results = google_search_service.cse().list(
-                q=request.query, cx=search_engine_id, num=3
-            ).execute()
-            
-            if not search_results.get('items'):
-                return {"answer": "I could not find any relevant information.", "citations": []}
-            
-            context_chunks = [item.get('snippet', '') for item in search_results.get('items', [])]
-            citations = [item.get('link', '') for item in search_results.get('items', [])]
-
-        # 4. Build context respecting token limits
-        prompt_template = """
-        You are an expert legal AI assistant. Based ONLY on the following legal context, provide a concise answer.
-        CONTEXT: {context}
-        QUESTION: {query}
-        CONCISE ANSWER:
-        """
-        TOKEN_LIMIT = 7000
-        for chunk in context_chunks:
-            temp_context = context + chunk + "\n\n---\n\n"
-            temp_prompt = prompt_template.format(context=temp_context, query=request.query)
-            if count_tokens(temp_prompt) <= TOKEN_LIMIT:
-                context = temp_context
-            else:
-                break
-        
-        if not context:
-            return {"answer": "Information found is too long to process. Try a more specific query.", "citations": []}
-
-        final_prompt = prompt_template.format(context=context, query=request.query)
-
-        # 5. Generate answer with Groq AI
-        completion = groq_client.chat.completions.create(
-            model="openai/gpt-oss-120b",
-            messages=[{"role": "user", "content": final_prompt}],
-            temperature=0.2
+# --- API ROUTES ---
+@app.post("/query")
+async def query_api(request: QueryRequest):
+    # Create embedding for query
+    query_embedding = embedder.encode(request.query).tolist()
+
+    # Search in ChromaDB
+    results = collection.query(query_embeddings=[query_embedding], n_results=3)
+
+    if results and results["documents"]:
+        context_docs = [doc for sublist in results["documents"] for doc in sublist]
+        response_text = f"Relevant documents found:\n{context_docs}"
+    else:
+        # fallback to Google search
+        response_text = "No relevant local docs found. Searching externally..."
+        google_results = (
+            google_search_service.cse()
+            .list(q=request.query, cx=search_engine_id, num=3)
+            .execute()
         )
-        answer = completion.choices[0].message.content
-
-        # 6. Return answer with citations
-        return {"answer": answer, "citations": list(set(citations))}
+        response_text += "\n\nExternal sources:\n"
+        for item in google_results.get("items", []):
+            response_text += f"- {item['title']}: {item['link']}\n"
 
-    except Exception as e:
-        print(f"Error during query processing: {e}")
-        return {"error": "Failed to process the request."}
+    return {"answer": response_text}