Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
+from typing import List, Dict, Literal
 from langgraph.graph import StateGraph, START
 from langgraph.types import Command
 from dotenv import load_dotenv
-from typing import Literal
 import gradio as gr
 from langchain_huggingface import HuggingFaceEndpoint
 
@@ -10,15 +10,7 @@ from langchain_huggingface import HuggingFaceEndpoint
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Define HuggingFaceEndpoint for urgency classification
-classifier = HuggingFaceEndpoint(
-    repo_id="mistralai/Mistral-7B-Instruct-v0.3",  # You can use any model fine-tuned for text classification
-    huggingfacehub_api_token=HF_TOKEN.strip(),
-    temperature=0.5,
-    max_new_tokens=10,  # For short classification outputs
-)
-
-# Define HuggingFaceEndpoint for response generation
+# Define HuggingFaceEndpoint
 llm = HuggingFaceEndpoint(
     repo_id="mistralai/Mistral-7B-Instruct-v0.3",
     huggingfacehub_api_token=HF_TOKEN.strip(),
@@ -26,87 +18,84 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=200,
 )
 
-# Define state
+# Define state class for managing the ticket state
 class State(dict):
-    issue: str
+    issue_description: str
     priority: str
-    response: str
     escalation_needed: bool
+    response: str
+    escalation: str
 
 # Create the graph
 builder = StateGraph(State)
 
-# Define nodes (agents)
+# Define the nodes for the multi-agent flow
+
+# Ticket creation agent
 def ticket_creation_agent(state: State) -> Command[Literal["priority_classification_agent"]]:
-    """Create a ticket based on customer issue."""
-    return Command(update={"issue": state["issue"]}, goto="priority_classification_agent")
+    """Capture ticket description and proceed to priority classification."""
+    return Command(update={"issue_description": state["issue_description"]}, goto="priority_classification_agent")
 
-def priority_classification_agent(state: State) -> Command[Literal["response_generation_agent"]]:
-    """Classify ticket priority based on NLP analysis of the issue."""
-    issue = state["issue"]
+# Priority classification agent
+def priority_classification_agent(state: State) -> Command[Literal["escalation_classification_agent"]]:
+    """Classify priority based on the issue description."""
+    prompt = f"Classify the following issue as urgent, critical, or normal: {state['issue_description']}"
+    priority = llm.invoke(prompt)  # Use .invoke() instead of __call__
     
-    # Use the model to classify urgency level
-    prompt = f"Classify the urgency of this issue: {issue}"
-    priority = classifier(prompt)
-    
-    # Use the classification result to assign priority
-    if "urgent" in priority.lower():
-        priority = "High"
-    elif "critical" in priority.lower():
-        priority = "Critical"
+    # Update the priority and proceed to escalation classification
+    return Command(update={"priority": priority}, goto="escalation_classification_agent")
+
+# Escalation classification agent
+def escalation_classification_agent(state: State) -> Command[Literal["generate_response_agent"]]:
+    """Classify whether escalation is needed based on priority."""
+    if state["priority"].lower() in ["urgent", "critical"]:
+        escalation_needed = True
     else:
-        priority = "Normal"
-    
-    return Command(update={"priority": priority}, goto="response_generation_agent")
+        escalation_needed = False
 
-def response_generation_agent(state: State) -> Command[Literal["escalation_agent"]]:
-    """Generate response based on issue and priority."""
-    prompt = f"""
-    Customer Issue: {state['issue']}
-    Priority: {state['priority']}
-    
-    You are a customer service representative. Provide a detailed response to this issue based on its priority.
-    """
-    response = llm(prompt)
-    return Command(update={"response": response, "escalation_needed": state['priority'] == "High"}, goto="escalation_agent")
+    return Command(update={"escalation_needed": escalation_needed}, goto="generate_response_agent")
+
+# Generate response agent
+def generate_response_agent(state: State) -> Dict[str, str]:
+    """Generate response based on ticket priority and escalation need."""
+    if state["escalation_needed"]:
+        escalation = "Escalate the issue to a senior team member immediately."
+    else:
+        escalation = "No escalation needed."
 
-def escalation_agent(state: State):
-    """Decide if the ticket needs to be escalated based on priority."""
-    escalation = "Yes" if state["escalation_needed"] else "No"
-    return {"response": state["response"], "escalation": escalation}
+    prompt = f"Generate a response for the following issue: {state['issue_description']}. The priority is {state['priority']}."
+    response = llm.invoke(prompt)  # Use .invoke() instead of __call__
+
+    return {
+        "response": response,
+        "escalation": escalation
+    }
 
 # Add nodes to the graph
 builder.add_edge(START, "ticket_creation_agent")
 builder.add_node("ticket_creation_agent", ticket_creation_agent)
 builder.add_node("priority_classification_agent", priority_classification_agent)
-builder.add_node("response_generation_agent", response_generation_agent)
-builder.add_node("escalation_agent", escalation_agent)
+builder.add_node("escalation_classification_agent", escalation_classification_agent)
+builder.add_node("generate_response_agent", generate_response_agent)
 
 # Compile the graph
 graph = builder.compile()
 
-# Gradio Interface
-def process_ticket(issue: str):
-    """Run the multi-agent customer support flow with user input."""
-    state = {"issue": issue}
+# Gradio Interface function to process the ticket
+def process_ticket(issue_description: str):
+    """Process the issue ticket through the multi-agent flow."""
+    state = {"issue_description": issue_description}
     result = graph.invoke(state)
 
-    # Return the customer service response and escalation info
-    return result["response"], result["escalation"]
+    # Return the response and escalation result
+    return result.get("response", "No response generated"), result.get("escalation", "No escalation specified")
 
+# Gradio Interface
 iface = gr.Interface(
     fn=process_ticket,
-    inputs=[
-        gr.Textbox(
-            label="Enter Customer Issue",
-            placeholder="Describe the issue that needs to be addressed...",
-        ),
-    ],
-    outputs=[
-        gr.Textbox(label="Generated Customer Service Response"),
-        gr.Textbox(label="Escalation Status"),
-    ],
-    title="Customer Support Multi-Agent System",
+    inputs=gr.Textbox(label="Describe the issue"),
+    outputs=[gr.Textbox(label="Response"), gr.Textbox(label="Escalation Decision")],
+    title="Ticket Handling System",
 )
 
 iface.launch()