Merge branch 'main' of github.com:sohank-17/Neon-AI-Project

multi_llm_chatbot_backend/app/api/routes.py

@@ -11,6 +11,9 @@ from pydantic import BaseModel
 from typing import Optional, List
 from fastapi import UploadFile, File
 from app.utils.document_extractor import extract_text_from_file
+from app.core.orchestrator import answer_with_persona_context
+from app.utils.chroma_client import add_persona_doc
+import hashlib
 from app.utils.file_limits import is_within_upload_limit
 
 router = APIRouter()
@@ -531,6 +534,22 @@ async def debug_personas():
         "current_provider": current_provider
     }
 
+class PersonaQuery(BaseModel):
+    question: str
+    persona: str
+
+@router.post("/ask/")
+async def ask_question(query: PersonaQuery):
+    response = await answer_with_persona_context(query.question, query.persona)
+    
+    # Store Q&A in vector DB
+    combined_text = f"Q: {query.question}\nA: {response}"
+    doc_id = hashlib.md5(combined_text.encode()).hexdigest()  # Create a unique doc ID
+    
+    add_persona_doc(combined_text, query.persona, doc_id)
+
+    return {"response": response}
+
 @router.get("/uploaded-files")
 def get_uploaded_filenames():
-    return {"files": session_context.uploaded_files}
+    return {"files": session_context.uploaded_files}

multi_llm_chatbot_backend/app/core/orchestrator.py

@@ -1,5 +1,8 @@
 from app.models.persona import Persona
 from typing import List
+from app.utils.chroma_client import query_persona_knowledge
+from app.llm.gemini_client import GeminiClient
+from app.llm.llm_client import LLMClient
 
 class ChatOrchestrator:
     def __init__(self):
@@ -25,3 +28,16 @@ class ChatOrchestrator:
             context.append({"role": persona.id, "content": reply})
 
         return responses
+    
+async def answer_with_persona_context(question: str, persona: str) -> str:
+    llm: LLMClient = GeminiClient()
+    context_chunks = query_persona_knowledge(question, persona)
+    context = "\n".join(context_chunks) if context_chunks else "No relevant info found."
+        
+    system_prompt = f"You are a helpful assistant responding as a {persona}."
+    message_context = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
+    ]
+
+    return await llm.generate(system_prompt, message_context)

multi_llm_chatbot_backend/app/llm/embedding_client.py

@@ -0,0 +1,11 @@
+from sentence_transformers import SentenceTransformer
+import os
+
+API_KEY = os.getenv("GEMINI_API_KEY")
+
+# Using a compact, fast model good for semantic search
+model = SentenceTransformer("all-MiniLM-L6-v2")
+
+def get_embedding(text: str) -> list[float]:
+    embedding = model.encode(text, convert_to_numpy=True)
+    return embedding.tolist()

multi_llm_chatbot_backend/app/utils/chroma_client.py

@@ -0,0 +1,25 @@
+import chromadb
+from chromadb.config import Settings
+from app.llm.embedding_client import get_embedding
+
+client = chromadb.PersistentClient(path="./chroma_storage")
+
+collection = client.get_or_create_collection("persona_knowledge")
+
+def add_persona_doc(text: str, persona: str, doc_id: str):
+    embedding = get_embedding(text)
+    collection.add(
+        documents=[text],
+        embeddings=[embedding],
+        metadatas=[{"persona": persona}],
+        ids=[doc_id]
+    )
+
+def query_persona_knowledge(query: str, persona: str):
+    query_embedding = get_embedding(query)
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=3,
+        where={"persona": persona}
+    )
+    return results['documents'][0] if results['documents'] else []