github/autogen
1
1
Fork 0
mirror of https://github.com/microsoft/autogen.git synced 2026-02-12 14:55:10 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add token usage to messages (#4028)

Browse Source 
* Add token usage to messages

* small test edit
This commit is contained in:
Eric Zhu
2024-11-01 13:20:25 -07:00
committed by GitHub
parent e9c16fe22e
commit ca7caa779d
3 changed files with 28 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
View File

@@ -266,8 +266,8 @@ class AssistantAgent(BaseChatAgent):
while isinstance(result.content, list) and all(isinstance(item, FunctionCall) for item in result.content):
event_logger.debug(ToolCallEvent(tool_calls=result.content, source=self.name))
# Add the tool call message to the output.
inner_messages.append(ToolCallMessage(content=result.content, source=self.name))
yield ToolCallMessage(content=result.content, source=self.name)
inner_messages.append(ToolCallMessage(content=result.content, source=self.name, model_usage=result.usage))
yield ToolCallMessage(content=result.content, source=self.name, model_usage=result.usage)
# Execute the tool calls.
results = await asyncio.gather(
@@ -303,7 +303,8 @@ class AssistantAgent(BaseChatAgent):
assert isinstance(result.content, str)
yield Response(
chat_message=TextMessage(content=result.content, source=self.name), inner_messages=inner_messages
chat_message=TextMessage(content=result.content, source=self.name, model_usage=result.usage),
inner_messages=inner_messages,
)
async def _execute_tool_call(

5
python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
View File

@@ -1,7 +1,7 @@
from typing import List
from autogen_core.components import FunctionCall, Image
from autogen_core.components.models import FunctionExecutionResult
from autogen_core.components.models import FunctionExecutionResult, RequestUsage
from pydantic import BaseModel
@@ -11,6 +11,9 @@ class BaseMessage(BaseModel):
source: str
"""The name of the agent that sent this message."""
model_usage: RequestUsage | None = None
"""The model client usage incurred when producing this message."""
class TextMessage(BaseMessage):
"""A text message."""