Update agent_langchain.py

agent_langchain.py

@@ -1,28 +1,29 @@
 import os
 import requests
 import torch
+import time
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import chromadb
 from chromadb.config import Settings
-from chromadb.utils import embedding_functions
-from langchain.agents import initialize_agent, Tool
-from langchain.agents import AgentType
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.agents import initialize_agent, Tool, AgentType
 from langchain.memory import ConversationBufferMemory
 
-# -------------------------------
-# Environment & URLs
-# -------------------------------
+# ==============================================================
+# 🌐 ENVIRONMENT & GLOBAL SETTINGS
+# ==============================================================
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
 ROUTING_URL = os.environ.get("ROUTING_URL")  # Space 2 URL
 SPACE_URL = os.environ.get("SPACE_URL", "http://localhost:7860")
 
+# Hugging Face Space writable paths
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface"
-# -------------------------------
-# Label Dictionary
-# -------------------------------
+
+# ==============================================================
+# 🏷️ LABEL DICTIONARY
+# ==============================================================
 LABEL_DICTIONARY = {
     "I1": "Low Impact",
     "I2": "Medium Impact",
@@ -39,32 +40,15 @@ LABEL_DICTIONARY = {
     "T5": "Question"
 }
 
-# -------------------------------
-# Load Classification Model
-# -------------------------------
+# ==============================================================
+# 🤖 LOAD CLASSIFICATION MODEL
+# ==============================================================
 clf_model_name = "DavinciTech/BERT_Categorizer"
 clf_tokenizer = AutoTokenizer.from_pretrained(clf_model_name)
 clf_model = AutoModelForSequenceClassification.from_pretrained(clf_model_name)
 
-# -------------------------------
-# Initialize ChromaDB Client for KB
-# -------------------------------
-# ✅ Use new API — persistent on Hugging Face writable directory
-chroma_client = chromadb.PersistentClient(path="/tmp/chroma")
-
-# ✅ Create or get your KB collection
-kb_collection = chroma_client.get_or_create_collection("Knowledge_Base")
-
-COLLECTION_NAME = "Knowledge_Base"
-try:
-    kb_collection = chroma_client.get_collection(COLLECTION_NAME)
-except:
-    kb_collection = None
-
-# -------------------------------
-# Classification Function
-# -------------------------------
 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]
@@ -79,10 +63,11 @@ def classify_ticket(text):
         "type": LABEL_DICTIONARY[f"T{type_idx}"]
     }
 
-# -------------------------------
-# Routing Function
-# -------------------------------
+# ==============================================================
+# 🧭 ROUTING FUNCTION (Space 2)
+# ==============================================================
 def call_routing(text, retries=3, delay=1):
+    """Call Space 2 routing endpoint."""
     url = ROUTING_URL if ROUTING_URL else f"{SPACE_URL}/route"
     for attempt in range(retries):
         try:
@@ -96,89 +81,86 @@ def call_routing(text, retries=3, delay=1):
             else:
                 return "General IT"
 
-# -------------------------------
-# KB Query
-# -------------------------------
+# ==============================================================
+# 📚 KNOWLEDGE BASE SETUP
+# ==============================================================
+# Persistent Chroma client (new API)
+chroma_client = chromadb.PersistentClient(path="/tmp/chroma")
+
+COLLECTION_NAME = "knowledge_base"
+
+try:
+    kb_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)
+except Exception as e:
+    kb_collection = None
+    print("⚠️ Could not initialize KB:", e)
+
 def query_kb(text, top_k=1):
+    """Query the knowledge base for relevant solutions."""
     if not kb_collection:
-        return {"answer": "⚠️ KB not set up. Call /setup first.", "confidence": 0.0}
+        return {"answer": "⚠️ KB not set up.", "confidence": 0.0}
 
     results = kb_collection.query(query_texts=[text], n_results=top_k)
-    if not results or len(results['documents'][0]) == 0:
+    if not results or not results.get("documents") or len(results["documents"][0]) == 0:
         return {"answer": "No relevant KB found.", "confidence": 0.0}
 
     return {
-        "answer": results['documents'][0][0],
-        "confidence": results['distances'][0][0] if results.get('distances') else 0.0,
-        "metadata": results['metadatas'][0][0] if results['metadatas'][0] else {}
+        "answer": results["documents"][0][0],
+        "confidence": results.get("distances", [[0.0]])[0][0],
+        "metadata": results.get("metadatas", [[{}]])[0][0]
     }
 
