Update app.py

app.py

@@ -1,6 +1,4 @@
 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"
@@ -10,32 +8,31 @@ os.environ["TORCH_HOME"] = "/tmp/torch"
 import json
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from pydantic import BaseModel
+from typing import Optional
 import chromadb
 from chromadb.config import Settings
 from sentence_transformers import SentenceTransformer
-import numpy as np
 
-# Import from agent_langchain
+# Import from the true agent
 from agent_langchain import (
-    process_ticket_langchain, 
-    classify_ticket, 
+    process_with_agent,
+    get_conversation_history,
+    classify_ticket,
     call_routing,
     get_kb_collection,
-    encoder
+    encoder,
+    conversations
 )
 
-app = FastAPI(title="Smart Helpdesk AI Agent LangChain")
+app = FastAPI(title="Smart Helpdesk AI Agent")
 
-# -------------------------------
 # Request Models
-# -------------------------------
 class TicketRequest(BaseModel):
     text: str
-    user_email: str = None
+    conversation_id: Optional[str] = None
+    user_email: Optional[str] = None
 
-# -------------------------------
-# Persistent Chroma client
-# -------------------------------
+# Persistent Chroma settings
 CHROMA_PATH = "/tmp/chroma"
 COLLECTION_NAME = "knowledge_base"
 
@@ -44,12 +41,8 @@ COLLECTION_NAME = "knowledge_base"
 # -------------------------------
 @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.
-    """
+    """Upload and index knowledge base."""
     try:
-        # Load JSON from uploaded file
         content_bytes = await kb_file.read()
         data = json.loads(content_bytes)
 
@@ -58,29 +51,19 @@ async def setup_kb(kb_file: UploadFile = File(...)):
 
         print(f"📘 Loaded {len(data)} items from {kb_file.filename}")
 
-        # Get or create collection using shared function
         chroma_client = chromadb.PersistentClient(
             path=CHROMA_PATH,
-            settings=Settings(
-                anonymized_telemetry=False,
-                allow_reset=True
-            )
+            settings=Settings(anonymized_telemetry=False, allow_reset=True)
         )
         collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
 
-        # Clear existing records
         if collection.count() > 0:
             print(f"🧹 Clearing {collection.count()} existing records...")
             collection.delete(ids=collection.get()['ids'])
 
-        # Prepare texts, ids, and metadata
         texts, ids, metadatas = [], [], []
         for i, item in enumerate(data):
-            # Try multiple possible field names for text content
-            text = (item.get("answer") or 
-                   item.get("text") or 
-                   item.get("content") or "")
-            
+            text = item.get("answer") or item.get("text") or item.get("content") or ""
             item_id = item.get("id") or str(i)
             category = item.get("category", "")
             
@@ -88,42 +71,26 @@ async def setup_kb(kb_file: UploadFile = File(...)):
                 print(f"⚠️ Skipping item {i} - no text content")
                 continue
             
-            # Create a richer text for embedding by combining relevant fields
-            # This helps with better semantic search
-            combined_text = text
-            if category:
-                combined_text = f"Category: {category}. {text}"
-                
+            combined_text = f"Category: {category}. {text}" if category else text
             texts.append(combined_text)
             ids.append(str(item_id))
-            metadatas.append({
-                "id": str(item_id), 
-                "category": category,
-                "original_index": i
-            })
+            metadatas.append({"id": str(item_id), "category": category, "original_index": i})
 
         if not texts:
             raise HTTPException(status_code=400, detail="No valid text content found in JSON.")
 
-        # 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, 
-            documents=texts,
-            metadatas=metadatas
-        )
+        collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)
 
-        # Update the global reference in agent_langchain
+        # Update global reference
         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()}
+        print(f"✅ Successfully added {collection.count()} records")
+        return {"message": "Knowledge base initialized", "count": collection.count()}
 
     except json.JSONDecodeError:
         raise HTTPException(status_code=400, detail="Invalid JSON file.")
@@ -133,39 +100,81 @@ async def setup_kb(kb_file: UploadFile = File(...)):
         raise HTTPException(status_code=500, detail=f"Setup failed: {str(e)}")
 
 # -------------------------------
-# Step-by-Step Endpoints
+# MAIN ORCHESTRATE ENDPOINT (True Agent)
+# -------------------------------
+@app.post("/orchestrate")
+async def orchestrate_endpoint(ticket: TicketRequest):
+    """
+    Main AI Agent endpoint - handles everything:
+    - Initial ticket processing
+    - Follow-up conversations
+    - Automatic escalation decisions
+    - Multi-turn dialogue
+    """
+    try:
+        result = process_with_agent(
+            user_message=ticket.text,
+            conversation_id=ticket.conversation_id
+        )
+        
+        return {
+            "conversation_id": result["conversation_id"],
+            "response": result["response"],
+            "status": result["status"],
+            "message_count": result["message_count"],
+            "can_continue": result.get("can_continue", True),
+            "instructions": {
+                "continue_conversation": "Send another request with the same conversation_id",
+                "new_ticket": "Send a request without conversation_id"
+            }
+        }
+        
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Agent failed: {str(e)}")
+
+# -------------------------------
+# Get Conversation History
 # -------------------------------
+@app.get("/conversation/{conversation_id}")
+async def get_conversation(conversation_id: str):
+    """Retrieve full conversation history."""
+    conv = get_conversation_history(conversation_id)
+    if not conv:
+        raise HTTPException(status_code=404, detail="Conversation not found")
+    
+    return {
+        "conversation_id": conversation_id,
+        "messages": conv["messages"],
+        "created_at": conv["created_at"],
+        "message_count": len(conv["messages"])
+    }
 
+# -------------------------------
+# Individual Tool Endpoints (for testing)
+# -------------------------------
 @app.post("/classify")
 async def classify_endpoint(ticket: TicketRequest):
-    """Classify the ticket (impact, urgency, type)"""
+    """Test classification only."""
     classification = classify_ticket(ticket.text)
     return {"classification": classification}
 
 @app.post("/route")
 async def route_endpoint(ticket: TicketRequest):
-    """Route the ticket to department (Space 2)"""
+    """Test routing only."""
     department = call_routing(ticket.text)
     return {"department": department}
 
 @app.post("/kb_query")
 async def kb_query_endpoint(ticket: TicketRequest):
-    """Query the flattened KB directly using embeddings and return the best match."""
+    """Test KB query only."""
     collection = get_kb_collection()
-    
-    if not collection:
+    if not collection or collection.count() == 0:
         raise HTTPException(status_code=400, detail="KB not set up. Call /setup first.")
 
     try:
-        # 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 = collection.query(
             query_embeddings=[query_embedding],
             n_results=1,
@@ -175,53 +184,17 @@ async def kb_query_endpoint(ticket: TicketRequest):
         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 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(float(confidence), 3)
-        }
+        return {"answer": best_doc, "confidence": round(float(confidence), 3)}
 
     except Exception as e:
         import traceback
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"KB query failed: {str(e)}")
 
-# -------------------------------
-# Full Ticket Orchestration
-# -------------------------------
-@app.post("/orchestrate")
-async def orchestrate_endpoint(ticket: TicketRequest):
-    """Full ticket orchestration via LangChain agent with nicely formatted reasoning 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
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        raise HTTPException(status_code=500, detail=f"Orchestration failed: {str(e)}")
-
 # -------------------------------
 # Health Check
 # -------------------------------
@@ -234,5 +207,6 @@ async def health():
     return {
         "status": "ok",
         "kb_status": kb_status,
-        "kb_records": kb_count
+        "kb_records": kb_count,
+        "active_conversations": len(conversations)
     }