Extract text and tool-call blocks from an AI message and render them. When streaming is enabled, text blocks are written to stdout immediately; otherwise they are accumulated in `state.full_response` for deferred output. Tool-call blocks are buffered and their names are printed to the
(
message_obj: AIMessage,
state: StreamState,
console: Console,
)
| 362 | |
| 363 | |
| 364 | def _process_ai_message( |
| 365 | message_obj: AIMessage, |
| 366 | state: StreamState, |
| 367 | console: Console, |
| 368 | ) -> None: |
| 369 | """Extract text and tool-call blocks from an AI message and render them. |
| 370 | |
| 371 | When streaming is enabled, text blocks are written to stdout immediately; |
| 372 | otherwise they are accumulated in `state.full_response` for deferred |
| 373 | output. Tool-call blocks are buffered and their names are printed to the |
| 374 | console. |
| 375 | |
| 376 | Args: |
| 377 | message_obj: The `AIMessage` received from the stream. |
| 378 | state: Stream state for accumulating response text and tool-call buffers. |
| 379 | console: Rich console for formatted output. |
| 380 | """ |
| 381 | # Extract token usage for stats accumulation |
| 382 | usage = getattr(message_obj, "usage_metadata", None) |
| 383 | if usage: |
| 384 | input_toks = usage.get("input_tokens", 0) |
| 385 | output_toks = usage.get("output_tokens", 0) |
| 386 | total_toks = usage.get("total_tokens", 0) |
| 387 | active_model = settings.model_name or "" |
| 388 | active_provider = settings.model_provider or "" |
| 389 | if input_toks or output_toks: |
| 390 | state.stats.record_request( |
| 391 | active_model, input_toks, output_toks, active_provider |
| 392 | ) |
| 393 | elif total_toks: |
| 394 | state.stats.record_request(active_model, total_toks, 0, active_provider) |
| 395 | |
| 396 | if not hasattr(message_obj, "content_blocks"): |
| 397 | logger.debug("AIMessage missing content_blocks attribute, skipping") |
| 398 | return |
| 399 | for block in message_obj.content_blocks: |
| 400 | if not isinstance(block, dict): |
| 401 | continue |
| 402 | block_type = block.get("type") |
| 403 | if block_type == "text": |
| 404 | text = block.get("text", "") |
| 405 | if text: |
| 406 | if state.stream: |
| 407 | if state.spinner: |
| 408 | state.spinner.stop() |
| 409 | _write_text(text) |
| 410 | state.full_response.append(text) |
| 411 | elif block_type in {"tool_call_chunk", "tool_call"}: |
| 412 | chunk_name = block.get("name") |
| 413 | chunk_id = block.get("id") |
| 414 | chunk_index = block.get("index") |
| 415 | |
| 416 | if chunk_index is not None: |
| 417 | buffer_key: int | str = chunk_index |
| 418 | elif chunk_id is not None: |
| 419 | buffer_key = chunk_id |
| 420 | else: |
| 421 | buffer_key = f"unknown-{len(state.tool_call_buffers)}" |
searching dependent graphs…