Update app.py

app.py

@@ -35,53 +35,55 @@ kb_collection = None
 # -------------------------------
 @app.post("/setup")
 async def setup_kb(kb_file: UploadFile = File(...)):
+    """
+    Uploads a JSON KB file (flattened), generates embeddings with SentenceTransformer,
+    and populates a persistent ChromaDB collection.
+    """
     global kb_collection
+
     try:
-        # Read JSON directly
+        # Load JSON from uploaded file
         content_bytes = await kb_file.read()
         data = json.loads(content_bytes)
 
         if not isinstance(data, list):
-            raise HTTPException(status_code=400, detail="JSON must contain a list of knowledge items.")
+            raise HTTPException(status_code=400, detail="JSON must be a list of items.")
 
-        print(f"📘 Loaded {len(data)} items from uploaded file '{kb_file.filename}'")
+        print(f"📘 Loaded {len(data)} items from {kb_file.filename}")
 
-        # Initialize Sentence Transformer
+        # Initialize encoder and Chroma
         encoder = SentenceTransformer("all-MiniLM-L6-v2")
+        chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+        collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
 
-        # Initialize ChromaDB (persistent)
-        chroma_client = chromadb.PersistentClient(path="/tmp/chroma")  # must be writable
-        kb_collection = chroma_client.get_or_create_collection(
-            name="knowledge_base",
-            embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
-                model_name="all-MiniLM-L6-v2"
-            )
-        )
-
-        # Clear old data
-        if kb_collection.count() > 0:
-            kb_collection.delete(ids=kb_collection.get()['ids'])
+        # Clear existing records
+        if collection.count() > 0:
+            print(f"🧹 Clearing {collection.count()} existing records...")
+            collection.delete(ids=collection.get()['ids'])
 
-        # Prepare data
+        # Prepare texts, ids, and metadata
         texts, ids, metadatas = [], [], []
         for i, item in enumerate(data):
-            content = item.get("answer") or item.get("content") or ""
-            title = item.get("id") or f"Document {i+1}"
-            texts.append(f"{title}. {content}")
-            ids.append(str(i))
-            # Only keep str/int/bool/float for metadata
-            metadata = {k: str(v) if isinstance(v, list) else v for k, v in item.items()}
-            metadatas.append(metadata)
+            text = item.get("text") or ""
+            item_id = item.get("id") or str(i)
+            texts.append(text)
+            ids.append(str(item_id))
+            metadatas.append({"id": str(item_id)})
 
         # Generate embeddings
+        print("🧠 Generating embeddings...")
         embeddings = encoder.encode(texts, show_progress_bar=True).tolist()
 
-        # Add to Chroma
-        kb_collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas)
-        print(f"✅ Successfully added {kb_collection.count()} records.")
+        # Add to ChromaDB
+        print("💾 Adding to ChromaDB...")
+        collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas)
+        kb_collection = collection  # assign to global
 
-        return {"message": "Knowledge base successfully initialized.", "count": kb_collection.count()}
+        print(f"✅ Successfully added {collection.count()} records to {COLLECTION_NAME}.")
+        return {"message": "Knowledge base successfully initialized.", "count": collection.count()}
 
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid JSON file.")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Setup failed: {e}")
 
@@ -103,13 +105,38 @@ async def route_endpoint(ticket: TicketRequest):
 
 @app.post("/kb_query")
 async def kb_query_endpoint(ticket: TicketRequest):
-    """Query KB directly"""
+    """Query the flattened KB directly using embeddings and return the best match."""
+    global kb_collection
     if not kb_collection:
         raise HTTPException(status_code=400, detail="KB not set up. Call /setup first.")
-    result = kb_collection.query(query_texts=[ticket.text], n_results=1)
-    if not result or len(result['documents'][0]) == 0:
-        return {"answer": "No relevant KB found."}
-    return {"answer": result['documents'][0][0], "confidence": result['distances'][0][0] if result.get('distances') else 0.0}
+
+    try:
+        # Use the same SentenceTransformer as in /setup
+        encoder = SentenceTransformer("all-MiniLM-L6-v2")
+        query_vec = encoder.encode([ticket.text])[0]
+
+        # Query ChromaDB
+        result = kb_collection.query(
+            query_embeddings=[query_vec],
+            n_results=1,
+            include=["documents", "distances", "metadatas"]
+        )
+
+        if not result or len(result['documents'][0]) == 0:
+            return {"answer": "No relevant KB found.", "confidence": 0.0}
+
+        # Extract best match
+        best_doc = result['documents'][0][0]
+        best_distance = result['distances'][0][0] if result.get('distances') else 0.0
+        confidence = float(1 - best_distance)  # Convert distance → confidence
+
+        return {
+            "answer": best_doc,
+            "confidence": round(confidence, 3)
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"KB query failed: {e}")
 
 # -------------------------------
 # Full Ticket Orchestration