Python: Fix response metadata construction (#6955)
* Python: Fix response metadata construction Propagate complete AgentResponse metadata through core response construction and provider finalization paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python: Refine Ollama response metadata parsing Filter Ollama usage details to real token counts and only propagate streaming finish metadata from final chunks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python: Address response metadata review comments Accumulate Copilot non-streaming usage events and keep structured response value parsing lazy for provider hooks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
GitHub
parent
7f551f057a
commit
e26c8591ae
@@ -23,6 +23,7 @@ from agent_framework import (
|
||||
Message,
|
||||
ResponseStream,
|
||||
ToolTypes,
|
||||
UsageDetails,
|
||||
load_settings,
|
||||
normalize_messages,
|
||||
normalize_tools,
|
||||
@@ -377,6 +378,7 @@ class RawClaudeAgent(BaseAgent, Generic[OptionsT]):
|
||||
self._default_options = opts
|
||||
self._started = False
|
||||
self._current_session_id: str | None = None
|
||||
self._structured_output: Any = None
|
||||
|
||||
def _normalize_tools(
|
||||
self,
|
||||
@@ -685,8 +687,7 @@ class RawClaudeAgent(BaseAgent, Generic[OptionsT]):
|
||||
Returns:
|
||||
An AgentResponse with structured_output set as value if present.
|
||||
"""
|
||||
structured_output = getattr(self, "_structured_output", None)
|
||||
return AgentResponse.from_updates(updates, value=structured_output)
|
||||
return AgentResponse.from_updates(updates, value=self._structured_output)
|
||||
|
||||
@overload
|
||||
def run(
|
||||
@@ -823,6 +824,34 @@ class RawClaudeAgent(BaseAgent, Generic[OptionsT]):
|
||||
raise AgentException(f"Claude API error: {error_msg}")
|
||||
session_id = message.session_id
|
||||
structured_output = message.structured_output
|
||||
usage = message.usage or {}
|
||||
input_tokens = usage.get("input_tokens")
|
||||
output_tokens = usage.get("output_tokens")
|
||||
total_token_count = (
|
||||
input_tokens + output_tokens
|
||||
if isinstance(input_tokens, int) and isinstance(output_tokens, int)
|
||||
else None
|
||||
)
|
||||
usage_details = UsageDetails(**{
|
||||
key: value
|
||||
for key, value in {
|
||||
"input_token_count": input_tokens,
|
||||
"output_token_count": output_tokens,
|
||||
"total_token_count": total_token_count,
|
||||
"cache_creation_input_token_count": usage.get("cache_creation_input_tokens"),
|
||||
"cache_read_input_token_count": usage.get("cache_read_input_tokens"),
|
||||
}.items()
|
||||
if isinstance(value, int)
|
||||
})
|
||||
finish_reason = message.stop_reason
|
||||
if usage_details or finish_reason:
|
||||
yield AgentResponseUpdate(
|
||||
contents=[Content.from_usage(usage_details, raw_representation=message)]
|
||||
if usage_details
|
||||
else None,
|
||||
finish_reason=cast(Any, finish_reason),
|
||||
raw_representation=message,
|
||||
)
|
||||
|
||||
# Update session with session ID
|
||||
if session_id:
|
||||
|
||||
@@ -271,6 +271,55 @@ class TestClaudeAgentRun:
|
||||
await agent.run("Hello", session=session)
|
||||
assert session.service_session_id == "test-session-id"
|
||||
|
||||
async def test_run_captures_result_message_usage_and_finish_reason(self) -> None:
|
||||
"""Test that ResultMessage metadata is propagated to the final AgentResponse."""
|
||||
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
|
||||
from claude_agent_sdk.types import StreamEvent
|
||||
|
||||
messages = [
|
||||
StreamEvent(
|
||||
event={
|
||||
"type": "content_block_delta",
|
||||
"delta": {"type": "text_delta", "text": "Response"},
|
||||
},
|
||||
uuid="event-1",
|
||||
session_id="test-session-id",
|
||||
),
|
||||
AssistantMessage(
|
||||
content=[TextBlock(text="Response")],
|
||||
model="claude-sonnet",
|
||||
),
|
||||
ResultMessage(
|
||||
subtype="success",
|
||||
duration_ms=100,
|
||||
duration_api_ms=50,
|
||||
is_error=False,
|
||||
num_turns=1,
|
||||
session_id="test-session-id",
|
||||
stop_reason="end_turn",
|
||||
usage={
|
||||
"input_tokens": 42,
|
||||
"output_tokens": 18,
|
||||
"cache_creation_input_tokens": 3,
|
||||
"cache_read_input_tokens": 5,
|
||||
},
|
||||
),
|
||||
]
|
||||
mock_client = self._create_mock_client(messages)
|
||||
|
||||
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
|
||||
agent = ClaudeAgent()
|
||||
response = await agent.run("Hello")
|
||||
|
||||
assert response.finish_reason == "end_turn"
|
||||
assert response.usage_details == {
|
||||
"input_token_count": 42,
|
||||
"output_token_count": 18,
|
||||
"total_token_count": 60,
|
||||
"cache_creation_input_token_count": 3,
|
||||
"cache_read_input_token_count": 5,
|
||||
}
|
||||
|
||||
async def test_run_with_session(self) -> None:
|
||||
"""Test run with existing session."""
|
||||
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
|
||||
@@ -378,6 +427,51 @@ class TestClaudeAgentRunStream:
|
||||
assert updates[0].text == "Streaming "
|
||||
assert updates[1].text == "response"
|
||||
|
||||
async def test_run_stream_final_response_captures_usage_and_finish_reason(self) -> None:
|
||||
"""Test run(stream=True) final response includes ResultMessage metadata."""
|
||||
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
|
||||
from claude_agent_sdk.types import StreamEvent
|
||||
|
||||
messages = [
|
||||
StreamEvent(
|
||||
event={
|
||||
"type": "content_block_delta",
|
||||
"delta": {"type": "text_delta", "text": "Streaming response"},
|
||||
},
|
||||
uuid="event-1",
|
||||
session_id="stream-session",
|
||||
),
|
||||
AssistantMessage(
|
||||
content=[TextBlock(text="Streaming response")],
|
||||
model="claude-sonnet",
|
||||
),
|
||||
ResultMessage(
|
||||
subtype="success",
|
||||
duration_ms=100,
|
||||
duration_api_ms=50,
|
||||
is_error=False,
|
||||
num_turns=1,
|
||||
session_id="stream-session",
|
||||
stop_reason="max_tokens",
|
||||
usage={"input_tokens": 7, "output_tokens": 9},
|
||||
),
|
||||
]
|
||||
mock_client = self._create_mock_client(messages)
|
||||
|
||||
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
|
||||
agent = ClaudeAgent()
|
||||
stream = agent.run("Hello", stream=True)
|
||||
async for _ in stream:
|
||||
pass
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert response.finish_reason == "max_tokens"
|
||||
assert response.usage_details == {
|
||||
"input_token_count": 7,
|
||||
"output_token_count": 9,
|
||||
"total_token_count": 16,
|
||||
}
|
||||
|
||||
async def test_run_stream_raises_on_assistant_message_error(self) -> None:
|
||||
"""Test run raises AgentException when AssistantMessage has an error."""
|
||||
from agent_framework.exceptions import AgentException
|
||||
|
||||
@@ -50,6 +50,7 @@ from ._types import (
|
||||
ChatResponseUpdate,
|
||||
Message,
|
||||
ResponseStream,
|
||||
_build_agent_response_from_chat_response, # pyright: ignore[reportPrivateUsage]
|
||||
map_chat_to_agent_update,
|
||||
normalize_messages,
|
||||
)
|
||||
@@ -1082,24 +1083,28 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):
|
||||
if not response:
|
||||
raise AgentInvalidResponseException("Chat client did not return a response.")
|
||||
|
||||
await self._finalize_response(
|
||||
response=response,
|
||||
agent_name=context["agent_name"],
|
||||
session=context["session"],
|
||||
session_context=context["session_context"],
|
||||
for message in response.messages:
|
||||
if message.author_name is None:
|
||||
message.author_name = context["agent_name"]
|
||||
|
||||
session = context["session"]
|
||||
if (
|
||||
session
|
||||
and response.conversation_id
|
||||
and not is_local_history_conversation_id(response.conversation_id)
|
||||
and session.service_session_id != response.conversation_id
|
||||
):
|
||||
session.service_session_id = response.conversation_id
|
||||
|
||||
agent_response = _build_agent_response_from_chat_response(
|
||||
response,
|
||||
response_format=context["chat_options"].get("response_format"),
|
||||
suppress_response_id=context["suppress_response_id"],
|
||||
)
|
||||
return AgentResponse(
|
||||
messages=response.messages,
|
||||
response_id=None if context["suppress_response_id"] else response.response_id,
|
||||
created_at=response.created_at,
|
||||
usage_details=response.usage_details,
|
||||
value=response.value,
|
||||
response_format=context["chat_options"].get("response_format"),
|
||||
continuation_token=response.continuation_token,
|
||||
raw_representation=response,
|
||||
additional_properties=response.additional_properties,
|
||||
)
|
||||
session_context = context["session_context"]
|
||||
session_context._response = agent_response # type: ignore[assignment]
|
||||
await self._run_after_providers(session=session, context=session_context)
|
||||
return agent_response
|
||||
|
||||
def _parse_streaming_response(
|
||||
self,
|
||||
@@ -1127,12 +1132,10 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):
|
||||
):
|
||||
session.service_session_id = conversation_id
|
||||
|
||||
suppress_response_id = context["suppress_response_id"]
|
||||
session_context = context["session_context"]
|
||||
session_context._response = AgentResponse( # type: ignore[assignment]
|
||||
messages=response.messages,
|
||||
response_id=None if suppress_response_id else response.response_id,
|
||||
)
|
||||
if context["suppress_response_id"]:
|
||||
response.response_id = None
|
||||
session_context._response = response # type: ignore[assignment]
|
||||
await self._run_after_providers(session=session, context=session_context)
|
||||
|
||||
def _propagate_conversation_id(update: AgentResponseUpdate) -> AgentResponseUpdate:
|
||||
@@ -1430,48 +1433,6 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]):
|
||||
"function_invocation_kwargs": additional_function_arguments,
|
||||
}
|
||||
|
||||
async def _finalize_response(
|
||||
self,
|
||||
response: ChatResponse,
|
||||
agent_name: str,
|
||||
session: AgentSession | None,
|
||||
session_context: SessionContext,
|
||||
suppress_response_id: bool = False,
|
||||
) -> None:
|
||||
"""Finalize response by setting author names and running after_run providers.
|
||||
|
||||
Args:
|
||||
response: The chat response to finalize.
|
||||
agent_name: The name of the agent to set as author.
|
||||
session: The conversation session.
|
||||
session_context: The invocation context.
|
||||
suppress_response_id: When True, omit the raw service response ID from the public response.
|
||||
"""
|
||||
# Ensure that the author name is set for each message in the response.
|
||||
for message in response.messages:
|
||||
if message.author_name is None:
|
||||
message.author_name = agent_name
|
||||
|
||||
# Propagate conversation_id back to session (e.g. thread ID from Assistants API).
|
||||
# For Responses-style APIs this can rotate every turn (response_id-based continuation),
|
||||
# so refresh when a newer value is returned.
|
||||
if (
|
||||
session
|
||||
and response.conversation_id
|
||||
and not is_local_history_conversation_id(response.conversation_id)
|
||||
and session.service_session_id != response.conversation_id
|
||||
):
|
||||
session.service_session_id = response.conversation_id
|
||||
|
||||
# Set the response on the context for after_run providers
|
||||
session_context._response = AgentResponse( # type: ignore[assignment]
|
||||
messages=response.messages,
|
||||
response_id=None if suppress_response_id else response.response_id,
|
||||
)
|
||||
|
||||
# Run after_run providers (reverse order)
|
||||
await self._run_after_providers(session=session, context=session_context)
|
||||
|
||||
async def _prepare_session_and_messages(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -29,7 +29,13 @@ from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, cast
|
||||
|
||||
from ._feature_stage import ExperimentalFeature, experimental
|
||||
from ._middleware import ChatContext, ChatMiddleware
|
||||
from ._types import AgentResponse, ChatResponse, Message, ResponseStream
|
||||
from ._types import (
|
||||
AgentResponse,
|
||||
ChatResponse,
|
||||
Message,
|
||||
ResponseStream,
|
||||
_build_agent_response_from_chat_response, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
from .exceptions import ChatClientInvalidResponseException
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -634,9 +640,9 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
|
||||
response: ChatResponse,
|
||||
) -> None:
|
||||
"""Persist a single model-call response through the configured history providers."""
|
||||
service_call_context._response = AgentResponse( # type: ignore[assignment]
|
||||
messages=response.messages,
|
||||
response_id=None,
|
||||
service_call_context._response = _build_agent_response_from_chat_response( # type: ignore[assignment]
|
||||
response,
|
||||
suppress_response_id=True,
|
||||
)
|
||||
for provider in reversed(self._providers):
|
||||
await provider.after_run(
|
||||
|
||||
@@ -2782,6 +2782,30 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
return self.text
|
||||
|
||||
|
||||
def _build_agent_response_from_chat_response( # pyright: ignore[reportUnusedFunction]
|
||||
response: ChatResponse[Any],
|
||||
*,
|
||||
response_format: StructuredResponseFormat = None,
|
||||
suppress_response_id: bool = False,
|
||||
) -> AgentResponse[Any]:
|
||||
"""Build the AgentResponse wrapper for a completed ChatResponse."""
|
||||
agent_response = AgentResponse(
|
||||
messages=response.messages,
|
||||
response_id=None if suppress_response_id else response.response_id,
|
||||
created_at=response.created_at,
|
||||
finish_reason=cast(FinishReasonLiteral | FinishReason | None, response.finish_reason),
|
||||
usage_details=response.usage_details,
|
||||
response_format=response_format,
|
||||
continuation_token=response.continuation_token,
|
||||
raw_representation=response,
|
||||
additional_properties=response.additional_properties,
|
||||
)
|
||||
if response._value_parsed: # pyright: ignore[reportPrivateUsage]
|
||||
agent_response._value = response._value # pyright: ignore[reportPrivateUsage]
|
||||
agent_response._value_parsed = True # pyright: ignore[reportPrivateUsage]
|
||||
return agent_response
|
||||
|
||||
|
||||
# region AgentResponseUpdate
|
||||
|
||||
|
||||
|
||||
@@ -122,6 +122,24 @@ class _ResponseIdRecordingHistoryProvider(_RecordingHistoryProvider):
|
||||
await super().after_run(agent=agent, session=session, context=context, state=state)
|
||||
|
||||
|
||||
class _ResponseMetadataRecordingHistoryProvider(_RecordingHistoryProvider):
|
||||
def __init__(self, source_id: str = "recording_response_metadata") -> None:
|
||||
super().__init__(source_id=source_id)
|
||||
self.responses: list[AgentResponse] = []
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
if context.response:
|
||||
self.responses.append(context.response)
|
||||
await super().after_run(agent=agent, session=session, context=context, state=state)
|
||||
|
||||
|
||||
def test_agent_session_type(agent_session: AgentSession) -> None:
|
||||
assert isinstance(agent_session, AgentSession)
|
||||
|
||||
@@ -493,6 +511,64 @@ async def test_chat_agent_persists_history_per_service_call(
|
||||
assert session.service_session_id is None
|
||||
|
||||
|
||||
async def test_per_service_call_history_provider_receives_full_agent_response_metadata(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
provider = _ResponseMetadataRecordingHistoryProvider()
|
||||
|
||||
@tool(name="lookup_weather", approval_mode="never_require")
|
||||
def lookup_weather(location: str) -> str:
|
||||
return f"Weather in {location}: sunny"
|
||||
|
||||
session = AgentSession()
|
||||
session.state[provider.source_id] = {"messages": []}
|
||||
first_response = ChatResponse(
|
||||
messages=Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1",
|
||||
name="lookup_weather",
|
||||
arguments='{"location": "Seattle"}',
|
||||
)
|
||||
],
|
||||
),
|
||||
response_id="resp_call_1",
|
||||
created_at="2026-07-07T12:02:00Z",
|
||||
finish_reason="tool_calls",
|
||||
usage_details={"input_token_count": 5, "output_token_count": 1, "total_token_count": 6},
|
||||
continuation_token=cast(Any, {"token": "psc-next"}),
|
||||
additional_properties={"provider": "psc-test"},
|
||||
)
|
||||
final_response = ChatResponse(
|
||||
messages=Message(role="assistant", contents=["It is sunny in Seattle."]),
|
||||
response_id="resp_call_2",
|
||||
finish_reason="stop",
|
||||
usage_details={"input_token_count": 6, "output_token_count": 4, "total_token_count": 10},
|
||||
)
|
||||
chat_client_base.run_responses = [first_response, final_response] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
|
||||
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[lookup_weather],
|
||||
context_providers=[provider],
|
||||
require_per_service_call_history_persistence=True,
|
||||
)
|
||||
|
||||
result = await agent.run("What's the weather in Seattle?", session=session)
|
||||
|
||||
assert result.text == "It is sunny in Seattle."
|
||||
assert len(provider.responses) == 2
|
||||
captured = provider.responses[0]
|
||||
assert captured.response_id is None
|
||||
assert captured.created_at == "2026-07-07T12:02:00Z"
|
||||
assert captured.finish_reason == "tool_calls"
|
||||
assert captured.usage_details == {"input_token_count": 5, "output_token_count": 1, "total_token_count": 6}
|
||||
assert captured.continuation_token == {"token": "psc-next"}
|
||||
assert captured.additional_properties == {"provider": "psc-test"}
|
||||
assert captured.raw_representation is first_response
|
||||
|
||||
|
||||
async def test_chat_agent_persists_history_per_service_call_streaming(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
@@ -1182,6 +1258,144 @@ async def test_chat_agent_context_providers_after_run(
|
||||
assert mock_provider.last_service_session_id == "test-thread-id"
|
||||
|
||||
|
||||
async def test_context_provider_receives_full_agent_response_non_streaming(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
"""Context providers should see the same full AgentResponse returned by non-streaming runs."""
|
||||
captured_response: AgentResponse | None = None
|
||||
|
||||
class CapturingContextProvider(ContextProvider):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(source_id="capture_response")
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
nonlocal captured_response
|
||||
captured_response = context.response
|
||||
|
||||
raw_response = ChatResponse(
|
||||
messages=[Message(role="assistant", contents=[Content.from_text('{"answer": "ok"}')])],
|
||||
response_id="resp_full",
|
||||
created_at="2026-07-07T12:00:00Z",
|
||||
finish_reason="stop",
|
||||
usage_details={"input_token_count": 3, "output_token_count": 2, "total_token_count": 5},
|
||||
continuation_token=cast(Any, {"token": "next"}),
|
||||
additional_properties={"provider": "test"},
|
||||
)
|
||||
chat_client_base.run_responses = [raw_response] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
|
||||
agent = Agent(client=chat_client_base, context_providers=[CapturingContextProvider()])
|
||||
|
||||
response = await agent.run(
|
||||
"Hello",
|
||||
options={"response_format": {"type": "object", "properties": {"answer": {"type": "string"}}}},
|
||||
)
|
||||
|
||||
assert captured_response is response
|
||||
assert response.response_id == "resp_full"
|
||||
assert response.created_at == "2026-07-07T12:00:00Z"
|
||||
assert response.finish_reason == "stop"
|
||||
assert response.usage_details == {"input_token_count": 3, "output_token_count": 2, "total_token_count": 5}
|
||||
assert response.value == {"answer": "ok"}
|
||||
assert response.continuation_token == {"token": "next"}
|
||||
assert response.additional_properties == {"provider": "test"}
|
||||
assert response.raw_representation is raw_response
|
||||
|
||||
|
||||
async def test_context_provider_after_run_preserves_lazy_structured_value_parsing(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
"""Context providers should see the response before malformed structured output raises for callers."""
|
||||
captured_response: AgentResponse | None = None
|
||||
|
||||
class CapturingContextProvider(ContextProvider):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(source_id="capture_lazy_value_response")
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
nonlocal captured_response
|
||||
captured_response = context.response
|
||||
|
||||
raw_response = ChatResponse(messages=[Message(role="assistant", contents=[Content.from_text("not-json")])])
|
||||
chat_client_base.run_responses = [raw_response] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
|
||||
agent = Agent(client=chat_client_base, context_providers=[CapturingContextProvider()])
|
||||
|
||||
response = await agent.run("Hello", options={"response_format": {"type": "object"}})
|
||||
|
||||
assert captured_response is response
|
||||
assert response.text == "not-json"
|
||||
with pytest.raises(ValueError, match="not valid JSON"):
|
||||
_ = response.value
|
||||
|
||||
|
||||
async def test_context_provider_receives_full_agent_response_streaming(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
"""Context providers should see the same full AgentResponse finalized by streaming runs."""
|
||||
captured_response: AgentResponse | None = None
|
||||
raw_update = ChatResponseUpdate(
|
||||
contents=[
|
||||
Content.from_text('{"answer": "ok"}'),
|
||||
Content.from_usage({"input_token_count": 4, "output_token_count": 3, "total_token_count": 7}),
|
||||
],
|
||||
role="assistant",
|
||||
response_id="resp_stream_full",
|
||||
created_at="2026-07-07T12:01:00Z",
|
||||
finish_reason="stop",
|
||||
continuation_token=cast(Any, {"token": "stream-next"}),
|
||||
additional_properties={"provider": "stream-test"},
|
||||
)
|
||||
|
||||
class CapturingContextProvider(ContextProvider):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(source_id="capture_stream_response")
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
nonlocal captured_response
|
||||
captured_response = context.response
|
||||
|
||||
chat_client_base.streaming_responses = [[raw_update]] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
|
||||
agent = Agent(client=chat_client_base, context_providers=[CapturingContextProvider()])
|
||||
|
||||
stream = agent.run(
|
||||
"Hello",
|
||||
stream=True,
|
||||
options={"response_format": {"type": "object", "properties": {"answer": {"type": "string"}}}},
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert captured_response is response
|
||||
assert response.response_id == "resp_stream_full"
|
||||
assert response.created_at == "2026-07-07T12:01:00Z"
|
||||
assert response.finish_reason == "stop"
|
||||
assert response.usage_details == {"input_token_count": 4, "output_token_count": 3, "total_token_count": 7}
|
||||
assert response.value == {"answer": "ok"}
|
||||
assert response.continuation_token == {"token": "stream-next"}
|
||||
assert response.additional_properties == {"provider": "stream-test"}
|
||||
assert response.raw_representation == [raw_update]
|
||||
|
||||
|
||||
async def test_chat_agent_context_providers_messages_adding(
|
||||
client: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
|
||||
@@ -9,12 +9,7 @@ import logging
|
||||
import sys
|
||||
import warnings
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence
|
||||
from typing import Any, ClassVar, Generic, Literal, TypedDict, overload
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
from typing import Any, ClassVar, Generic, Literal, TypedDict, cast, overload
|
||||
|
||||
from agent_framework import (
|
||||
AgentMiddlewareLayer,
|
||||
@@ -29,6 +24,8 @@ from agent_framework import (
|
||||
Message,
|
||||
ResponseStream,
|
||||
SessionContext,
|
||||
UsageDetails,
|
||||
add_usage_details,
|
||||
normalize_messages,
|
||||
)
|
||||
from agent_framework._settings import load_settings
|
||||
@@ -37,6 +34,15 @@ from agent_framework._types import AgentRunInputs, normalize_tools
|
||||
from agent_framework.exceptions import AgentException
|
||||
from agent_framework.observability import AgentTelemetryLayer
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
|
||||
try:
|
||||
from copilot import CopilotClient, CopilotSession, RuntimeConnection
|
||||
from copilot.generated.rpc import PermissionDecisionUserNotAvailable
|
||||
@@ -49,7 +55,7 @@ try:
|
||||
SessionHooks,
|
||||
SystemMessageConfig,
|
||||
)
|
||||
from copilot.session_events import PermissionRequest, SessionEvent, SessionEventType
|
||||
from copilot.session_events import AssistantUsageData, PermissionRequest, SessionEvent, SessionEventType
|
||||
from copilot.tools import Tool as CopilotTool
|
||||
from copilot.tools import ToolInvocation, ToolResult
|
||||
except ImportError as _copilot_import_error:
|
||||
@@ -58,12 +64,6 @@ except ImportError as _copilot_import_error:
|
||||
"Please use Python 3.11 or later."
|
||||
) from _copilot_import_error
|
||||
|
||||
if sys.version_info >= (3, 13):
|
||||
from typing import TypeVar # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import TypeVar # pragma: no cover
|
||||
|
||||
|
||||
DEFAULT_TIMEOUT_SECONDS: float = 60.0
|
||||
"""Default timeout in seconds for Copilot requests."""
|
||||
|
||||
@@ -563,6 +563,27 @@ class RawGitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
|
||||
)
|
||||
return self._run_impl(messages=messages, session=session, options=options)
|
||||
|
||||
@staticmethod
|
||||
def _parse_usage_details_from_copilot(data: AssistantUsageData) -> UsageDetails | None:
|
||||
total_token_count = (
|
||||
data.input_tokens + data.output_tokens
|
||||
if data.input_tokens is not None and data.output_tokens is not None
|
||||
else None
|
||||
)
|
||||
usage_details = UsageDetails(**{
|
||||
key: value
|
||||
for key, value in {
|
||||
"input_token_count": data.input_tokens,
|
||||
"output_token_count": data.output_tokens,
|
||||
"total_token_count": total_token_count,
|
||||
"cache_read_input_token_count": data.cache_read_tokens,
|
||||
"cache_creation_input_token_count": data.cache_write_tokens,
|
||||
"reasoning_output_token_count": data.reasoning_tokens,
|
||||
}.items()
|
||||
if value is not None
|
||||
})
|
||||
return usage_details or None
|
||||
|
||||
async def _run_impl(
|
||||
self,
|
||||
messages: AgentRunInputs | None = None,
|
||||
@@ -602,6 +623,27 @@ class RawGitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
|
||||
opts["tools"] = existing + list(session_context.tools)
|
||||
|
||||
copilot_session = await self._get_or_create_session(session, streaming=False, runtime_options=opts)
|
||||
usage_details: UsageDetails | None = None
|
||||
finish_reason: str | None = None
|
||||
model: str | None = None
|
||||
|
||||
def usage_event_handler(event: SessionEvent) -> None:
|
||||
nonlocal usage_details, finish_reason, model
|
||||
if event.type != SessionEventType.ASSISTANT_USAGE:
|
||||
return
|
||||
if isinstance(event.data, AssistantUsageData):
|
||||
parsed_usage_details = self._parse_usage_details_from_copilot(event.data)
|
||||
if parsed_usage_details:
|
||||
usage_details = add_usage_details(usage_details, parsed_usage_details)
|
||||
if event.data.finish_reason:
|
||||
finish_reason = event.data.finish_reason
|
||||
if event.data.model:
|
||||
model = event.data.model
|
||||
else:
|
||||
logger.warning(
|
||||
"Ignoring GitHub Copilot assistant usage event with unexpected payload type: %s",
|
||||
type(event.data).__name__,
|
||||
)
|
||||
|
||||
# Build the prompt from the full set of messages in the session context,
|
||||
# so that any context/history provider-injected messages are included.
|
||||
@@ -610,10 +652,13 @@ class RawGitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
|
||||
if session_context.instructions:
|
||||
prompt = "\n".join(session_context.instructions) + "\n" + prompt
|
||||
|
||||
unsubscribe = copilot_session.on(usage_event_handler)
|
||||
try:
|
||||
response_event = await copilot_session.send_and_wait(prompt, timeout=timeout)
|
||||
except Exception as ex:
|
||||
raise AgentException(f"GitHub Copilot request failed: {ex}") from ex
|
||||
finally:
|
||||
unsubscribe()
|
||||
|
||||
response_messages: list[Message] = []
|
||||
response_id: str | None = None
|
||||
@@ -635,7 +680,13 @@ class RawGitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
|
||||
)
|
||||
response_id = message_id
|
||||
|
||||
response = AgentResponse(messages=response_messages, response_id=response_id)
|
||||
response = AgentResponse(
|
||||
messages=response_messages,
|
||||
response_id=response_id,
|
||||
finish_reason=cast(Any, finish_reason),
|
||||
usage_details=usage_details,
|
||||
additional_properties={"model": model} if model else None,
|
||||
)
|
||||
session_context._response = response # type: ignore[assignment]
|
||||
await self._run_after_providers(session=session, context=session_context)
|
||||
return response
|
||||
@@ -721,6 +772,26 @@ class RawGitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
|
||||
raw_representation=event,
|
||||
)
|
||||
queue.put_nowait(update)
|
||||
elif event.type == SessionEventType.ASSISTANT_USAGE:
|
||||
if not isinstance(event.data, AssistantUsageData):
|
||||
logger.warning(
|
||||
"Ignoring GitHub Copilot assistant usage event with unexpected payload type: %s",
|
||||
type(event.data).__name__,
|
||||
)
|
||||
return
|
||||
usage_details = self._parse_usage_details_from_copilot(event.data)
|
||||
finish_reason = event.data.finish_reason or None
|
||||
model = event.data.model or None
|
||||
if usage_details or finish_reason or model:
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_usage(usage_details, raw_representation=event.data)]
|
||||
if usage_details
|
||||
else None,
|
||||
finish_reason=cast(Any, finish_reason),
|
||||
additional_properties={"model": model} if model else None,
|
||||
raw_representation=event,
|
||||
)
|
||||
queue.put_nowait(update)
|
||||
elif event.type == SessionEventType.TOOL_EXECUTION_START:
|
||||
tool_call_id = getattr(event.data, "tool_call_id", None) or ""
|
||||
tool_name = getattr(event.data, "tool_name", None) or ""
|
||||
|
||||
@@ -27,6 +27,7 @@ from agent_framework import (
|
||||
from agent_framework.exceptions import AgentException
|
||||
from copilot.session import PermissionHandler, PreToolUseHookInput
|
||||
from copilot.session_events import (
|
||||
AssistantUsageData,
|
||||
Data,
|
||||
SessionEvent,
|
||||
SessionEventType,
|
||||
@@ -459,6 +460,68 @@ class TestGitHubCopilotAgentRun:
|
||||
|
||||
assert isinstance(response, AgentResponse)
|
||||
|
||||
async def test_run_captures_assistant_usage_event(
|
||||
self,
|
||||
mock_client: MagicMock,
|
||||
mock_session: MagicMock,
|
||||
assistant_message_event: SessionEvent,
|
||||
) -> None:
|
||||
"""Test non-streaming run captures assistant usage metadata from session events."""
|
||||
usage_data = AssistantUsageData(
|
||||
model="claude-sonnet-4-5",
|
||||
input_tokens=120,
|
||||
output_tokens=40,
|
||||
cache_read_tokens=7,
|
||||
cache_write_tokens=3,
|
||||
reasoning_tokens=11,
|
||||
finish_reason="stop",
|
||||
)
|
||||
usage_event = SessionEvent(
|
||||
data=usage_data,
|
||||
id=uuid4(),
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
type=SessionEventType.ASSISTANT_USAGE,
|
||||
)
|
||||
second_usage_event = SessionEvent(
|
||||
data=AssistantUsageData(
|
||||
model="gpt-5.1-mini",
|
||||
input_tokens=5,
|
||||
output_tokens=2,
|
||||
finish_reason="length",
|
||||
),
|
||||
id=uuid4(),
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
type=SessionEventType.ASSISTANT_USAGE,
|
||||
)
|
||||
usage_handler: Any = None
|
||||
|
||||
def mock_on(handler: Any) -> Any:
|
||||
nonlocal usage_handler
|
||||
usage_handler = handler
|
||||
return lambda: None
|
||||
|
||||
async def mock_send_and_wait(*args: Any, **kwargs: Any) -> SessionEvent:
|
||||
usage_handler(usage_event)
|
||||
usage_handler(second_usage_event)
|
||||
return assistant_message_event
|
||||
|
||||
mock_session.on = mock_on
|
||||
mock_session.send_and_wait = AsyncMock(side_effect=mock_send_and_wait)
|
||||
|
||||
agent = GitHubCopilotAgent(client=mock_client)
|
||||
response = await agent.run("Hello")
|
||||
|
||||
assert response.finish_reason == "length"
|
||||
assert response.usage_details == {
|
||||
"input_token_count": 125,
|
||||
"output_token_count": 42,
|
||||
"total_token_count": 167,
|
||||
"cache_read_input_token_count": 7,
|
||||
"cache_creation_input_token_count": 3,
|
||||
"reasoning_output_token_count": 11,
|
||||
}
|
||||
assert response.additional_properties["model"] == "gpt-5.1-mini"
|
||||
|
||||
async def test_run_empty_response(
|
||||
self,
|
||||
mock_client: MagicMock,
|
||||
@@ -521,6 +584,50 @@ class TestGitHubCopilotAgentRunStreaming:
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].contents[0].text == "Hello"
|
||||
|
||||
async def test_run_streaming_captures_assistant_usage_event(
|
||||
self,
|
||||
mock_client: MagicMock,
|
||||
mock_session: MagicMock,
|
||||
assistant_delta_event: SessionEvent,
|
||||
session_idle_event: SessionEvent,
|
||||
) -> None:
|
||||
"""Test streaming final response includes assistant usage metadata."""
|
||||
usage_data = AssistantUsageData(
|
||||
model="gpt-5.1-mini",
|
||||
input_tokens=10,
|
||||
output_tokens=4,
|
||||
finish_reason="stop",
|
||||
)
|
||||
usage_event = SessionEvent(
|
||||
data=usage_data,
|
||||
id=uuid4(),
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
type=SessionEventType.ASSISTANT_USAGE,
|
||||
)
|
||||
events = [assistant_delta_event, usage_event, session_idle_event]
|
||||
|
||||
def mock_on(handler: Any) -> Any:
|
||||
for event in events:
|
||||
handler(event)
|
||||
return lambda: None
|
||||
|
||||
mock_session.on = mock_on
|
||||
|
||||
agent = GitHubCopilotAgent(client=mock_client)
|
||||
stream = agent.run("Hello", stream=True)
|
||||
async for _ in stream:
|
||||
pass
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert response.text == "Hello"
|
||||
assert response.finish_reason == "stop"
|
||||
assert response.usage_details == {
|
||||
"input_token_count": 10,
|
||||
"output_token_count": 4,
|
||||
"total_token_count": 14,
|
||||
}
|
||||
assert response.additional_properties["model"] == "gpt-5.1-mini"
|
||||
|
||||
async def test_run_streaming_with_session(
|
||||
self,
|
||||
mock_client: MagicMock,
|
||||
|
||||
@@ -540,11 +540,30 @@ class OllamaChatClient(
|
||||
|
||||
def _parse_streaming_response_from_ollama(self, response: OllamaChatResponse) -> ChatResponseUpdate:
|
||||
contents = self._parse_contents_from_ollama(response)
|
||||
finish_reason = None
|
||||
if response.done:
|
||||
usage_details = UsageDetails(
|
||||
**{
|
||||
key: value
|
||||
for key, value in {
|
||||
"input_token_count": response.prompt_eval_count,
|
||||
"output_token_count": response.eval_count,
|
||||
"total_token_count": response.prompt_eval_count + response.eval_count
|
||||
if isinstance(response.prompt_eval_count, int) and isinstance(response.eval_count, int)
|
||||
else None,
|
||||
}.items()
|
||||
if isinstance(value, int)
|
||||
}
|
||||
)
|
||||
if usage_details:
|
||||
contents.append(Content.from_usage(usage_details, raw_representation=response))
|
||||
finish_reason = response.done_reason if response.done_reason in ("stop", "length") else None
|
||||
return ChatResponseUpdate(
|
||||
contents=contents,
|
||||
role="assistant",
|
||||
model=response.model,
|
||||
created_at=response.created_at,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
|
||||
def _parse_response_from_ollama(
|
||||
@@ -554,15 +573,27 @@ class OllamaChatClient(
|
||||
response_format: Any | None = None,
|
||||
) -> ChatResponse:
|
||||
contents = self._parse_contents_from_ollama(response)
|
||||
usage_details = UsageDetails(
|
||||
**{
|
||||
key: value
|
||||
for key, value in {
|
||||
"input_token_count": response.prompt_eval_count,
|
||||
"output_token_count": response.eval_count,
|
||||
"total_token_count": response.prompt_eval_count + response.eval_count
|
||||
if isinstance(response.prompt_eval_count, int) and isinstance(response.eval_count, int)
|
||||
else None,
|
||||
}.items()
|
||||
if isinstance(value, int)
|
||||
}
|
||||
)
|
||||
finish_reason = response.done_reason if response.done_reason in ("stop", "length") else None
|
||||
|
||||
return ChatResponse(
|
||||
messages=[Message(role="assistant", contents=contents)],
|
||||
model=response.model,
|
||||
created_at=response.created_at,
|
||||
usage_details=UsageDetails(
|
||||
input_token_count=response.prompt_eval_count,
|
||||
output_token_count=response.eval_count,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
usage_details=usage_details or None,
|
||||
response_format=response_format,
|
||||
)
|
||||
|
||||
|
||||
@@ -262,6 +262,71 @@ async def test_cmc(
|
||||
assert result.text == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_maps_done_reason_to_finish_reason(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[Message],
|
||||
) -> None:
|
||||
mock_chat.return_value = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
eval_count=2,
|
||||
prompt_eval_count=3,
|
||||
done_reason="length",
|
||||
)
|
||||
chat_history.append(Message(contents=["hello world"], role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert result.finish_reason == "length"
|
||||
assert result.usage_details == {
|
||||
"input_token_count": 3,
|
||||
"output_token_count": 2,
|
||||
"total_token_count": 5,
|
||||
}
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_leaves_unknown_done_reason_unset(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[Message],
|
||||
) -> None:
|
||||
mock_chat.return_value = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
done_reason="load",
|
||||
)
|
||||
chat_history.append(Message(contents=["hello world"], role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert result.finish_reason is None
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_omits_usage_when_token_counts_are_missing(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[Message],
|
||||
) -> None:
|
||||
mock_chat.return_value = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
done_reason="stop",
|
||||
)
|
||||
chat_history.append(Message(contents=["hello world"], role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = await ollama_client.get_response(messages=chat_history)
|
||||
|
||||
assert result.finish_reason == "stop"
|
||||
assert not result.usage_details
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_response_format_dict(
|
||||
mock_chat: AsyncMock,
|
||||
@@ -380,6 +445,72 @@ async def test_cmc_streaming(
|
||||
assert chunk.text == "test"
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_maps_done_reason_and_usage(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[Message],
|
||||
) -> None:
|
||||
response = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
done=True,
|
||||
done_reason="stop",
|
||||
eval_count=4,
|
||||
prompt_eval_count=6,
|
||||
created_at="2024-01-01T00:00:00Z",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [response]
|
||||
mock_chat.return_value = stream
|
||||
chat_history.append(Message(contents=["hello world"], role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = ollama_client.get_response(messages=chat_history, stream=True)
|
||||
async for _ in result:
|
||||
pass
|
||||
final_response = await result.get_final_response()
|
||||
|
||||
assert final_response.text == "test"
|
||||
assert final_response.finish_reason == "stop"
|
||||
assert final_response.usage_details == {
|
||||
"input_token_count": 6,
|
||||
"output_token_count": 4,
|
||||
"total_token_count": 10,
|
||||
}
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_ignores_done_reason_and_usage_before_final_chunk(
|
||||
mock_chat: AsyncMock,
|
||||
ollama_unit_test_env: dict[str, str],
|
||||
chat_history: list[Message],
|
||||
) -> None:
|
||||
response = OllamaChatResponse(
|
||||
message=OllamaMessage(content="test", role="assistant"),
|
||||
model="test",
|
||||
done=False,
|
||||
done_reason="stop",
|
||||
eval_count=4,
|
||||
prompt_eval_count=6,
|
||||
created_at="2024-01-01T00:00:00Z",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [response]
|
||||
mock_chat.return_value = stream
|
||||
chat_history.append(Message(contents=["hello world"], role="user"))
|
||||
|
||||
ollama_client = OllamaChatClient()
|
||||
result = ollama_client.get_response(messages=chat_history, stream=True)
|
||||
async for _ in result:
|
||||
pass
|
||||
final_response = await result.get_final_response()
|
||||
|
||||
assert final_response.text == "test"
|
||||
assert final_response.finish_reason is None
|
||||
assert final_response.usage_details is None
|
||||
|
||||
|
||||
@patch.object(AsyncClient, "chat", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming_reasoning(
|
||||
mock_chat: AsyncMock,
|
||||
|
||||
Reference in New Issue
Block a user