Chore: improvements for ticket classification sample

samples/ticket-classification/README.md

@@ -1,6 +1,6 @@
 # Support Ticket Classification System
 
-Use LangGraph with Azure OpenAI to automatically classify support tickets into predefined categories with confidence scores. UiPath Orchestrator API integration for human approval step.
+Use LangGraph with Azure OpenAI to automatically classify support tickets into predefined categories with confidence scores. UiPath Action Center integration for human approval step.
 
 ## Debug
 
@@ -52,7 +52,8 @@ The input ticket should be in the following format:
 ```json
 {
     "message": "The ticket message or description",
-    "ticket_id": "Unique ticket identifier"
+    "ticket_id": "Unique ticket identifier",
+    "assignee"[optional]: "username or email of the person assigned to handle escalations"
 }
 ```
 

samples/ticket-classification/agent.mermaid

@@ -1,16 +1,19 @@
-%%{init: {'flowchart': {'curve': 'linear'}}}%%
+---
+config:
+  flowchart:
+    curve: linear
+---
 graph TD;
 	__start__([<p>__start__</p>]):::first
 	classify(classify)
-	create_action(create_action)
-	human_approval(human_approval)
+	human_approval_node(human_approval_node)
 	notify_team(notify_team)
 	__end__([<p>__end__</p>]):::last
 	__start__ --> classify;
-	classify --> create_action;
-	create_action --> human_approval;
-	human_approval --> notify_team;
+	classify --> human_approval_node;
 	notify_team --> __end__;
+	human_approval_node -.-> classify;
+	human_approval_node -.-> notify_team;
 	classDef default fill:#f2f0ff,line-height:1.2
 	classDef first fill-opacity:0
 	classDef last fill:#bfb6fc

samples/ticket-classification/main.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 from langchain_openai import AzureChatOpenAI
 from langchain_core.output_parsers import PydanticOutputParser
@@ -10,14 +10,15 @@
 from pydantic import BaseModel, Field
 
 from uipath_sdk import UiPathSDK
-
+from uipath_sdk._models import CreateAction
 logger = logging.getLogger(__name__)
 
 uipath = UiPathSDK()
 
 class GraphInput(BaseModel):
     message: str
     ticket_id: str
+    assignee: Optional[str]
 
 class GraphOutput(BaseModel):
     label: str
@@ -26,9 +27,11 @@ class GraphOutput(BaseModel):
 class GraphState(BaseModel):
     message: str
     ticket_id: str
+    assignee: Optional[str] = None
     label: Optional[str] = None
     confidence: Optional[float] = None
-
+    predicted_categories: List[str] = []
+    human_approval: Optional[bool] = None
 
 class TicketClassification(BaseModel):
     label: Literal["security", "error", "system", "billing", "performance"] = Field(
@@ -78,6 +81,10 @@ def get_azure_openai_api_key() -> str:
 
     return api_key
 
+def decide_next_node(state: GraphState) -> Literal["classify", "notify_team"]:
+    if state.human_approval is True:
+        return "notify_team"
+    return "classify"
 
 async def classify(state: GraphState) -> GraphState:
     """Classify the support ticket using LLM."""
@@ -87,51 +94,77 @@ async def classify(state: GraphState) -> GraphState:
         azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
         api_version="2024-10-21"
     )
+    new_state = GraphState(
+        message=state.message,
+        ticket_id=state.ticket_id,
+        assignee=state.assignee,
+        predicted_categories=state.predicted_categories.copy(),
+        human_approval=state.human_approval
+    )
+
+    if len(new_state.predicted_categories) > 0:
+        prompt.append(("user", f"The ticket is 100% not part of the following categories '{new_state.predicted_categories}'. Choose another one."))
+
     _prompt = prompt.partial(
         format_instructions=output_parser.get_format_instructions()
     )
     chain = _prompt | llm | output_parser
 
     try:
-        result = await chain.ainvoke({"ticket_text": state.message})
-        print(result)
-        state.label = result.label
-        state.confidence = result.confidence
+        result = await chain.ainvoke({"ticket_text": new_state.message})
+        new_state.label = result.label
+        new_state.predicted_categories.append(result.label)
+        new_state.confidence = result.confidence
         logger.info(
             f"Ticket classified with label: {result.label} confidence score: {result.confidence}"
         )
-        return state
+        return new_state
     except Exception as e:
         logger.error(f"Classification failed: {str(e)}")
-        state.label = "error"
-        state.confidence = 0.0
-        return state
+        return GraphState(
+            message=new_state.message,
+            ticket_id=new_state.ticket_id,
+            assignee=new_state.assignee,
+            predicted_categories=new_state.predicted_categories,
+            human_approval=new_state.human_approval,
+            label="error",
+            confidence=0.0
+        )
 
 async def wait_for_human(state: GraphState) -> GraphState:
     logger.info("Wait for human approval")
-    feedback = interrupt(f"Label: {state.label} Confidence: {state.confidence}")
-
-    if isinstance(feedback, bool) and feedback is True:
-        return Command(goto="notify_team")
-    else:
-        return Command(goto=END)
+    action_data = interrupt(CreateAction(name="escalation_agent_app",
+                                         title="Action Required: Review classification",
+                                         data={
+                                             "AgentOutput": f"This is how I classified the ticket: '{state.ticket_id}', with message '{state.message}' \n Label: '{state.label}' Confidence: '{state.confidence}'",
+                                             "AgentName": "ticket-classification "},
+                                         app_version=1,
+                                         assignee=state.assignee,
+                                         ))
+    new_state = GraphState(
+        message=state.message,
+        ticket_id=state.ticket_id,
+        assignee=state.assignee,
+        predicted_categories=state.predicted_categories.copy(),
+        human_approval=isinstance(action_data["Answer"], bool) and action_data["Answer"] is True
+    )
+    return new_state
 
-async def notify_team(state: GraphState) -> GraphState:
+async def notify_team(state: GraphState) -> GraphOutput:
     logger.info("Send team email notification")
-    print(state)
-    return state
+    return GraphOutput(label=state.label, confidence=state.confidence)
 
 """Process a support ticket through the workflow."""
 
 builder = StateGraph(GraphState, input=GraphInput, output=GraphOutput)
 
 builder.add_node("classify", classify)
-builder.add_node("human_approval", wait_for_human)
+builder.add_node("human_approval_node", wait_for_human)
 builder.add_node("notify_team", notify_team)
 
 builder.add_edge(START, "classify")
-builder.add_edge("classify", "human_approval")
-builder.add_edge("human_approval", "notify_team")
+builder.add_edge("classify", "human_approval_node")
+builder.add_conditional_edges("human_approval_node", decide_next_node)
 builder.add_edge("notify_team", END)