Update app.py

app.py

@@ -1,14 +1,28 @@
 import os
+
+# SET CACHE PATHS BEFORE ANY IMPORTS
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers"
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/sentence_transformers"
+os.environ["TORCH_HOME"] = "/tmp/torch"
+
 import json
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from pydantic import BaseModel
-from agent_langchain import process_ticket_langchain, classify_ticket, call_routing, kb_collection
 import chromadb
 from chromadb.config import Settings
-from chromadb.api.models import Collection
-from chromadb.utils import embedding_functions
 from sentence_transformers import SentenceTransformer
+import numpy as np
 
+# Import from agent_langchain
+from agent_langchain import (
+    process_ticket_langchain, 
+    classify_ticket, 
+    call_routing,
+    get_kb_collection,
+    encoder
+)
 
 app = FastAPI(title="Smart Helpdesk AI Agent LangChain")
 
@@ -19,16 +33,11 @@ class TicketRequest(BaseModel):
     text: str
     user_email: str = None
 
-class SetupRequest(BaseModel):
-    kb_file: str  # path to KB.json
-
 # -------------------------------
 # Persistent Chroma client
 # -------------------------------
 CHROMA_PATH = "/tmp/chroma"
 COLLECTION_NAME = "knowledge_base"
-# Global variable for the running app
-kb_collection = None
 
 # -------------------------------
 # KB Setup Endpoint
@@ -39,8 +48,6 @@ 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:
         # Load JSON from uploaded file
         content_bytes = await kb_file.read()
@@ -51,9 +58,14 @@ async def setup_kb(kb_file: UploadFile = File(...)):
 
         print(f"📘 Loaded {len(data)} items from {kb_file.filename}")
 
-        # Initialize encoder and Chroma
-        encoder = SentenceTransformer("all-MiniLM-L6-v2")
-        chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
+        # Get or create collection using shared function
+        chroma_client = chromadb.PersistentClient(
+            path=CHROMA_PATH,
+            settings=Settings(
+                anonymized_telemetry=False,
+                allow_reset=True
+            )
+        )
         collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
 
         # Clear existing records
@@ -64,20 +76,36 @@ async def setup_kb(kb_file: UploadFile = File(...)):
         # Prepare texts, ids, and metadata
         texts, ids, metadatas = [], [], []
         for i, item in enumerate(data):
-            text = item.get("text") or ""
+            text = item.get("text") or item.get("content") or ""
             item_id = item.get("id") or str(i)
+            
+            if not text:
+                print(f"⚠️ Skipping item {i} - no text content")
+                continue
+                
             texts.append(text)
             ids.append(str(item_id))
-            metadatas.append({"id": str(item_id)})
+            metadatas.append({"id": str(item_id), "original_index": i})
+
+        if not texts:
+            raise HTTPException(status_code=400, detail="No valid text content found in JSON.")
 
-        # Generate embeddings
+        # Generate embeddings using the shared encoder
         print("🧠 Generating embeddings...")
         embeddings = encoder.encode(texts, show_progress_bar=True).tolist()
 
         # Add to ChromaDB
         print("💾 Adding to ChromaDB...")
-        collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas)
-        kb_collection = collection  # assign to global
+        collection.add(
+            ids=ids, 
+            embeddings=embeddings, 
+            documents=texts,
+            metadatas=metadatas
+        )
+
+        # Update the global reference in agent_langchain
+        import agent_langchain
+        agent_langchain.kb_collection = collection
 
         print(f"✅ Successfully added {collection.count()} records to {COLLECTION_NAME}.")
         return {"message": "Knowledge base successfully initialized.", "count": collection.count()}
@@ -85,7 +113,9 @@ async def setup_kb(kb_file: UploadFile = File(...)):
     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}")
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Setup failed: {str(e)}")
 
 # -------------------------------
 # Step-by-Step Endpoints
@@ -106,37 +136,47 @@ async def route_endpoint(ticket: TicketRequest):
 @app.post("/kb_query")
 async def kb_query_endpoint(ticket: TicketRequest):
     """Query the flattened KB directly using embeddings and return the best match."""
-    global kb_collection
-    if not kb_collection:
+    collection = get_kb_collection()
+    
+    if not collection:
         raise HTTPException(status_code=400, detail="KB not set up. Call /setup first.")
 
     try:
-        # Use the same SentenceTransformer as in /setup
-        encoder = SentenceTransformer("all-MiniLM-L6-v2")
-        query_vec = encoder.encode([ticket.text])[0]
+        # Check if collection has data
+        count = collection.count()
+        if count == 0:
+            raise HTTPException(status_code=400, detail="KB is empty. Please upload data via /setup.")
+
+        # Encode query using the shared encoder
+        query_embedding = encoder.encode([ticket.text])[0].tolist()
 
         # Query ChromaDB
-        result = kb_collection.query(
-            query_embeddings=[query_vec],
+        result = collection.query(
+            query_embeddings=[query_embedding],
             n_results=1,
             include=["documents", "distances", "metadatas"]
         )
 
-        if not result or len(result['documents'][0]) == 0:
+        if not result or not result.get('documents') 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
+        best_distance = result['distances'][0][0] if result.get('distances') else 1.0
+        
+        # Convert L2 distance to confidence score
+        # For normalized embeddings, L2 distance ranges from 0 (identical) to ~2.0 (opposite)
+        confidence = max(0.0, 1.0 - (best_distance / 2.0))
 
         return {
             "answer": best_doc,
-            "confidence": round(confidence, 3)
+            "confidence": round(float(confidence), 3)
         }
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"KB query failed: {e}")
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"KB query failed: {str(e)}")
 
 # -------------------------------
 # Full Ticket Orchestration
@@ -144,24 +184,40 @@ async def kb_query_endpoint(ticket: TicketRequest):
 @app.post("/orchestrate")
 async def orchestrate_endpoint(ticket: TicketRequest):
     """Full ticket orchestration via LangChain agent with nicely formatted reasoning trace"""
-    result = process_ticket_langchain(ticket.text)
-
-    # Format reasoning trace for readability
-    formatted_trace = [{"step": idx + 1, "description": line} for idx, line in enumerate(result.get("reasoning_trace", []))]
-
-    response = {
-        "status": result["status"],
-        "classification": result["classification"],
-        "department": result["department"],
-        "answer": result["answer"],
-        "reasoning_trace": formatted_trace
-    }
+    try:
+        result = process_ticket_langchain(ticket.text)
+
+        # Format reasoning trace for readability
+        formatted_trace = [
+            {"step": idx + 1, "description": line} 
+            for idx, line in enumerate(result.get("reasoning_trace", []))
+        ]
+
+        response = {
+            "status": result["status"],
+            "classification": result["classification"],
+            "department": result["department"],
+            "answer": result["answer"],
+            "reasoning_trace": formatted_trace
+        }
 
-    return response
+        return response
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Orchestration failed: {str(e)}")
 
 # -------------------------------
 # Health Check
 # -------------------------------
 @app.get("/health")
 async def health():
-    return {"status": "ok"}
+    collection = get_kb_collection()
+    kb_status = "initialized" if collection and collection.count() > 0 else "not initialized"
+    kb_count = collection.count() if collection else 0
+    
+    return {
+        "status": "ok",
+        "kb_status": kb_status,
+        "kb_records": kb_count
+    }