Update agent_langchain.py

agent_langchain.py

@@ -3,7 +3,10 @@ import requests
 import torch
 import time
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import numpy as np
+
 import chromadb
+from chromadb.utils import embedding_functions
 from chromadb.config import Settings
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.agents import initialize_agent, Tool, AgentType
@@ -87,30 +90,61 @@ def call_routing(text, retries=3, delay=5):
 # 📚 KNOWLEDGE BASE SETUP
 # ==============================================================
 # Persistent Chroma client (new API)
-chroma_client = chromadb.PersistentClient(path="/tmp/chroma")
-
+CHROMA_PATH = "/data/chroma"
 COLLECTION_NAME = "knowledge_base"
 
 try:
+    chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
     kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
 except Exception as e:
     kb_collection = None
     print("⚠️ Could not initialize KB:", e)
 
-def query_kb(text, top_k=1):
-    """Query the knowledge base for relevant solutions."""
-    if not kb_collection:
-        return {"answer": "⚠️ KB not set up.", "confidence": 0.0}
-
-    results = kb_collection.query(query_texts=[text], n_results=top_k)
-    if not results or not results.get("documents") or len(results["documents"][0]) == 0:
-        return {"answer": "No relevant KB found.", "confidence": 0.0}
+# Sentence Transformer Embedding Function
+embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
+    model_name="sentence-transformers/all-MiniLM-L6-v2"
+)
 
-    return {
-        "answer": results["documents"][0][0],
-        "confidence": results.get("distances", [[0.0]])[0][0],
-        "metadata": results.get("metadatas", [[{}]])[0][0]
-    }
+def cosine_similarity(vec1, vec2):
+    """Compute cosine similarity between two vectors."""
+    if np.linalg.norm(vec1) == 0 or np.linalg.norm(vec2) == 0:
+        return 0.0
+    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
+
+def query_kb(text: str, top_k: int = 1):
+    """
+    Query the knowledge base for relevant solutions using Chroma.
+    Returns: {"answer": str, "confidence": float, "metadata": dict}
+    """
+    if not kb_collection:
+        return {"answer": "⚠️ Knowledge Base not initialized.", "confidence": 0.0, "metadata": {}}
+
+    try:
+        # Embed the query
+        query_vector = embedding_func([text])[0]
+
+        # Query Chroma
+        results = kb_collection.query(query_embeddings=[query_vector], n_results=top_k)
+        if not results or not results.get("documents") or len(results["documents"][0]) == 0:
+            return {"answer": "No relevant KB entry found.", "confidence": 0.0, "metadata": {}}
+
+        # Extract best match
+        best_doc = results["documents"][0][0]
+        best_meta = results["metadatas"][0][0]
+        best_vec = np.array(results["embeddings"][0][0]) if "embeddings" in results else query_vector
+
+        # Compute similarity confidence manually
+        confidence = float(cosine_similarity(query_vector, best_vec))
+
+        return {
+            "answer": best_doc,
+            "confidence": round(confidence, 3),
+            "metadata": best_meta
+        }
+
+    except Exception as e:
+        print(f"⚠️ KB query failed: {e}")
+        return {"answer": "Error accessing KB.", "confidence": 0.0, "metadata": {}}
 
 # ==============================================================
 # 🧠 GEMINI LLM (Official LangChain Integration)