Update agent_langchain.py

agent_langchain.py

@@ -147,29 +147,37 @@ def initialize_kb(json_texts: list, metadatas: list):
 
 def query_kb(text: str, top_k: int = 1):
     """
-    Query the knowledge base for relevant solutions.
-    Returns {"answer": str, "confidence": float, "metadata": dict}
+    Query the Chroma knowledge base using SentenceTransformer embeddings.
+    Returns answer + confidence.
     """
     global kb_collection
-    if not kb_collection:
-        return {"answer": "⚠️ KB not initialized. Please call /setup first.", "confidence": 0.0, "metadata": {}}
+    if not kb_collection or kb_collection.count() == 0:
+        return {"answer": None, "confidence": 0.0}
 
-    try:
-        query_vector = embedding_func([text])[0]
-        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": {}}
+    # Embed the query manually using SentenceTransformer
+    encoder = SentenceTransformer("all-MiniLM-L6-v2")
+    query_embedding = encoder.encode([text])[0]
 
-        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
+    # Query Chroma by embeddings
+    results = kb_collection.query(
+        query_embeddings=[query_embedding],
+        n_results=top_k
+    )
+
+    if not results or not results.get("documents") or len(results["documents"][0]) == 0:
+        return {"answer": None, "confidence": 0.0}
+
+    # Extract top document and metadata
+    answer = results["documents"][0][0]
+    metadata = results["metadatas"][0][0]
 
-        confidence = float(np.dot(query_vector, best_vec) / (np.linalg.norm(query_vector) * np.linalg.norm(best_vec) + 1e-8))
+    # Compute cosine similarity manually for confidence
+    stored_embedding = np.array(results["embeddings"][0][0])
+    query_vec = np.array(query_embedding)
+    confidence = float(np.dot(query_vec, stored_embedding) /
+                       (np.linalg.norm(query_vec) * np.linalg.norm(stored_embedding)))
 
-        return {"answer": best_doc, "confidence": round(confidence, 3), "metadata": best_meta}
-    except Exception as e:
-        print("⚠️ KB query failed:", e)
-        return {"answer": "Error querying KB.", "confidence": 0.0, "metadata": {}}
+    return {"answer": answer, "confidence": round(confidence, 3), "metadata": metadata}
 
 
 # ==============================================================