1
1
from langchain_core .messages import AIMessage
2
2
from langgraph .graph import END , START , StateGraph
3
3
4
- from template_langgraph .agents .basic_workflow_agent .models import AgentInput , AgentOutput , AgentState
4
+ from template_langgraph .agents .basic_workflow_agent .models import AgentInput , AgentOutput , AgentState , Profile
5
5
from template_langgraph .llms .azure_openais import AzureOpenAiWrapper
6
6
from template_langgraph .loggers import get_logger
7
7
@@ -20,12 +20,14 @@ def create_graph(self):
20
20
# Create nodes
21
21
workflow .add_node ("initialize" , self .initialize )
22
22
workflow .add_node ("do_something" , self .do_something )
23
+ workflow .add_node ("extract_profile" , self .extract_profile )
23
24
workflow .add_node ("finalize" , self .finalize )
24
25
25
26
# Create edges
26
27
workflow .add_edge (START , "initialize" )
27
28
workflow .add_edge ("initialize" , "do_something" )
28
- workflow .add_edge ("do_something" , "finalize" )
29
+ workflow .add_edge ("do_something" , "extract_profile" )
30
+ workflow .add_edge ("extract_profile" , "finalize" )
29
31
workflow .add_edge ("finalize" , END )
30
32
31
33
# Compile the graph
@@ -55,6 +57,15 @@ def do_something(self, state: AgentState) -> AgentState:
55
57
56
58
return state
57
59
60
+ def extract_profile (self , state : AgentState ) -> AgentState :
61
+ """Extract profile information from the state."""
62
+ logger .info (f"Extracting profile from state: { state } " )
63
+ profile = self .llm .with_structured_output (Profile ).invoke (
64
+ input = state ["messages" ],
65
+ )
66
+ state ["profile" ] = profile
67
+ return state
68
+
58
69
def finalize (self , state : AgentState ) -> AgentState :
59
70
"""Finalize the agent's work and prepare the output."""
60
71
logger .info (f"Finalizing BasicWorkflowAgent with state: { state } " )
@@ -75,7 +86,10 @@ def run_agent(self, input: AgentInput) -> AgentOutput:
75
86
}
76
87
final_state = app .invoke (initial_state )
77
88
logger .info (f"Final state after running agent: { final_state } " )
78
- return AgentOutput (response = final_state ["messages" ][- 1 ].content )
89
+ return AgentOutput (
90
+ response = final_state ["messages" ][- 1 ].content ,
91
+ profile = final_state ["profile" ],
92
+ )
79
93
80
94
def draw_mermaid_png (self ) -> bytes :
81
95
"""Draw the graph in Mermaid format."""
0 commit comments