Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Dict, Literal
+from typing import Dict, Literal
 from langgraph.graph import StateGraph, START
 from langgraph.types import Command
 from dotenv import load_dotenv
@@ -40,7 +40,7 @@ def ticket_creation_agent(state: State) -> Command[Literal["priority_classificat
 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__
+    priority = llm.invoke(prompt).strip()  # Use .invoke() instead of __call__
     
     # Update the priority and proceed to escalation classification
     return Command(update={"priority": priority}, goto="escalation_classification_agent")
@@ -48,28 +48,17 @@ def priority_classification_agent(state: State) -> Command[Literal["escalation_c
 # 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:
-        escalation_needed = False
-
+    escalation_needed = state["priority"].lower() in ["urgent", "critical"]
     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."
-
+    escalation = "Escalate the issue to a senior team member immediately." if state["escalation_needed"] else "No escalation needed."
     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__
+    response = llm.invoke(prompt).strip()  # Use .invoke() and ensure the response is clean
 
-    return {
-        "response": response,
-        "escalation": escalation
-    }
+    return {"response": response, "escalation": escalation}
 
 # Add nodes to the graph
 builder.add_edge(START, "ticket_creation_agent")
@@ -85,17 +74,24 @@ graph = builder.compile()
 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 response and escalation result
-    return result.get("response", "No response generated"), result.get("escalation", "No escalation specified")
+    try:
+        result = graph.invoke(state)
+        response = result.get("response", "No response generated")
+        escalation = result.get("escalation", "No escalation specified")
+        return response, escalation
+    except Exception as e:
+        return f"Error occurred: {e}", "Unable to determine escalation"
 
 # Gradio Interface
 iface = gr.Interface(
     fn=process_ticket,
     inputs=gr.Textbox(label="Describe the issue"),
-    outputs=[gr.Textbox(label="Response"), gr.Textbox(label="Escalation Decision")],
+    outputs=[
+        gr.Textbox(label="Response"),
+        gr.Textbox(label="Escalation Decision"),
+    ],
     title="Ticket Handling System",
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()