Upload main.py

main.py

@@ -0,0 +1,113 @@
+# main.py (Final Version with Updated Response Structure)
+import os
+import chromadb
+from fastapi import FastAPI
+from pydantic import BaseModel
+from dotenv import load_dotenv
+from sentence_transformers import SentenceTransformer
+import tiktoken
+from groq import Groq
+from googleapiclient.discovery import build
+
+# --- INITIALIZATION ---
+load_dotenv()
+app = FastAPI()
+groq_client = Groq()
+
+google_api_key = os.getenv("GOOGLE_API_KEY")
+search_engine_id = os.getenv("SEARCH_ENGINE_ID")
+google_search_service = build("customsearch", "v1", developerKey=google_api_key)
+
+print("Loading embedding model: 'nomic-ai/nomic-embed-text-v1.5'...")
+model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)
+print("Model loaded.")
+
+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 ---
+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 = model.encode(request.query).tolist()
+        results = collection.query(
+            query_embeddings=[query_embedding],
+            n_results=4,
+            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. If not, fallback to Google Custom Search
+            print(f"INFO: No relevant results found locally. Falling back to Google Search.")
+            search_results = google_search_service.cse().list(
+                q=request.query, cx=search_engine_id, num=4
+            ).execute()
+            
+            if not search_results.get('items'):
+                return {"answer": "I could not find any relevant information to answer your question.", "citations": []}
+            
+            context_chunks = [item.get('snippet', '') for item in search_results.get('items', [])]
+            citations = [item.get('link', '') for item in search_results.get('items', [])]
+
+        prompt_template = """
+        You are an expert legal AI assistant. Based ONLY on the following legal context, provide a detailed answer to the user's question. Structure your answer in multiple paragraphs and use bullet points for key points. Each bullet should be clear and informative. Do not use any outside knowledge.
+
+        CONTEXT:
+        {context}
+
+        QUESTION:
+        {query}
+
+        DETAILED ANSWER (use paragraphs and bullet points):
+        """
+
+        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": "I found some information, but it was too long to process. Please try a more specific query.", "citations": []}
+
+        final_prompt = prompt_template.format(context=context, query=request.query)
+
+        # 4. Generate answer using the chosen context
+        completion = groq_client.chat.completions.create(
+            model="openai/gpt-oss-120b",
+            messages=[{"role": "user", "content": final_prompt}],
+            temperature=0.2
+        )
+        answer = completion.choices[0].message.content
+
+        # --- THIS IS THE UPDATED RETURN STATEMENT ---
+        # Ensure only 4 unique citations are returned
+        citations = list(dict.fromkeys(citations))[:4]
+        return { "answer": answer, "citations": citations }
+
+    except Exception as e:
+        print(f"Error during query processing: {e}")
+        return {"error": "Failed to process the request."}