Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
thought="Calling pass function",
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"pass",
"TERMINATE",
Expand Down Expand Up @@ -142,18 +143,21 @@ async def test_run_with_tools_and_reflection() -> None:
content=[FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function")],
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
CreateResult(
finish_reason="stop",
content="Hello",
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
CreateResult(
finish_reason="stop",
content="TERMINATE",
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
],
model_info={
Expand Down Expand Up @@ -244,6 +248,7 @@ async def test_run_with_parallel_tools() -> None:
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
thought="Calling pass and echo functions",
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"pass",
"TERMINATE",
Expand Down Expand Up @@ -329,6 +334,7 @@ async def test_run_with_parallel_tools_with_empty_call_ids() -> None:
],
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"pass",
"TERMINATE",
Expand Down Expand Up @@ -411,6 +417,7 @@ async def test_handoffs() -> None:
],
usage=RequestUsage(prompt_tokens=42, completion_tokens=43),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
)
],
model_info={
Expand Down Expand Up @@ -557,6 +564,7 @@ async def test_list_chat_messages(monkeypatch: pytest.MonkeyPatch) -> None:
content="Response to message 1",
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
)
]
)
Expand Down Expand Up @@ -759,6 +767,7 @@ async def test_model_client_stream_with_tool_calls() -> None:
finish_reason="function_calls",
usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"Example response 2 to task",
]
Expand Down
3 changes: 3 additions & 0 deletions python/packages/autogen-agentchat/tests/test_group_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,7 @@ async def test_round_robin_group_chat_with_tools(runtime: AgentRuntime | None) -
content=[FunctionCall(id="1", name="pass", arguments=json.dumps({"input": "pass"}))],
usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"Hello",
"TERMINATE",
Expand Down Expand Up @@ -848,6 +849,7 @@ async def test_swarm_handoff_using_tool_calls(runtime: AgentRuntime | None) -> N
content=[FunctionCall(id="1", name="handoff_to_agent2", arguments=json.dumps({}))],
usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"Hello",
"TERMINATE",
Expand Down Expand Up @@ -938,6 +940,7 @@ async def test_swarm_with_parallel_tool_calls(runtime: AgentRuntime | None) -> N
],
usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
cached=False,
raw_response={"id": "mock-id", "provider": "replay"},
),
"Hello",
"TERMINATE",
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from dataclasses import dataclass
from typing import List, Literal, Optional, Union
from typing import List, Literal, Optional, Union, Dict, Any

from pydantic import BaseModel, Field
from typing_extensions import Annotated
Expand Down Expand Up @@ -125,3 +125,8 @@ class CreateResult(BaseModel):
thought: Optional[str] = None
"""The reasoning text for the completion if available. Used for reasoning models
and additional text content besides function calls."""

raw_response: Optional[Dict[str, Any]] = None
"""Raw response from the model API, useful for custom field access."""


2 changes: 2 additions & 0 deletions python/packages/autogen-core/tests/test_tool_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,13 +113,15 @@ async def create(
usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
cached=False,
logprobs=None,
raw_response={"id": "mock-id", "provider": "replay"},
)
return CreateResult(
content="Done",
finish_reason="stop",
usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
cached=False,
logprobs=None,
raw_response={"id": "mock-id", "provider": "replay"},
)

def create_stream(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,7 @@ async def create(
finish_reason=data.get("finish_reason", "stop"),
usage=data.get("usage", RequestUsage(prompt_tokens=0, completion_tokens=0)),
cached=True,
raw_response=data.get("raw_response", {"id": "mock-id", "provider": "replay"}),
)
return result

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -557,6 +557,7 @@ async def create(
usage=usage,
cached=False,
thought=thought,
raw_response=result,
)

# Update usage statistics
Expand Down Expand Up @@ -775,6 +776,7 @@ async def create_stream(
usage=usage,
cached=False,
thought=thought,
raw_response=result,
)

# Emit the end event.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -438,6 +438,7 @@ async def create(
usage=usage,
cached=False,
thought=thought,
raw_response=result,
)

self.add_usage(usage)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,7 @@ async def create(
if not response_tool_calls and not response_text:
logger.debug("DEBUG: No response text found. Returning empty response.")
return CreateResult(
content="", usage=RequestUsage(prompt_tokens=0, completion_tokens=0), finish_reason="stop", cached=False
content="", usage=RequestUsage(prompt_tokens=0, completion_tokens=0), finish_reason="stop", cached=False, raw_response=response
)

# Create a CreateResult object
Expand All @@ -373,6 +373,7 @@ async def create(
usage=cast(RequestUsage, response["usage"]),
finish_reason=normalize_stop_reason(finish_reason), # type: ignore
cached=False,
raw_response=response,
)

# If we are running in the context of a handler we can get the agent_id
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -602,6 +602,7 @@ async def create(
usage=usage,
cached=False,
logprobs=None,
raw_response=result,
)

self._total_usage = _add_usage(self._total_usage, usage)
Expand Down Expand Up @@ -737,6 +738,7 @@ async def create_stream(
usage=usage,
cached=False,
logprobs=None,
raw_response=result,
)

# Emit the end event.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -653,6 +653,7 @@ async def create(
cached=False,
logprobs=logprobs,
thought=thought,
raw_response=result,
)

self._total_usage = _add_usage(self._total_usage, usage)
Expand Down Expand Up @@ -879,6 +880,7 @@ async def create_stream(
cached=False,
logprobs=logprobs,
thought=thought,
raw_response=result,
)

# Log the end of the stream.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ async def create(
_, output_token_count = self._tokenize(response)
self._cur_usage = RequestUsage(prompt_tokens=prompt_token_count, completion_tokens=output_token_count)
response = CreateResult(
finish_reason="stop", content=response, usage=self._cur_usage, cached=self._cached_bool_value
finish_reason="stop", content=response, usage=self._cur_usage, cached=self._cached_bool_value, raw_response=response
)
else:
self._cur_usage = RequestUsage(
Expand Down Expand Up @@ -221,7 +221,7 @@ async def create_stream(
else:
yield token
yield CreateResult(
finish_reason="stop", content=response, usage=self._cur_usage, cached=self._cached_bool_value
finish_reason="stop", content=response, usage=self._cur_usage, cached=self._cached_bool_value, raw_response=response
)
self._update_total_usage()
else:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -521,6 +521,7 @@ async def create(
usage=RequestUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
cached=False,
thought=thought,
raw_response=result
)

@staticmethod
Expand Down Expand Up @@ -676,6 +677,7 @@ async def create_stream(
finish_reason="function_calls",
usage=RequestUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
cached=False,
raw_response=result,
)
return

Expand All @@ -698,6 +700,7 @@ async def create_stream(
usage=RequestUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens),
cached=False,
thought=thought,
raw_response=result,
)

# Emit the end event.
Expand Down
15 changes: 3 additions & 12 deletions python/uv.lock
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't want to delete this file. I meant we should keep this file unchanged.

Copy link
Contributor Author

@tongyu0924 tongyu0924 Apr 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't want to delete this file. I meant we should keep this file unchanged.

I’ve reverted uv.lock to match main exactly.
Could you please confirm that everything looks correct? Thanks!

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.