|
9 | 9 |
|
10 | 10 | import json |
11 | 11 | import os |
| 12 | +import traceback |
12 | 13 | from collections.abc import AsyncIterator |
13 | 14 | from textwrap import dedent |
14 | 15 | from typing import Any |
@@ -477,64 +478,76 @@ async def run_agent( |
477 | 478 | if not disable_tools: |
478 | 479 | stream_kwargs["tools"] = TOOLS # type: ignore[assignment] |
479 | 480 |
|
480 | | - async with client.messages.stream(**stream_kwargs) as stream: |
481 | | - current_tool_id: str | None = None |
482 | | - current_tool_name: str | None = None |
483 | | - current_tool_args = "" |
484 | | - |
485 | | - async for event in stream: |
486 | | - etype = type(event).__name__ |
487 | | - |
488 | | - if etype == "RawContentBlockStartEvent": |
489 | | - block = event.content_block # type: ignore[attr-defined] |
490 | | - if block.type == "text": |
491 | | - pass # streaming text chunks follow |
492 | | - elif block.type == "tool_use": |
493 | | - current_tool_id = block.id |
494 | | - current_tool_name = block.name |
495 | | - current_tool_args = "" |
496 | | - yield encoder.encode(ToolCallStartEvent( |
497 | | - type=EventType.TOOL_CALL_START, |
498 | | - tool_call_id=current_tool_id, |
499 | | - tool_call_name=current_tool_name, |
500 | | - parent_message_id=msg_id, |
501 | | - )) |
502 | | - |
503 | | - elif etype == "RawContentBlockDeltaEvent": |
504 | | - delta = event.delta # type: ignore[attr-defined] |
505 | | - if delta.type == "text_delta": |
506 | | - response_text += delta.text |
507 | | - yield encoder.encode(TextMessageContentEvent( |
508 | | - type=EventType.TEXT_MESSAGE_CONTENT, |
509 | | - message_id=msg_id, |
510 | | - delta=delta.text, |
511 | | - )) |
512 | | - elif delta.type == "input_json_delta": |
513 | | - current_tool_args += delta.partial_json |
514 | | - yield encoder.encode(ToolCallArgsEvent( |
515 | | - type=EventType.TOOL_CALL_ARGS, |
516 | | - tool_call_id=current_tool_id or "", |
517 | | - delta=delta.partial_json, |
518 | | - )) |
519 | | - |
520 | | - elif etype == "RawContentBlockStopEvent": |
521 | | - if current_tool_id and current_tool_name: |
522 | | - yield encoder.encode(ToolCallEndEvent( |
523 | | - type=EventType.TOOL_CALL_END, |
524 | | - tool_call_id=current_tool_id, |
525 | | - )) |
526 | | - try: |
527 | | - parsed_args = json.loads(current_tool_args) if current_tool_args else {} |
528 | | - except json.JSONDecodeError: |
529 | | - parsed_args = {} |
530 | | - tool_calls.append({ |
531 | | - "id": current_tool_id, |
532 | | - "name": current_tool_name, |
533 | | - "input": parsed_args, |
534 | | - }) |
535 | | - current_tool_id = None |
536 | | - current_tool_name = None |
537 | | - current_tool_args = "" |
| 481 | + try: |
| 482 | + async with client.messages.stream(**stream_kwargs) as stream: |
| 483 | + current_tool_id: str | None = None |
| 484 | + current_tool_name: str | None = None |
| 485 | + current_tool_args = "" |
| 486 | + |
| 487 | + async for event in stream: |
| 488 | + etype = type(event).__name__ |
| 489 | + |
| 490 | + if etype == "RawContentBlockStartEvent": |
| 491 | + block = event.content_block # type: ignore[attr-defined] |
| 492 | + if block.type == "text": |
| 493 | + pass # streaming text chunks follow |
| 494 | + elif block.type == "tool_use": |
| 495 | + current_tool_id = block.id |
| 496 | + current_tool_name = block.name |
| 497 | + current_tool_args = "" |
| 498 | + yield encoder.encode(ToolCallStartEvent( |
| 499 | + type=EventType.TOOL_CALL_START, |
| 500 | + tool_call_id=current_tool_id, |
| 501 | + tool_call_name=current_tool_name, |
| 502 | + parent_message_id=msg_id, |
| 503 | + )) |
| 504 | + |
| 505 | + elif etype == "RawContentBlockDeltaEvent": |
| 506 | + delta = event.delta # type: ignore[attr-defined] |
| 507 | + if delta.type == "text_delta": |
| 508 | + response_text += delta.text |
| 509 | + yield encoder.encode(TextMessageContentEvent( |
| 510 | + type=EventType.TEXT_MESSAGE_CONTENT, |
| 511 | + message_id=msg_id, |
| 512 | + delta=delta.text, |
| 513 | + )) |
| 514 | + elif delta.type == "input_json_delta": |
| 515 | + current_tool_args += delta.partial_json |
| 516 | + yield encoder.encode(ToolCallArgsEvent( |
| 517 | + type=EventType.TOOL_CALL_ARGS, |
| 518 | + tool_call_id=current_tool_id or "", |
| 519 | + delta=delta.partial_json, |
| 520 | + )) |
| 521 | + |
| 522 | + elif etype == "RawContentBlockStopEvent": |
| 523 | + if current_tool_id and current_tool_name: |
| 524 | + yield encoder.encode(ToolCallEndEvent( |
| 525 | + type=EventType.TOOL_CALL_END, |
| 526 | + tool_call_id=current_tool_id, |
| 527 | + )) |
| 528 | + try: |
| 529 | + parsed_args = json.loads(current_tool_args) if current_tool_args else {} |
| 530 | + except json.JSONDecodeError: |
| 531 | + parsed_args = {} |
| 532 | + tool_calls.append({ |
| 533 | + "id": current_tool_id, |
| 534 | + "name": current_tool_name, |
| 535 | + "input": parsed_args, |
| 536 | + }) |
| 537 | + current_tool_id = None |
| 538 | + current_tool_name = None |
| 539 | + current_tool_args = "" |
| 540 | + except Exception: |
| 541 | + # Surface the error as visible text in the chat so D5 |
| 542 | + # probes see a non-empty assistant response instead of a |
| 543 | + # silent broken SSE stream. Full traceback is logged |
| 544 | + # server-side by FastAPI's exception handler. |
| 545 | + err_text = f"Agent error: {traceback.format_exc()}" |
| 546 | + yield encoder.encode(TextMessageContentEvent( |
| 547 | + type=EventType.TEXT_MESSAGE_CONTENT, |
| 548 | + message_id=msg_id, |
| 549 | + delta=err_text, |
| 550 | + )) |
538 | 551 |
|
539 | 552 | yield encoder.encode(TextMessageEndEvent( |
540 | 553 | type=EventType.TEXT_MESSAGE_END, |
@@ -571,6 +584,7 @@ async def run_agent( |
571 | 584 | yield encoder.encode(ToolCallResultEvent( |
572 | 585 | type=EventType.TOOL_CALL_RESULT, |
573 | 586 | tool_call_id=tc["id"], |
| 587 | + message_id=f"{msg_id}-tool-result-{tc['id']}", |
574 | 588 | content=result_text, |
575 | 589 | )) |
576 | 590 | tool_results.append({ |
|
0 commit comments