Update agent_langchain.py

agent_langchain.py

@@ -1,8 +1,281 @@
-"""
-Complete fixed version of process_with_agent function in agent_langchain.py
-Replace the existing function (around line 180-330) with this version.
-"""
+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 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,
@@ -161,7 +434,7 @@ def process_with_agent(
                 "resolution": agent_response,
                 "created_at_iso": conv["created_at"],
                 "escalated": escalated,
-                "reasoning_trace": reasoning_trace  # Include for debugging
+                "reasoning_trace": reasoning_trace
             }
             
             # Add escalation reason if escalated
@@ -180,7 +453,7 @@ def process_with_agent(
             if firestore_id:
                 print(f"✅ Successfully saved to Firestore with ID: {firestore_id}")
             else:
-                print("❌ Failed to save to Firestore - check Firebase connection")
+                print("❌ Failed to save to Firestore")
             
             if callback:
                 callback({
@@ -199,7 +472,7 @@ def process_with_agent(
             "reasoning_trace": reasoning_trace,
             "ticket_info": ticket_info,
             "firestore_id": firestore_id,
-            "escalated": escalated  # Explicitly return this
+            "escalated": escalated
         }
         
     except Exception as e:
@@ -222,4 +495,8 @@ 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)