-# -------------------------------
-# Gemini LLM Wrapper
-# -------------------------------
-class GeminiLLM:
-    def __init__(self, api_key=GEMINI_API_KEY):
-        self.api_key = api_key
-        self.api_url = GEMINI_API_URL
-
-    def __call__(self, prompt: str):
-        if not self.api_key:
-            return {"text": "⚠️ Gemini API key not set."}
-        payload = {"contents": [{"parts": [{"text": prompt}]}]}
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        try:
-            resp = requests.post(self.api_url, json=payload, headers=headers)
-            resp.raise_for_status()
-            data = resp.json()
-            text = data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
-            return text
-        except:
-            return "⚠️ Gemini API call failed."
-
-# -------------------------------
-# Define LangChain Tools
-# -------------------------------
+# ==============================================================
+# 🧠 GEMINI LLM (Official LangChain Integration)
+# ==============================================================
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.5-flash",
+    temperature=0.3,
+    google_api_key=GEMINI_API_KEY
+)
+
+# ==============================================================
+# 🧰 DEFINE LANGCHAIN TOOLS
+# ==============================================================
 tools = [
     Tool(
         name="TicketClassifier",
         func=lambda text: classify_ticket(text),
-        description="Classifies a ticket into impact, urgency, and type. Mandatory tool."
+        description="Classifies the ticket into impact, urgency, and type. Mandatory tool."
     ),
     Tool(
         name="RoutingTool",
         func=lambda text: call_routing(text),
-        description="Assigns a department for the ticket via Space 2. Mandatory tool."
+        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 KB for relevant solution. Returns answer text."
+        description="Searches the KB for relevant solutions. Returns a descriptive answer."
     )
 ]
 
-# -------------------------------
-# Initialize Memory
-# -------------------------------
+# ==============================================================
+# 💬 MEMORY & AGENT INITIALIZATION
+# ==============================================================
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-# -------------------------------
-# Initialize Agent
-# -------------------------------
 agent_executor = initialize_agent(
     tools=tools,
-    llm=GeminiLLM(),
+    llm=llm,
     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
     memory=memory,
     verbose=False
 )
 
-# -------------------------------
-# Process Ticket Function
-# -------------------------------
+# ==============================================================
+# 🧾 MAIN TICKET PROCESSOR
+# ==============================================================
 def process_ticket_langchain(ticket_text):
+    """Full pipeline: classify → route → query KB → decide KB vs Gemini."""
     reasoning_trace = []
 
-    # Step 1: Classifier
+    # Step 1: Classification
     classification = classify_ticket(ticket_text)
     reasoning_trace.append(f"[Classifier] Impact: {classification['impact']}, Urgency: {classification['urgency']}, Type: {classification['type']}")
 
@@ -188,26 +170,27 @@ def process_ticket_langchain(ticket_text):
 
     # Step 3: KB Search
     kb_result = query_kb(ticket_text)
-    reasoning_trace.append(f"[KB Search] Top answer: '{kb_result['answer']}' (confidence: {kb_result['confidence']})")
+    reasoning_trace.append(f"[KB Search] Top Answer: '{kb_result['answer']}' (confidence: {kb_result['confidence']})")
 
-    # Step 4: Decision KB vs LLM
+    # Step 4: KB vs LLM Decision
     if kb_result["confidence"] >= 0.75:
         final_answer = kb_result["answer"]
         status = "resolved"
-        reasoning_trace.append("[Decision] KB confidence high → ticket resolved via KB.")
+        reasoning_trace.append("[Decision] High KB confidence → ticket resolved via KB.")
     else:
         llm_prompt = f"""
-You are a professional IT helpdesk assistant.
+You are a professional IT helpdesk agent.
 A user submitted the following ticket: "{ticket_text}"
+
 Ticket classification: {classification}
 Assigned department: {department}
-KB Search result: {kb_result['answer']} (confidence: {kb_result['confidence']})
+Knowledge base result: {kb_result['answer']} (confidence: {kb_result['confidence']})
 
-Provide a professional and descriptive solution or guidance based on this information.
+Please provide a clear, descriptive, and professional IT helpdesk response.
 """
-        final_answer = GeminiLLM()(llm_prompt)
+        final_answer = llm.invoke(llm_prompt).content
         status = "escalated"
-        reasoning_trace.append("[Decision] KB confidence low → ticket escalated via Gemini LLM.")
+        reasoning_trace.append("[Decision] Low KB confidence → fallback to Gemini LLM for escalation.")
 
     return {
         "status": status,