|
8 | 8 | import re |
9 | 9 | from typing import Any |
10 | 10 |
|
11 | | -from agent_framework import ChatMessage |
| 11 | +from agent_framework import Message |
12 | 12 |
|
13 | 13 | from v4.config.settings import connection_config |
14 | 14 | from v4.models.messages import ( |
@@ -64,22 +64,22 @@ def _extract_tool_calls_from_contents(contents: list[Any]) -> list[AgentToolCall |
64 | 64 |
|
65 | 65 | def agent_response_callback( |
66 | 66 | agent_id: str, |
67 | | - message: ChatMessage, |
| 67 | + message: Message, |
68 | 68 | user_id: str | None = None, |
69 | 69 | ) -> None: |
70 | 70 | """ |
71 | | - Final (non-streaming) agent response callback using agent_framework ChatMessage. |
| 71 | + Final (non-streaming) agent response callback using agent_framework Message. |
72 | 72 | """ |
73 | 73 | agent_name = getattr(message, "author_name", None) or agent_id or "Unknown Agent" |
74 | 74 | role = getattr(message, "role", "assistant") |
75 | 75 |
|
76 | | - # FIX: Properly extract text from ChatMessage |
77 | | - # ChatMessage has a .text property that concatenates all TextContent items |
| 76 | + # FIX: Properly extract text from Message |
| 77 | + # Message has a .text property that concatenates all TextContent items |
78 | 78 | text = "" |
79 | | - if isinstance(message, ChatMessage): |
| 79 | + if isinstance(message, Message): |
80 | 80 | text = message.text # Use the property directly |
81 | 81 | else: |
82 | | - # Fallback for non-ChatMessage objects |
| 82 | + # Fallback for non-Message objects |
83 | 83 | text = str(getattr(message, "text", "")) |
84 | 84 |
|
85 | 85 | text = clean_citations(text or "") |
@@ -125,8 +125,8 @@ async def streaming_agent_response_callback( |
125 | 125 | # If text is None, don't fall back to str(update) as that would show object repr |
126 | 126 | # Just skip if there's no actual text content |
127 | 127 | if chunk_text is None: |
128 | | - # Check if update is a ChatMessage |
129 | | - if isinstance(update, ChatMessage): |
| 128 | + # Check if update is a Message |
| 129 | + if isinstance(update, Message): |
130 | 130 | chunk_text = update.text or "" |
131 | 131 | elif hasattr(update, "content"): |
132 | 132 | chunk_text = str(update.content) if update.content else "" |
|
0 commit comments