streaming: switch to stream_async + iter_chunks response drain

This commit is contained in:
daniel
2026-05-07 16:27:26 -05:00
parent 7f7f555983
commit 40a942b506
2 changed files with 32 additions and 32 deletions

View File

@@ -185,8 +185,8 @@ _current_actor_id: str = ''
@app.entrypoint @app.entrypoint
def main(payload: dict, context) -> dict: async def main(payload: dict, context):
"""Handle an invocation from agent-runner Lambda.""" """Handle an invocation from agent-runner Lambda (streaming)."""
global _current_actor_id global _current_actor_id
# Set up channel adapter # Set up channel adapter
@@ -267,15 +267,6 @@ def main(payload: dict, context) -> dict:
_code_interpreter.code_interpreter, home_assistant, connect_google_account, _code_interpreter.code_interpreter, home_assistant, connect_google_account,
manage_service] manage_service]
def _run_agent(tools):
agent = Agent(
model=model,
system_prompt=system_prompt,
session_manager=session_manager,
tools=tools,
)
return agent(payload.get('prompt', ''))
workspace_mcp_client = MCPClient( workspace_mcp_client = MCPClient(
lambda: streamablehttp_client(WORKSPACE_MCP_URL, timeout=20, auth=_SigV4HttpxAuth(actor_id=actor_id)) lambda: streamablehttp_client(WORKSPACE_MCP_URL, timeout=20, auth=_SigV4HttpxAuth(actor_id=actor_id))
) )
@@ -290,17 +281,26 @@ def main(payload: dict, context) -> dict:
else: else:
print(f'[main] actor={actor_id} has no google_email — skipping workspace_mcp') print(f'[main] actor={actor_id} has no google_email — skipping workspace_mcp')
agent = Agent(
model=model,
system_prompt=system_prompt,
session_manager=session_manager,
tools=base_tools + list(workspace_tools),
)
final_message = None
try: try:
result = _run_agent(base_tools + list(workspace_tools)) async for event in agent.stream_async(payload.get('prompt', '')):
if 'result' in event:
final_message = event['result'].message
yield event
finally: finally:
_typing_active = False _typing_active = False
# Flush buffered memory events
session_manager.close() session_manager.close()
# Deliver final response # Deliver final response if agent didn't call send_message
if not messaging.was_sent() and result.message: if not messaging.was_sent() and final_message:
msg = result.message msg = final_message
if isinstance(msg, dict): if isinstance(msg, dict):
content = msg.get('content', {}) content = msg.get('content', {})
if isinstance(content, dict): if isinstance(content, dict):
@@ -311,7 +311,5 @@ def main(payload: dict, context) -> dict:
msg = str(content) msg = str(content)
adapter.send(str(msg)) adapter.send(str(msg))
return {'result': result.message}
app.run() app.run()

View File

@@ -187,8 +187,10 @@ def handler(event, context):
payload=json.dumps(payload).encode(), payload=json.dumps(payload).encode(),
) )
# Consume streaming response (agent delivers to Telegram via send_message tool) # Drain streaming response body (agent delivers to Telegram via send_message tool)
for chunk in response.get('response', []): body = response.get('response')
pass # intentional no-op — agent handles delivery internally if body is not None:
for _ in body.iter_chunks():
pass
print(f"[agent-runner] Completed session={session_id} actor={actor_id}") print(f"[agent-runner] Completed session={session_id} actor={actor_id}")