88import os
99import logging
1010
11- from ag_ui .core import RunAgentInput , RunErrorEvent , RunFinishedEvent
11+ from ag_ui .core import RunAgentInput , RunErrorEvent
1212from bedrock_agentcore .identity .auth import requires_access_token
1313from bedrock_agentcore .runtime import BedrockAgentCoreApp , RequestContext
1414from copilotkit import CopilotKitMiddleware , LangGraphAGUIAgent
@@ -123,45 +123,6 @@ def _build_checkpointer() -> AgentCoreMemorySaver:
123123 region_name = os .environ .get ("AWS_DEFAULT_REGION" , "us-east-1" ),
124124 )
125125
126-
127- async def create_langgraph_agent (tools : list ):
128- try :
129- return create_agent (
130- model = _build_model (streaming = True ),
131- tools = [* tools , query_data , * todo_tools ], # MCP tools + data + todo tools
132- checkpointer = _build_checkpointer (),
133- middleware = [CopilotKitMiddleware ()],
134- system_prompt = SYSTEM_PROMPT ,
135- state_schema = AgentState , # extends BaseAgentState with todos: list[Todo]
136- )
137- except Exception :
138- logging .exception ("Error creating LangGraph agent" )
139- raise
140-
141-
142- async def create_runtime_graph ():
143- mcp_client = await create_gateway_mcp_client ()
144- tools = await mcp_client .get_tools ()
145- return await create_langgraph_agent (tools )
146-
147-
148- class ActorAwareLangGraphAgent (LangGraphAGUIAgent ):
149- async def run (self , input : RunAgentInput ):
150- actor_id = (
151- self .config .get ("configurable" , {}).get ("actor_id" ) if self .config else None
152- )
153- if not actor_id :
154- raise ValueError (
155- "Missing actor identity. Provide forwardedProps.actor_id/user_id "
156- "or include sub claim in the bearer token."
157- )
158-
159- self .graph = await create_runtime_graph ()
160- self .config = {"configurable" : {"actor_id" : actor_id }}
161- async for event in super ().run (input ):
162- yield event
163-
164-
165126@app .entrypoint
166127async def invocations (payload : dict , context : RequestContext ):
167128 input_data = RunAgentInput .model_validate (payload )
@@ -176,14 +137,24 @@ async def invocations(payload: dict, context: RequestContext):
176137 "or include sub claim in the bearer token."
177138 )
178139
179- request_agent = ActorAwareLangGraphAgent (
180- name = "LangGraphSingleAgent" ,
181- description = "LangGraph single agent exposed via AG-UI" ,
182- graph = None ,
183- config = {"configurable" : {"actor_id" : actor_id }},
184- )
185-
186140 try :
141+ mcp_client = await create_gateway_mcp_client ()
142+ tools = await mcp_client .get_tools ()
143+ graph = await create_agent (
144+ model = _build_model (streaming = True ),
145+ tools = [* tools , query_data , * todo_tools ], # MCP tools + data + todo tools
146+ checkpointer = _build_checkpointer (),
147+ middleware = [CopilotKitMiddleware ()],
148+ system_prompt = SYSTEM_PROMPT ,
149+ state_schema = AgentState , # extends BaseAgentState with todos: list[Todo]
150+ )
151+
152+ request_agent = LangGraphAGUIAgent (
153+ name = "LangGraphSingleAgent" ,
154+ description = "LangGraph single agent exposed via AG-UI" ,
155+ graph = graph ,
156+ config = {"configurable" : {"actor_id" : actor_id }},
157+ )
187158 async for event in request_agent .run (input_data ):
188159 if event is not None :
189160 yield event .model_dump (mode = "json" , by_alias = True , exclude_none = True )
0 commit comments