Update agent_langchain.py

agent_langchain.py

@@ -1,281 +1,8 @@
-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"
+"""
+Complete fixed version of process_with_agent function in agent_langchain.py
+Replace the existing function (around line 180-330) with this version.
+"""
 
-import json
-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")
-firebase_creds_json = os.getenv("FIREBASE_CREDS_JSON")
-
-# Initialize Firebase
-db = None
-if firebase_creds_json:
-    try:
-        creds_dict = json.loads(firebase_creds_json)
-        cred = credentials.Certificate(creds_dict)
-
-        if not firebase_admin._apps:
-            firebase_admin.initialize_app(cred)
-        db = firestore.client()
-        print("✅ Firebase initialized from FIREBASE_CREDS_JSON")
-    except Exception as e:
-        import traceback
-        print(f"⚠️ Firebase init failed: {e}")
-        traceback.print_exc()
-else:
-    print("⚠️ FIREBASE_CREDS_JSON not found in environment variables")
-
-# 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
-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 ONLY when KB confidence is below 0.6 AND issue is truly complex. Always try KB first!"""
-    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
-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). ALWAYS USE THIS FIRST before escalating. 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. CRITICAL: Use this ONLY as a LAST RESORT when: 1) KB confidence score is below 0.6 AND you've already tried KB, 2) Issue is extremely complex and unusual, 3) User explicitly confirms the KB solution failed after trying it. DO NOT escalate if KB has a reasonable solution (confidence > 0.6)."
-    )
-]
-
-# IMPROVED Agent Prompt
-AGENT_PROMPT = """You are an intelligent IT Helpdesk AI Agent. Your PRIMARY goal is to resolve tickets using the Knowledge Base. Escalation is a LAST RESORT.
-
-AVAILABLE TOOLS:
-{tools}
-
-TOOL NAMES: {tool_names}
-
-CRITICAL RULES:
-1. **ALWAYS search Knowledge Base FIRST** - This is your primary tool for resolution
-2. **Trust KB solutions with confidence >= 0.6** - These are reliable solutions, provide them to users
-3. **ONLY escalate when ABSOLUTELY necessary**:
-   - KB confidence is below 0.6 AND no solution found
-   - Issue is extremely unusual or complex beyond KB scope
-   - User explicitly tried your KB solution and reports it failed
-4. **Be thorough with KB** - If first search doesn't work, try rephrasing the query
-5. **Maintain context** - Remember conversation history for follow-ups
-
-DECISION WORKFLOW:
-NEW TICKET → Search KB → If confidence >= 0.6 → Provide solution → Mark RESOLVED
-         ↓
-   If confidence < 0.6 → Try rephrasing search → Still low? → Classify & Route → THEN escalate
-
-FOLLOW-UP → Check if user tried solution → Worked? → Mark RESOLVED
-                                      ↓
-                              Failed? → Search KB again with different query → Still failing? → THEN escalate
-
-IMPORTANT: 
-- Don't mention tool names or confidence scores to users
-- Provide clear, step-by-step instructions from KB
-- Be conversational and helpful
-- Escalation means you couldn't solve it - avoid this outcome!
-
-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
-
-Begin!
-
-Question: {input}
-Thought: {agent_scratchpad}"""
-
-prompt = PromptTemplate.from_template(AGENT_PROMPT)
-
-# Create Agent
-agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
-agent_executor = AgentExecutor(
-    agent=agent,
-    tools=tools,
-    verbose=True,
-    max_iterations=6,
-    handle_parsing_errors="Check your output and make sure it conforms to the format instructions!",
-    return_intermediate_steps=True,
-    early_stopping_method="force"
-)
-
-# Main Processing Function
 def process_with_agent(
     user_message: str, 
     conversation_id: str = None,
@@ -323,26 +50,48 @@ def process_with_agent(
         agent_response = result.get("output", "I apologize, I encountered an error.")
         intermediate_steps = result.get("intermediate_steps", [])
         
-        # Determine status based on agent's response
+        # IMPROVED: Check tool usage to determine status
         status = "in_progress"
         should_save = False
         escalated = False
+        used_escalation_tool = False
         
-        # Check for escalation FIRST (most important)
-        if "ESCALATED" in agent_response or "TKT-" in agent_response:
+        # First, check if EscalateToHuman tool was actually used
+        for action, observation in intermediate_steps:
+            if action.tool == "EscalateToHuman":
+                used_escalation_tool = True
+                print("🔍 Detected EscalateToHuman tool usage")
+                break
+        
+        # Determine final status
+        if used_escalation_tool:
+            # Tool was used - definitely escalated
+            status = "escalated"
+            should_save = True
+            escalated = True
+            print("🔴 Ticket ESCALATED (tool used) - Saving to Firebase")
+        elif "ESCALATED" in agent_response or "TKT-" in agent_response:
+            # Fallback: Check response text
             status = "escalated"
             should_save = True
             escalated = True
-            print("🔴 Ticket ESCALATED - Saving to Firebase")
-        # Then check for resolution
-        elif any(phrase in agent_response.lower() for phrase in ["resolved", "you're all set", "should work now", "problem solved", "this should fix", "try these steps"]):
+            print("🔴 Ticket ESCALATED (text match) - Saving to Firebase")
+        elif any(phrase in agent_response.lower() for phrase in [
+            "resolved", "you're all set", "should work now", "problem solved", 
+            "this should fix", "try these steps", "follow these steps"
+        ]):
+            # Resolution detected
             status = "resolved"
             should_save = True
             print("✅ Ticket RESOLVED - Saving to Firebase")
+        else:
+            # Still in progress
+            print("⏳ Ticket IN PROGRESS - Not saving yet")
         
         # Extract ticket info from tools
         ticket_info = conv.get("ticket_info", {})
         kb_confidence = None
+        escalation_reason = None
         
         for action, observation in intermediate_steps:
             if action.tool == "ClassifyTicket":
@@ -354,8 +103,10 @@ def process_with_agent(
                         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: ", "")
+            
             elif action.tool == "SearchKnowledgeBase":
                 # Extract confidence from KB response
                 if "[KB Confidence:" in str(observation):
@@ -365,10 +116,15 @@ def process_with_agent(
                         ticket_info["kb_confidence"] = kb_confidence
                     except:
                         pass
+            
+            elif action.tool == "EscalateToHuman":
+                # Capture escalation reason
+                escalation_reason = action.tool_input
         
         conv["ticket_info"] = ticket_info
         conv["status"] = status
         
+        # Build reasoning trace for debugging
         reasoning_trace = []
         for action, observation in intermediate_steps:
             reasoning_trace.append({
@@ -394,6 +150,8 @@ def process_with_agent(
         # Save to Firestore if resolved OR escalated
         firestore_id = None
         if should_save:
+            print(f"💾 Preparing to save ticket - Status: {status}, Escalated: {escalated}")
+            
             firestore_data = {
                 "conversation_id": conversation_id,
                 "status": status,
@@ -402,21 +160,27 @@ def process_with_agent(
                 "messages": conv["messages"],
                 "resolution": agent_response,
                 "created_at_iso": conv["created_at"],
-                "escalated": escalated
+                "escalated": escalated,
+                "reasoning_trace": reasoning_trace  # Include for debugging
             }
             
             # Add escalation reason if escalated
             if escalated:
-                escalation_reason = "Low KB confidence" if kb_confidence and kb_confidence < 0.6 else "Complex issue requiring human intervention"
-                firestore_data["escalation_reason"] = escalation_reason
-                print(f"📤 Saving escalated ticket to Firebase - Reason: {escalation_reason}")
+                if escalation_reason:
+                    firestore_data["escalation_reason"] = escalation_reason
+                elif kb_confidence is not None and kb_confidence < 0.6:
+                    firestore_data["escalation_reason"] = f"Low KB confidence: {kb_confidence}"
+                else:
+                    firestore_data["escalation_reason"] = "Complex issue requiring human intervention"
+                
+                print(f"🔼 Saving escalated ticket - Reason: {firestore_data['escalation_reason']}")
             
             firestore_id = save_ticket_to_firestore(firestore_data)
             
             if firestore_id:
                 print(f"✅ Successfully saved to Firestore with ID: {firestore_id}")
             else:
-                print("❌ Failed to save to Firestore")
+                print("❌ Failed to save to Firestore - check Firebase connection")
             
             if callback:
                 callback({
@@ -424,6 +188,8 @@ def process_with_agent(
                     "firestore_id": firestore_id,
                     "status": status
                 })
+        else:
+            print(f"⏭️  Not saving - Status: {status}, Should save: {should_save}")
         
         return {
             "conversation_id": conversation_id,
@@ -432,11 +198,12 @@ def process_with_agent(
             "message_count": len(conv["messages"]),
             "reasoning_trace": reasoning_trace,
             "ticket_info": ticket_info,
-            "firestore_id": firestore_id
+            "firestore_id": firestore_id,
+            "escalated": escalated  # Explicitly return this
         }
         
     except Exception as e:
-        print(f"Agent error: {e}")
+        print(f"❌ Agent error: {e}")
         import traceback
         traceback.print_exc()
         
@@ -455,8 +222,4 @@ def process_with_agent(
             "response": error_response,
             "status": "error",
             "error": str(e)
-        }
-
-def get_conversation_history(conversation_id: str):
-    """Get conversation history."""
-    return conversations.get(conversation_id)
+        }