Update agent_langchain.py

agent_langchain.py

@@ -88,64 +88,84 @@ def call_routing(text, retries=3, delay=5):
                 return "General IT"
 
 # ==============================================================
-# 📚 KNOWLEDGE BASE SETUP
+# 📚 KNOWLEDGE BASE SETUP (PATCHED)
 # ==============================================================
-# Persistent Chroma client (new API)
-CHROMA_PATH = "/data/chroma"
-COLLECTION_NAME = "knowledge_base"
+import threading
 
-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)
+CHROMA_PATH = "/data/chroma"  # Persistent path in Space
+COLLECTION_NAME = "knowledge_base"
 
-# Sentence Transformer Embedding Function
+kb_collection = None  # Global placeholder, initialized via /setup
 embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
     model_name="sentence-transformers/all-MiniLM-L6-v2"
 )
+kb_lock = threading.Lock()  # Ensure thread-safe updates
+
+def initialize_kb(json_texts: list, metadatas: list):
+    """
+    Initialize or populate the knowledge base with embeddings.
+    This is called from /setup endpoint.
+    """
+    global kb_collection
+    with kb_lock:
+        try:
+            # Create persistent client (only now)
+            chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+            kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+        except Exception as e:
+            print("⚠️ Could not initialize persistent KB, using in-memory fallback:", e)
+            chroma_client = chromadb.Client()  # fallback in-memory
+            kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+
+        # Clear existing entries safely
+        try:
+            if kb_collection.count() > 0:
+                # Use a safe delete by fetching all ids first
+                ids_to_delete = kb_collection.get()['ids']
+                kb_collection.delete(ids=ids_to_delete)
+        except Exception as e:
+            print("⚠️ Could not clear KB:", e)
+
+        # Generate embeddings
+        try:
+            vectors = embedding_func(json_texts)
+            kb_collection.add(
+                ids=[str(i) for i in range(len(json_texts))],
+                embeddings=vectors,
+                metadatas=metadatas,
+                documents=json_texts
+            )
+            print(f"✅ KB initialized with {len(json_texts)} entries")
+        except Exception as e:
+            print("⚠️ Failed to populate KB:", e)
 
-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}
+    Query the knowledge base for relevant solutions.
+    Returns {"answer": str, "confidence": float, "metadata": dict}
     """
+    global kb_collection
     if not kb_collection:
-        return {"answer": "⚠️ Knowledge Base not initialized.", "confidence": 0.0, "metadata": {}}
+        return {"answer": "⚠️ KB not initialized. Please call /setup first.", "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
-        }
+        confidence = float(np.dot(query_vector, best_vec) / (np.linalg.norm(query_vector) * np.linalg.norm(best_vec) + 1e-8))
 
+        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": {}}
+        print("⚠️ KB query failed:", e)
+        return {"answer": "Error querying KB.", "confidence": 0.0, "metadata": {}}
+
 
 # ==============================================================
 # 🧠 GEMINI LLM (Official LangChain Integration)