Update agent_langchain.py

agent_langchain.py

@@ -1,282 +1,325 @@
 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"
-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
+from typing import Optional, Dict, Any
+from datetime import datetime
 import chromadb
 from chromadb.config import Settings
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.agents import initialize_agent, Tool, AgentType
-from langchain.memory import ConversationBufferMemory
-import threading
-
-# ==============================================================
-# 🌎 ENVIRONMENT & GLOBAL SETTINGS
-# ==============================================================
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-ROUTING_URL = os.environ.get("ROUTING_URL")  # Space 2 URL
-SPACE_URL = os.environ.get("SPACE_URL", "http://localhost:7860")
+from sentence_transformers import SentenceTransformer
+from langchain_groq import ChatGroq
+from langchain_core.prompts import PromptTemplate
 
-# ==============================================================
-# 🏷️ 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"
-}
+# Environment setup
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# ==============================================================
-# 🤖 LOAD 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"
+# Initialize models
+encoder = SentenceTransformer('all-MiniLM-L6-v2')
+llm = ChatGroq(
+    model="llama-3.3-70b-versatile",
+    temperature=0.3,
+    api_key=os.getenv("GROQ_API_KEY")
 )
 
-def classify_ticket(text):
-    """Classify the 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 (Space 2)
-# ==============================================================
-def call_routing(text, retries=3, delay=5):
-    """Call Space 2 routing endpoint and get department only."""
-    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()
-            data = resp.json()
-            return data.get("department", "General IT")
-        except Exception as e:
-            print(f"Routing attempt {attempt+1} failed: {e}")
-            if attempt < retries - 1:
-                time.sleep(delay)
-            else:
-                return "General IT"
+# Global storage
+conversations = {}
+kb_collection = None
 
-# ==============================================================
-# 📚 KNOWLEDGE BASE SETUP
-# ==============================================================
+# Chroma settings
 CHROMA_PATH = "/tmp/chroma"
 COLLECTION_NAME = "knowledge_base"
 
-# Shared KB collection reference
-kb_collection = None
-kb_lock = threading.Lock()
-
-# Initialize encoder once
-encoder = SentenceTransformer("all-MiniLM-L6-v2", cache_folder="/tmp/sentence_transformers")
+# ===========================
+# Knowledge Base Functions
+# ===========================
 
 def get_kb_collection():
-    """Get or create the KB collection."""
+    """Get or initialize 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}")
+        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"Warning: Could not initialize KB collection: {e}")
     return kb_collection
 
-
-def query_kb(text: str, top_k: int = 1):
-    """
-    Query the Chroma knowledge base using SentenceTransformer embeddings.
-    Returns answer + confidence.
-    """
+def query_kb(query: str, n_results: int = 1) -> Dict[str, Any]:
+    """Query knowledge base for relevant information."""
     collection = get_kb_collection()
     
-    if not collection:
-        return {"answer": None, "confidence": 0.0, "metadata": {}}
+    if not collection or collection.count() == 0:
+        return {"answer": None, "confidence": 0.0}
     
     try:
