Create agent_langchain.py

agent_langchain.py

@@ -0,0 +1,420 @@
+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 requests
+import torch
+import time
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import chromadb
+from chromadb.config import Settings
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.agents import AgentExecutor, create_react_agent
+from langchain.tools import Tool
+from langchain.prompts import PromptTemplate
+import threading
+from datetime import datetime
+import firebase_admin
+from firebase_admin import credentials, firestore
+from typing import Optional, Dict, Any
+
+# Environment
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+ROUTING_URL = os.environ.get("ROUTING_URL")
+SPACE_URL = os.environ.get("SPACE_URL", "http://localhost:7860")
+FIREBASE_CREDS_PATH = os.environ.get("FIREBASE_CREDS_PATH")
+
+# Initialize Firebase
+db = None
+if FIREBASE_CREDS_PATH and os.path.exists(FIREBASE_CREDS_PATH):
+    try:
+        if not firebase_admin._apps:
+            cred = credentials.Certificate(FIREBASE_CREDS_PATH)
+            firebase_admin.initialize_app(cred)
+        db = firestore.client()
+        print("✅ Firebase initialized")
+    except Exception as e:
+        print(f"⚠️ Firebase init failed: {e}")
+
+# Label Dictionary
+LABEL_DICTIONARY = {
+    "I1": "Low Impact", "I2": "Medium Impact", "I3": "High Impact", "I4": "Critical Impact",
+    "U1": "Low Urgency", "U2": "Medium Urgency", "U3": "High Urgency", "U4": "Critical Urgency",
+    "T1": "Information", "T2": "Incident", "T3": "Problem", "T4": "Request", "T5": "Question"
+}
+
+# Classification Model
+clf_model_name = "DavinciTech/BERT_Categorizer"
+clf_tokenizer = AutoTokenizer.from_pretrained(clf_model_name, cache_dir="/tmp/transformers")
+clf_model = AutoModelForSequenceClassification.from_pretrained(clf_model_name, cache_dir="/tmp/transformers")
+
+def classify_ticket(text):
+    """Classify ticket into Impact, Urgency, and Type."""
+    inputs = clf_tokenizer(text, return_tensors="pt", truncation=True)
+    outputs = clf_model(**inputs)
+    logits = outputs.logits[0]
+    impact_idx = torch.argmax(logits[:4]).item() + 1
+    urgency_idx = torch.argmax(logits[4:8]).item() + 1
+    type_idx = torch.argmax(logits[8:]).item() + 1
+    return {
+        "impact": LABEL_DICTIONARY[f"I{impact_idx}"],
+        "urgency": LABEL_DICTIONARY[f"U{urgency_idx}"],
+        "type": LABEL_DICTIONARY[f"T{type_idx}"]
+    }
+
+# Routing Function
+def call_routing(text, retries=3, delay=5):
+    """Route ticket to appropriate department."""
+    url = ROUTING_URL if ROUTING_URL else f"{SPACE_URL}/route"
+    for attempt in range(retries):
+        try:
+            resp = requests.post(url, json={"text": text}, timeout=30)
+            resp.raise_for_status()
+            return resp.json().get("department", "General IT")
+        except Exception as e:
+            print(f"Routing attempt {attempt+1} failed: {e}")
+            if attempt < retries - 1:
+                time.sleep(delay)
+    return "General IT"
+
+# Knowledge Base
+CHROMA_PATH = "/tmp/chroma"
+COLLECTION_NAME = "knowledge_base"
+kb_collection = None
+kb_lock = threading.Lock()
+encoder = SentenceTransformer("all-MiniLM-L6-v2", cache_folder="/tmp/sentence_transformers")
+
+def get_kb_collection():
+    global kb_collection
+    if kb_collection is None:
+        with kb_lock:
+            if kb_collection is None:
+                try:
+                    chroma_client = chromadb.PersistentClient(
+                        path=CHROMA_PATH,
+                        settings=Settings(anonymized_telemetry=False, allow_reset=True)
+                    )
+                    kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+                except Exception as e:
+                    print(f"Could not get KB collection: {e}")
+    return kb_collection
+
+def query_kb(text: str, top_k: int = 1):
+    """Query KB and return answer with confidence."""
+    collection = get_kb_collection()
+    if not collection or collection.count() == 0:
+        return {"answer": None, "confidence": 0.0}
+    
+    try:
+        query_embedding = encoder.encode([text])[0].tolist()
+        results = collection.query(
+            query_embeddings=[query_embedding],
+            n_results=top_k,
+            include=["documents", "distances", "metadatas"]
+        )
+        
+        if not results or not results.get("documents") or len(results["documents"][0]) == 0:
+            return {"answer": None, "confidence": 0.0}
+        
+        answer = results["documents"][0][0]
+        distance = results["distances"][0][0] if results.get("distances") else 1.0
+        confidence = max(0.0, 1.0 - (distance / 2.0))
+        
+        return {"answer": answer, "confidence": round(float(confidence), 3)}
+    except Exception as e:
+        print(f"KB query failed: {e}")
+        return {"answer": None, "confidence": 0.0}
+
+# Firestore Helper
+def save_ticket_to_firestore(ticket_data: Dict[str, Any]):
+    """Save resolved/escalated ticket to Firestore."""
+    if not db:
+        print("⚠️ Firestore not initialized, skipping save")
+        return None
+    
+    try:
+        ticket_ref = db.collection('tickets').document()
+        ticket_data['created_at'] = firestore.SERVER_TIMESTAMP
+        ticket_data['updated_at'] = firestore.SERVER_TIMESTAMP
+        ticket_ref.set(ticket_data)
+        print(f"✅ Ticket saved to Firestore: {ticket_ref.id}")
+        return ticket_ref.id
+    except Exception as e:
+        print(f"❌ Firestore save failed: {e}")
+        return None
+
+# Gemini LLM
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash-exp",
+    temperature=0.3,
+    google_api_key=GEMINI_API_KEY
+)
+
+# Global conversation storage
+conversations = {}
+
+# Tool Functions for Agent (TRULY AUTONOMOUS NOW)
+def classify_tool(query: str) -> str:
+    """Analyzes ticket severity, impact, urgency, and type. Use when you need to understand ticket priority."""
+    result = classify_ticket(query)
+    return f"Impact: {result['impact']}, Urgency: {result['urgency']}, Type: {result['type']}"
+
+def routing_tool(query: str) -> str:
+    """Identifies which IT department should handle this issue. Use when you need to know responsible team."""
+    dept = call_routing(query)
+    return f"Department: {dept}"
+
+def kb_tool(query: str) -> str:
+    """Searches knowledge base for solutions. Returns answer with confidence score. Use when you need technical solutions."""
+    result = query_kb(query)
+    if result["answer"] and result["confidence"] > 0.5:
+        return f"[KB Confidence: {result['confidence']}]\n{result['answer']}"
+    return f"[KB Confidence: {result['confidence']}] No relevant solution found in knowledge base."
+
+def escalation_tool(reason: str) -> str:
+    """Creates escalation ticket for human agent. Use when: KB confidence is low, issue is complex, or user reports solution failed."""
+    ticket_id = f"TKT-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+    return f"ESCALATED: Ticket {ticket_id} created. Reason: {reason}. Human agent will respond in 2-4 hours."
+
+# Define Tools with better descriptions
+tools = [
+    Tool(
+        name="ClassifyTicket",
+        func=classify_tool,
+        description="Analyzes ticket to determine impact level, urgency, and type. Use this when you need to understand the severity or priority of an issue."
+    ),
+    Tool(
+        name="RouteTicket",
+        func=routing_tool,
+        description="Determines which IT department should handle this ticket. Use this when you need to identify the responsible team."
+    ),
+    Tool(
+        name="SearchKnowledgeBase",
+        func=kb_tool,
+        description="Searches internal knowledge base for solutions. Returns answer with confidence score (0-1). Use this when you need to find technical solutions or troubleshooting steps."
+    ),
+    Tool(
+        name="EscalateToHuman",
+        func=escalation_tool,
+        description="Creates an escalation ticket for human agent review. Use this ONLY when: 1) KB confidence score is below 0.75, 2) Issue is highly complex or unusual, 3) User confirms solution didn't work."
+    )
+]
+
+# IMPROVED Agent Prompt - More Autonomous
+AGENT_PROMPT = """You are an intelligent IT Helpdesk AI Agent. Your goal is to efficiently resolve IT support tickets using the tools available to you.
+
+AVAILABLE TOOLS:
+{tools}
+
+TOOL NAMES: {tool_names}
+
+GUIDING PRINCIPLES:
+1. **Think autonomously** - Decide which tools you need based on the situation, not a fixed sequence
+2. **Be efficient** - Only use tools when they add value to solving the user's problem
+3. **Trust high-confidence solutions** - If KB returns confidence >= 0.75, provide that solution
+4. **Escalate wisely** - Only escalate when truly necessary (low KB confidence, complex issues, or failed solutions)
+5. **Maintain context** - Remember previous conversation history when handling follow-ups
+6. **Be empathetic** - Users are frustrated when things break; be professional and supportive
+
+DECISION FRAMEWORK:
+- For a NEW ticket: You might need to classify, route, search KB - but decide based on what's needed
+- For FOLLOW-UPS: If user says solution worked → close positively. If it failed → search KB again or escalate
+- For SIMPLE questions: You may not need all tools - use your judgment
+- For COMPLEX issues: Use multiple tools to gather information before providing solution
+
+FORMAT:
+Question: the user's input
+Thought: your reasoning about what to do next
+Action: the tool to use (must be one of [{tool_names}])
+Action Input: the input for that tool
+Observation: the tool's output
+... (repeat Thought/Action/Observation as needed)
+Thought: I now have enough information to respond
+Final Answer: your complete response to the user
+
+IMPORTANT: 
+- Don't mention tool names or technical process to users
+- Provide clear, step-by-step instructions
+- Ask clarifying questions if needed
+- Be conversational and helpful
+
+Begin!
+
+Question: {input}
+Thought: {agent_scratchpad}"""
+
+prompt = PromptTemplate.from_template(AGENT_PROMPT)
+
+# Create Agent with more flexibility
+agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
+agent_executor = AgentExecutor(
+    agent=agent,
+    tools=tools,
+    verbose=True,
+    max_iterations=8,  # Increased for more autonomy
+    handle_parsing_errors=True,
+    return_intermediate_steps=True,
+    early_stopping_method="generate"  # Allow agent to decide when to stop
+)
+
+# Main Processing Function
+def process_with_agent(
+    user_message: str, 
+    conversation_id: str = None,
+    user_email: str = None,
+    callback=None  # For streaming updates via WebSocket
+):
+    """Process user message through autonomous AI agent."""
+    
+    if not conversation_id:
+        conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hash(user_message) % 10000}"
+    
+    if conversation_id not in conversations:
+        conversations[conversation_id] = {
+            "messages": [],
+            "ticket_info": {},
+            "created_at": datetime.now().isoformat(),
+            "user_email": user_email,
+            "status": "open"
+        }
+    
+    conv = conversations[conversation_id]
+    
+    conv["messages"].append({
+        "role": "user",
+        "content": user_message,
+        "timestamp": datetime.now().isoformat()
+    })
+    
+    if callback:
+        callback({"type": "status", "message": "Agent is thinking..."})
+    
+    # Build context for follow-ups
+    if len(conv["messages"]) > 1:
+        context = f"CONVERSATION HISTORY:\n"
+        for msg in conv["messages"][-6:-1]:  # Last 5 messages for context
+            context += f"{msg['role'].upper()}: {msg['content']}\n"
+        context += f"\nCURRENT MESSAGE: {user_message}"
+        agent_input = context
+    else:
+        agent_input = user_message
+    
+    try:
+        result = agent_executor.invoke({"input": agent_input})
+        
+        agent_response = result.get("output", "I apologize, I encountered an error.")
+        intermediate_steps = result.get("intermediate_steps", [])
+        
+        # Determine status and handle Firestore
+        status = "in_progress"
+        should_save = False
+        
+        # Check for resolution indicators
+        if any(phrase in agent_response.lower() for phrase in ["resolved", "you're all set", "should work now", "problem solved"]):
+            status = "resolved"
+            should_save = True
+        elif "ESCALATED" in agent_response or "TKT-" in agent_response:
+            status = "escalated"
+            should_save = True
+        
+        # Extract ticket info from tools
+        ticket_info = conv.get("ticket_info", {})
+        for action, observation in intermediate_steps:
+            if action.tool == "ClassifyTicket":
+                # Parse classification
+                parts = str(observation).split(", ")
+                for part in parts:
+                    if "Impact:" in part:
+                        ticket_info["impact"] = part.split(": ")[1]
+                    elif "Urgency:" in part:
+                        ticket_info["urgency"] = part.split(": ")[1]
+                    elif "Type:" in part:
+                        ticket_info["type"] = part.split(": ")[1]
+            elif action.tool == "RouteTicket":
+                ticket_info["department"] = str(observation).replace("Department: ", "")
+        
+        conv["ticket_info"] = ticket_info
+        conv["status"] = status
+        
+        reasoning_trace = []
+        for action, observation in intermediate_steps:
+            reasoning_trace.append({
+                "tool": action.tool,
+                "input": action.tool_input,
+                "output": str(observation)[:200]
+            })
+            
+            if callback:
+                callback({
+                    "type": "tool_use",
+                    "tool": action.tool,
+                    "input": action.tool_input
+                })
+        
+        conv["messages"].append({
+            "role": "assistant",
+            "content": agent_response,
+            "timestamp": datetime.now().isoformat(),
+            "reasoning": reasoning_trace
+        })
+        
+        # Save to Firestore if resolved/escalated
+        firestore_id = None
+        if should_save:
+            firestore_data = {
+                "conversation_id": conversation_id,
+                "status": status,
+                "user_email": user_email or "anonymous",
+                "ticket_info": ticket_info,
+                "messages": conv["messages"],
+                "resolution": agent_response,
+                "created_at_iso": conv["created_at"]
+            }
+            firestore_id = save_ticket_to_firestore(firestore_data)
+            
+            if callback:
+                callback({
+                    "type": "saved",
+                    "firestore_id": firestore_id
+                })
+        
+        return {
+            "conversation_id": conversation_id,
+            "response": agent_response,
+            "status": status,
+            "message_count": len(conv["messages"]),
+            "reasoning_trace": reasoning_trace,
+            "ticket_info": ticket_info,
+            "firestore_id": firestore_id
+        }
+        
+    except Exception as e:
+        print(f"Agent error: {e}")
+        import traceback
+        traceback.print_exc()
+        
+        error_response = "I apologize, I encountered an error. Please try again or I can escalate this to a human agent."
+        conv["messages"].append({
+            "role": "assistant",
+            "content": error_response,
+            "timestamp": datetime.now().isoformat()
+        })
+        
+        if callback:
+            callback({"type": "error", "message": str(e)})
+        
+        return {
+            "conversation_id": conversation_id,
+            "response": error_response,
+            "status": "error",
+            "error": str(e)
+        }
+
+def get_conversation_history(conversation_id: str):
+    """Get conversation history."""
+    return conversations.get(conversation_id)