Create app.py

app.py

@@ -0,0 +1,371 @@
+import os
+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
+import asyncio
+from fastapi import FastAPI, HTTPException, UploadFile, File, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional, Dict, Set
+import chromadb
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+
+# Import from autonomous agent
+from agent_langchain import (
+    process_with_agent,
+    get_conversation_history,
+    classify_ticket,
+    call_routing,
+    get_kb_collection,
+    encoder,
+    conversations
+)
+
+app = FastAPI(title="Smart Helpdesk AI Agent - Autonomous + WebSocket")
+
+# CORS for frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Update with your frontend URL in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Request Models
+class TicketRequest(BaseModel):
+    text: str
+    conversation_id: Optional[str] = None
+    user_email: Optional[str] = None
+
+# WebSocket Connection Manager
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: Dict[str, WebSocket] = {}
+    
+    async def connect(self, websocket: WebSocket, conversation_id: str):
+        await websocket.accept()
+        self.active_connections[conversation_id] = websocket
+        print(f"🔌 WebSocket connected: {conversation_id}")
+    
+    def disconnect(self, conversation_id: str):
+        if conversation_id in self.active_connections:
+            del self.active_connections[conversation_id]
+            print(f"🔌 WebSocket disconnected: {conversation_id}")
+    
+    async def send_message(self, conversation_id: str, message: dict):
+        if conversation_id in self.active_connections:
+            try:
+                await self.active_connections[conversation_id].send_json(message)
+            except Exception as e:
+                print(f"Error sending message: {e}")
+                self.disconnect(conversation_id)
+
+manager = ConnectionManager()
+
+# Persistent Chroma settings
+CHROMA_PATH = "/tmp/chroma"
+COLLECTION_NAME = "knowledge_base"
+
+# -------------------------------
+# KB Setup Endpoint
+# -------------------------------
+@app.post("/setup")
+async def setup_kb(kb_file: UploadFile = File(...)):
+    """Upload and index knowledge base."""
+    try:
+        content_bytes = await kb_file.read()
+        data = json.loads(content_bytes)
+
+        if not isinstance(data, list):
+            raise HTTPException(status_code=400, detail="JSON must be a list of items.")
+
+        print(f"📘 Loaded {len(data)} items from {kb_file.filename}")
+
+        chroma_client = chromadb.PersistentClient(
+            path=CHROMA_PATH,
+            settings=Settings(anonymized_telemetry=False, allow_reset=True)
+        )
+        collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+
+        if collection.count() > 0:
+            print(f"🧹 Clearing {collection.count()} existing records...")
+            collection.delete(ids=collection.get()['ids'])
+
+        texts, ids, metadatas = [], [], []
+        for i, item in enumerate(data):
+            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", "")
+            
+            if not text:
+                print(f"⚠️ Skipping item {i} - no text content")
+                continue
+            
+            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})
+
+        if not texts:
+            raise HTTPException(status_code=400, detail="No valid text content found in JSON.")
+
+        print("🧠 Generating embeddings...")
+        embeddings = encoder.encode(texts, show_progress_bar=True).tolist()
+
+        print("💾 Adding to ChromaDB...")
+        collection.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)
+
+        # Update global reference
+        import agent_langchain
+        agent_langchain.kb_collection = collection
+
+        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.")
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=f"Setup failed: {str(e)}")
+
+# -------------------------------
+# WebSocket Endpoint (REAL-TIME BIDIRECTIONAL)
+# -------------------------------
+@app.websocket("/ws/{conversation_id}")
+async def websocket_endpoint(websocket: WebSocket, conversation_id: str):
+    """
+    WebSocket endpoint for real-time agent communication.
+    
+    Client sends: {"text": "My issue description", "user_email": "user@example.com"}
+    Server streams: 
+      - {"type": "status", "message": "Agent is thinking..."}
+      - {"type": "tool_use", "tool": "SearchKnowledgeBase", "input": "..."}
+      - {"type": "response", "content": "Here's the solution..."}
+      - {"type": "saved", "firestore_id": "abc123"}
+    """
+    await manager.connect(websocket, conversation_id)
+    
+    try:
+        while True:
+            # Receive message from client
+            data = await websocket.receive_json()
+            user_message = data.get("text")
+            user_email = data.get("user_email")
+            
+            if not user_message:
+                await manager.send_message(conversation_id, {
+                    "type": "error",
+                    "message": "No text provided"
+                })
+                continue
+            
+            # Send thinking status
+            await manager.send_message(conversation_id, {
+                "type": "status",
+                "message": "🤔 Analyzing your request..."
+            })
+            
+            # Callback for streaming updates
+            async def ws_callback(update):
+                await manager.send_message(conversation_id, update)
+            
+            # Process with agent (in thread to avoid blocking)
+            loop = asyncio.get_event_loop()
+            result = await loop.run_in_executor(
+                None,
+                lambda: process_with_agent(
+                    user_message=user_message,
+                    conversation_id=conversation_id,
+                    user_email=user_email,
+                    callback=lambda msg: asyncio.run_coroutine_threadsafe(ws_callback(msg), loop)
+                )
+            )
+            
+            # Send final response
+            await manager.send_message(conversation_id, {
+                "type": "response",
+                "conversation_id": result["conversation_id"],
+                "content": result["response"],
+                "status": result["status"],
+                "ticket_info": result.get("ticket_info", {}),
+                "message_count": result["message_count"],
+                "firestore_id": result.get("firestore_id")
+            })
+            
+    except WebSocketDisconnect:
+        manager.disconnect(conversation_id)
+        print(f"Client disconnected: {conversation_id}")
+    except Exception as e:
+        print(f"WebSocket error: {e}")
+        import traceback
+        traceback.print_exc()
+        try:
+            await manager.send_message(conversation_id, {
+                "type": "error",
+                "message": str(e)
+            })
+        except:
+            pass
+        manager.disconnect(conversation_id)
+
+# -------------------------------
+# REST Endpoint (backward compatible)
+# -------------------------------
+@app.post("/orchestrate")
+async def orchestrate_endpoint(ticket: TicketRequest):
+    """
+    REST endpoint for agent interaction (backward compatible).
+    Use WebSocket for real-time experience.
+    """
+    try:
+        result = process_with_agent(
+            user_message=ticket.text,
+            conversation_id=ticket.conversation_id,
+            user_email=ticket.user_email
+        )
+        
+        return {
+            "conversation_id": result["conversation_id"],
+            "response": result["response"],
+            "status": result["status"],
+            "ticket_info": result.get("ticket_info", {}),
+            "message_count": result["message_count"],
+            "reasoning_trace": result.get("reasoning_trace", []),
+            "firestore_id": result.get("firestore_id"),
+            "instructions": {
+                "websocket": f"ws://your-domain/ws/{result['conversation_id']}",
+                "continue_conversation": "Include the conversation_id in your next request"
+            }
+        }
+        
+    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"],
+        "ticket_info": conv.get("ticket_info", {}),
+        "status": conv.get("status", "unknown"),
+        "created_at": conv["created_at"],
+        "message_count": len(conv["messages"])
+    }
+
+# -------------------------------
+# List Active Conversations
+# -------------------------------
+@app.get("/conversations")
+async def list_conversations():
+    """List all active conversations."""
+    conv_list = []
+    for conv_id, conv_data in conversations.items():
+        conv_list.append({
+            "conversation_id": conv_id,
+            "status": conv_data.get("status", "unknown"),
+            "message_count": len(conv_data["messages"]),
+            "created_at": conv_data["created_at"],
+            "user_email": conv_data.get("user_email", "anonymous"),
+            "last_message": conv_data["messages"][-1]["content"][:100] if conv_data["messages"] else None
+        })
+    
+    return {
+        "total": len(conv_list),
+        "conversations": sorted(conv_list, key=lambda x: x["created_at"], reverse=True)
+    }
+
+# -------------------------------
+# Individual Tool Endpoints (for testing)
+# -------------------------------
+@app.post("/classify")
+async def classify_endpoint(ticket: TicketRequest):
+    """Test classification only."""
+    classification = classify_ticket(ticket.text)
+    return {"classification": classification}
+
+@app.post("/route")
+async def route_endpoint(ticket: TicketRequest):
+    """Test routing only."""
+    department = call_routing(ticket.text)
+    return {"department": department}
+
+@app.post("/kb_query")
+async def kb_query_endpoint(ticket: TicketRequest):
+    """Test KB query only."""
+    collection = get_kb_collection()
+    if not collection or collection.count() == 0:
+        raise HTTPException(status_code=400, detail="KB not set up. Call /setup first.")
+
+    try:
+        query_embedding = encoder.encode([ticket.text])[0].tolist()
+        result = collection.query(
+            query_embeddings=[query_embedding],
+            n_results=1,
+            include=["documents", "distances", "metadatas"]
+        )
+
+        if not result or not result.get('documents') or len(result['documents'][0]) == 0:
+            return {"answer": "No relevant KB found.", "confidence": 0.0}
+
+        best_doc = result['documents'][0][0]
+        best_distance = result['distances'][0][0] if result.get('distances') else 1.0
+        confidence = max(0.0, 1.0 - (best_distance / 2.0))
+
+        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)}")
+
+# -------------------------------
+# Health Check
+# -------------------------------
+@app.get("/health")
+async def health():
+    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,
+        "active_conversations": len(conversations),
+        "active_websockets": len(manager.active_connections),
+        "agent_type": "Autonomous ReAct Agent with Gemini + WebSocket"
+    }
+
+# -------------------------------
+# Root endpoint
+# -------------------------------
+@app.get("/")
+async def root():
+    return {
+        "message": "Smart Helpdesk AI Agent API",
+        "endpoints": {
+            "websocket": "/ws/{conversation_id}",
+            "rest": "/orchestrate",
+            "setup_kb": "/setup",
+            "conversations": "/conversations",
+            "health": "/health"
+        },
+        "documentation": "/docs"
+    }