-        # Check if collection has data
-        count = collection.count()
-        if count == 0:
-            return {"answer": None, "confidence": 0.0, "metadata": {}}
-        
-        # Embed the query
-        query_embedding = encoder.encode([text])[0].tolist()
-
-        # Query Chroma by embeddings
-        results = collection.query(
+        query_embedding = encoder.encode([query])[0].tolist()
+        result = collection.query(
             query_embeddings=[query_embedding],
-            n_results=top_k,
-            include=["documents", "distances", "metadatas", "embeddings"]
+            n_results=n_results,
+            include=["documents", "distances", "metadatas"]
         )
-
-        if not results or not results.get("documents") or len(results["documents"][0]) == 0:
-            return {"answer": None, "confidence": 0.0, "metadata": {}}
-
-        # Extract top document and metadata
-        answer = results["documents"][0][0]
-        metadata = results["metadatas"][0][0] if results.get("metadatas") and results["metadatas"][0] else {}
         
-        # Calculate confidence from distance
-        # ChromaDB returns L2 (Euclidean) distance by default - lower is better
-        if results.get("distances") and len(results["distances"][0]) > 0:
-            distance = results["distances"][0][0]
-            # Convert L2 distance to similarity score (0-1 range)
-            # For normalized embeddings, max L2 distance is ~2.0
-            confidence = max(0.0, 1.0 - (distance / 2.0))
-        else:
-            # Fallback: compute cosine similarity if embeddings are available
-            if results.get("embeddings") and len(results["embeddings"][0]) > 0:
-                stored_embedding = np.array(results["embeddings"][0][0])
-                query_vec = np.array(query_embedding)
-                confidence = float(np.dot(query_vec, stored_embedding) /
-                                 (np.linalg.norm(query_vec) * np.linalg.norm(stored_embedding) + 1e-8))
-            else:
-                confidence = 0.5
-
+        if not result or not result.get('documents') or len(result['documents'][0]) == 0:
+            return {"answer": None, "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": answer, 
-            "confidence": round(float(confidence), 3), 
-            "metadata": metadata
+            "answer": best_doc,
+            "confidence": float(confidence),
+            "metadata": result['metadatas'][0][0] if result.get('metadatas') else {}
         }
+    except Exception as e:
+        print(f"KB query error: {e}")
+        return {"answer": None, "confidence": 0.0}
+
+# ===========================
+# Classification & Routing
+# ===========================
+
+def classify_ticket(ticket_text: str) -> str:
+    """Classify ticket into priority/category."""
+    prompt = PromptTemplate(
+        input_variables=["ticket"],
+        template="""Classify this IT support ticket into ONE of these categories:
+- password_reset
+- software_issue
+- hardware_problem
+- network_issue
+- access_request
+- general_inquiry
+
+Ticket: {ticket}
+
+Category (one word only):"""
+    )
+    
+    try:
+        response = llm.invoke(prompt.format(ticket=ticket_text))
+        classification = response.content.strip().lower()
         
+        valid_categories = ["password_reset", "software_issue", "hardware_problem", 
+                          "network_issue", "access_request", "general_inquiry"]
+        
+        return classification if classification in valid_categories else "general_inquiry"
     except Exception as e:
-        print(f"⚠️ KB query failed: {e}")
-        import traceback
-        traceback.print_exc()
-        return {"answer": None, "confidence": 0.0, "metadata": {}}
-
-
-# ==============================================================
-# 🧠 GEMINI LLM (Official LangChain Integration)
-# ==============================================================
-llm = ChatGoogleGenerativeAI(
-    model="gemini-2.0-flash-exp",
-    temperature=0.3,
-    google_api_key=GEMINI_API_KEY
-)
-
-# ==============================================================
-# 🧰 DEFINE LANGCHAIN TOOLS
-# ==============================================================
-tools = [
-    Tool(
-        name="TicketClassifier",
-        func=lambda text: classify_ticket(text),
-        description="Classifies the ticket into impact, urgency, and type. Mandatory tool."
-    ),
-    Tool(
-        name="RoutingTool",
-        func=lambda text: call_routing(text),
-        description="Determines which department should handle the ticket (via Space 2). Mandatory tool."
-    ),
-    Tool(
-        name="KnowledgeBaseTool",
-        func=lambda text: query_kb(text)["answer"],
-        description="Searches the KB for relevant solutions. Returns a descriptive answer."
+        print(f"Classification error: {e}")
+        return "general_inquiry"
+
+def call_routing(ticket_text: str) -> str:
+    """Route ticket to appropriate department."""
+    prompt = PromptTemplate(
+        input_variables=["ticket"],
+        template="""Route this IT ticket to the correct department. Choose ONE:
+- IT Support
+- Network Team
+- Security Team
+- Hardware Team
+- Access Management
+
+Ticket: {ticket}
+
+Department (exact name only):"""
     )
-]
-
-# ==============================================================
-# 💬 MEMORY & AGENT INITIALIZATION
-# ==============================================================
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-agent_executor = initialize_agent(
-    tools=tools,
-    llm=llm,
-    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-    memory=memory,
-    verbose=False
-)
+    
+    try:
+        response = llm.invoke(prompt.format(ticket=ticket_text))
+        department = response.content.strip()
+        
+        valid_depts = ["IT Support", "Network Team", "Security Team", "Hardware Team", "Access Management"]
+        
+        return department if department in valid_depts else "IT Support"
+    except Exception as e:
+        print(f"Routing error: {e}")
+        return "IT Support"
 
-# ==============================================================
-# 🧾 MAIN TICKET PROCESSOR
-# ==============================================================
-def process_ticket_langchain(ticket_text):
-    """Full pipeline: classify → route → query KB → decide KB vs Gemini."""
-    reasoning_trace = []
+# ===========================
+# Main Processing Functions
+# ===========================
 
-    # Step 1: Classification
+def process_ticket_langchain(ticket_text: str) -> Dict[str, Any]:
+    """
+    Original single-turn ticket processing.
+    Used for initial ticket intake.
+    """
+    # Step 1: Classify
     classification = classify_ticket(ticket_text)
-    reasoning_trace.append(f"[Classifier] Impact: {classification['impact']}, Urgency: {classification['urgency']}, Type: {classification['type']}")
-
-    # Step 2: Routing
+    
+    # Step 2: Route
     department = call_routing(ticket_text)
-    reasoning_trace.append(f"[Routing] Assigned Department: {department}")
-
-    # Step 3: KB Search
+    
+    # Step 3: Query KB
     kb_result = query_kb(ticket_text)
-    reasoning_trace.append(f"[KB Search] Top Answer: '{kb_result['answer']}' (confidence: {kb_result['confidence']})")
-
-    # Step 4: KB vs LLM Decision
-    if kb_result["answer"] and kb_result["confidence"] >= 0.75:
-        final_answer = kb_result["answer"]
+    
+    # Step 4: Generate response
+    if kb_result["answer"] and kb_result["confidence"] >= 0.7:
+        answer = kb_result["answer"]
         status = "resolved"
-        reasoning_trace.append("[Decision] High KB confidence → ticket resolved via KB.")
     else:
-        llm_prompt = f"""
-You are a professional IT helpdesk agent.
-A user submitted the following ticket: "{ticket_text}"
+        # Generate ticket ID and escalate
+        ticket_id = f"TKT-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+        answer = f"""I couldn't find a confident answer in our knowledge base.
 
-Ticket classification: {classification}
-Assigned department: {department}
-Knowledge base result: {kb_result['answer']} (confidence: {kb_result['confidence']})
+**Ticket Created:** {ticket_id}
+**Department:** {department}
+**Classification:** {classification}
 
-Please provide a clear, descriptive, and professional IT helpdesk response.
-"""
-        final_answer = llm.invoke(llm_prompt).content
+A specialist will review your ticket and respond within 2-4 business hours."""
         status = "escalated"
-        reasoning_trace.append("[Decision] Low KB confidence → fallback to Gemini LLM for escalation.")
-
+    
     return {
-        "status": status,
         "classification": classification,
         "department": department,
-        "answer": final_answer,
-        "reasoning_trace": reasoning_trace
-    }
+        "answer": answer,
+        "confidence": kb_result["confidence"],
+        "status": status
+    }
+
+def process_with_agent(user_message: str, conversation_id: Optional[str] = None) -> Dict[str, Any]:
+    """
+    Full agentic conversation handler with memory and escalation logic.
+    This is the main function used by the /orchestrate endpoint.
+    """
+    # Generate conversation ID if not provided
+    if not conversation_id:
+        conversation_id = f"conv_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hash(user_message) % 10000}"
+    
+    # Initialize conversation if new
+    if conversation_id not in conversations:
+        conversations[conversation_id] = {
+            "messages": [],
+            "ticket_info": None,
+            "created_at": datetime.now().isoformat(),
+            "escalated": False
+        }
+    
+    conv = conversations[conversation_id]
+    
+    # Don't continue if already escalated
+    if conv.get("escalated", False):
+        return {
+            "conversation_id": conversation_id,
+            "response": "This ticket has been escalated to a human agent. They will contact you soon.",
+            "status": "escalated",
+            "message_count": len(conv["messages"]),
+            "can_continue": False
+        }
+    
+    # Add user message
+    conv["messages"].append({
+        "role": "user",
+        "content": user_message,
+        "timestamp": datetime.now().isoformat()
+    })
+    
+    # First message - full orchestration
+    if len(conv["messages"]) == 1:
+        result = process_ticket_langchain(user_message)
+        conv["ticket_info"] = {
+            "classification": result["classification"],
+            "department": result["department"],
+            "initial_query": user_message
+        }
+        
+        response_text = result["answer"]
+        status = result["status"]
+        
+        if status == "escalated":
+            conv["escalated"] = True
+    
+    # Follow-up messages
+    else:
+        # Check for escalation keywords
+        escalation_keywords = ["not working", "didn't work", "still broken", "still not", 
+                              "escalate", "human", "agent", "supervisor", "still having"]
+        wants_escalation = any(kw in user_message.lower() for kw in escalation_keywords)
+        
+        if wants_escalation:
+            # Try KB one more time with refined query
+            kb_result = query_kb(user_message)
+            
+            if kb_result["answer"] and kb_result["confidence"] >= 0.75:
+                response_text = f"Let me try a different solution:\n\n{kb_result['answer']}\n\nPlease try this and let me know if it works."
+                status = "in_progress"
+            else:
+                # Escalate to human
+                ticket_id = f"TKT-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+                conv["escalated"] = True
+                response_text = f"""I understand the previous solutions haven't resolved your issue. I'm escalating this to a human specialist.
+
+**Escalation Ticket:** {ticket_id}
+**Department:** {conv['ticket_info']['department']}
+**Priority:** High
+**Issue:** {conv['ticket_info']['classification']}
+
+A {conv['ticket_info']['department']} specialist will contact you within 2-4 business hours. They will have full access to our conversation history.
+
+Your ticket reference: {ticket_id}"""
+                status = "escalated"
+        else:
+            # Continue conversation with full context
+            context = f"""You are a helpful IT helpdesk AI agent. Provide clear, concise troubleshooting help.
+
+**Conversation Context:**
+- Initial Issue: {conv['ticket_info']['initial_query']}
+- Classification: {conv['ticket_info']['classification']}
+- Department: {conv['ticket_info']['department']}
+
+**Recent Conversation:**
+"""
+            # Include last 6 messages for context
+            for msg in conv["messages"][-6:]:
+                context += f"{msg['role'].upper()}: {msg['content']}\n"
+            
+            context += f"""\n**Current User Message:** {user_message}
+
+Instructions:
+- Provide helpful, specific guidance
+- If user confirms something worked, congratulate them
+- If unclear, ask clarifying questions
+- Keep responses concise (2-3 paragraphs max)
+- Don't repeat previous solutions
+
+Your response:"""
+            
+            try:
+                response_text = llm.invoke(context).content
+                status = "in_progress"
+            except Exception as e:
+                print(f"LLM error: {e}")
+                response_text = "I'm having trouble processing that. Could you rephrase your question?"
+                status = "in_progress"
+    
+    # Add assistant response
+    conv["messages"].append({
+        "role": "assistant",
+        "content": response_text,
+        "status": status,
+        "timestamp": datetime.now().isoformat()
+    })
+    
+    return {
+        "conversation_id": conversation_id,
+        "response": response_text,
+        "status": status,
+        "message_count": len(conv["messages"]),
+        "can_continue": not conv.get("escalated", False)
+    }
+
+def get_conversation_history(conversation_id: str) -> Optional[Dict[str, Any]]:
+    """Retrieve conversation history by ID."""
+    return conversations.get(conversation_id)
+
+# ===========================
+# Initialization
+# ===========================
+
+# Initialize KB collection on module load
+get_kb_collection()
+
+print("✅ Agent LangChain module loaded successfully")
+print(f"📊 KB Collection: {'initialized' if kb_collection else 'not initialized'}")
+if kb_collection:
+    print(f"📚 KB Records: {kb_collection.count()}")