diff --git a/tests/agents/core/test_llm_processor.py b/tests/agents/core/test_llm_processor.py index 27bc20ee2..03829b264 100644 --- a/tests/agents/core/test_llm_processor.py +++ b/tests/agents/core/test_llm_processor.py @@ -21,6 +21,7 @@ from trpc_agent_sdk.models import LLMModel, LlmRequest, LlmResponse, ModelRegistry from trpc_agent_sdk.sessions import InMemorySessionService from trpc_agent_sdk.types import Content, Part +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason class _StubAgent(BaseAgent): @@ -113,6 +114,18 @@ def test_preserves_error_fields(self, model, invocation_context): assert event.error_code == "validation_error" assert event.error_message == "Bad request" + def test_preserves_finish_reason(self, model, invocation_context): + proc = LlmProcessor(model) + response = LlmResponse( + content=None, + partial=False, + finish_reason=GenaiFinishReason.STOP, + finish_reason_raw="end_turn", + ) + event = proc._create_event_from_response(invocation_context, "evt-3", response) + assert event.finish_reason == GenaiFinishReason.STOP + assert event.finish_reason_raw == "end_turn" + # --------------------------------------------------------------------------- # _create_error_event diff --git a/tests/models/test_anthropic_model.py b/tests/models/test_anthropic_model.py index 58b5f7ec4..ba15e0637 100644 --- a/tests/models/test_anthropic_model.py +++ b/tests/models/test_anthropic_model.py @@ -16,6 +16,7 @@ from trpc_agent_sdk.models import AnthropicModel from trpc_agent_sdk.models import LlmRequest from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import FunctionDeclaration from trpc_agent_sdk.types import GenerateContentConfig from trpc_agent_sdk.types import Part @@ -437,6 +438,7 @@ async def mock_aiter(self): cache_creation_input_tokens=None, ) mock_final_message.content = [anthropic_types.TextBlock(text="Hello world!", type="text")] + mock_final_message.stop_reason = "end_turn" mock_stream.get_final_message = AsyncMock(return_value=mock_final_message) mock_client = AsyncMock() @@ -463,6 +465,8 @@ async def mock_aiter(self): assert final_response.usage_metadata is not None assert final_response.usage_metadata.prompt_token_count == 5 assert final_response.usage_metadata.candidates_token_count == 3 + assert final_response.finish_reason_raw == "end_turn" + assert final_response.finish_reason == GenaiFinishReason.STOP class TestAnthropicInjectCacheControl: diff --git a/tests/models/test_anthropic_model_ext.py b/tests/models/test_anthropic_model_ext.py index 2f02b0d9a..76283bd28 100644 --- a/tests/models/test_anthropic_model_ext.py +++ b/tests/models/test_anthropic_model_ext.py @@ -17,6 +17,7 @@ from trpc_agent_sdk.models._anthropic_model import _FinishReason from trpc_agent_sdk.types import ( Content, + FinishReason as GenaiFinishReason, FunctionDeclaration, GenerateContentConfig, Part, @@ -25,11 +26,11 @@ Type, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- + def _model(**kwargs): defaults = dict(model_name="claude-3-5-sonnet-20241022", api_key="test-key") defaults.update(kwargs) @@ -42,6 +43,7 @@ def _model(**kwargs): class TestToClaudeRole: + def test_model_role_maps_to_assistant(self): """'model' role is converted to 'assistant'.""" model = _model() @@ -74,6 +76,7 @@ def test_unknown_role_defaults_to_user(self): class TestIsImagePart: + def test_image_part_detected(self): """Image inline data is recognised as image.""" model = _model() @@ -99,6 +102,7 @@ def test_text_part_not_image(self): class TestPartToMessageBlock: + def test_text_part(self): """Text part converts to TextBlockParam.""" model = _model() @@ -122,7 +126,13 @@ def test_function_response_with_content_array(self): model = _model() part = Part.from_function_response( name="fn", - response={"content": [{"type": "text", "text": "ok"}, {"type": "other", "data": 1}]}, + response={"content": [{ + "type": "text", + "text": "ok" + }, { + "type": "other", + "data": 1 + }]}, ) part.function_response.id = "id1" block = model._part_to_message_block(part) @@ -174,6 +184,7 @@ def test_unsupported_part_raises(self): class TestParseFinishReason: + def test_end_turn_maps_to_stop(self): """'end_turn' maps to STOP.""" model = _model() @@ -205,12 +216,59 @@ def test_none_maps_to_error(self): assert model._parse_finish_reason(None) == _FinishReason.ERROR +class TestLlmResponseFinishReason: + """finish_reason / finish_reason_raw on Anthropic responses.""" + + @staticmethod + def _message(stop_reason): + return anthropic_types.Message( + id="msg_1", + type="message", + role="assistant", + content=[anthropic_types.TextBlock(text="hi", type="text")], + model="claude-3-5-sonnet-20241022", + stop_reason=stop_reason, + stop_sequence=None, + usage=anthropic_types.Usage(input_tokens=1, output_tokens=2), + ) + + def test_max_tokens_is_carried(self): + model = _model() + result = model._message_to_llm_response(self._message("max_tokens")) + assert result.finish_reason_raw == "max_tokens" + assert result.finish_reason == GenaiFinishReason.MAX_TOKENS + + def test_end_turn_raw_is_preserved(self): + model = _model() + result = model._message_to_llm_response(self._message("end_turn")) + assert result.finish_reason_raw == "end_turn" + assert result.finish_reason == GenaiFinishReason.STOP + + def test_tool_use_raw_is_preserved(self): + model = _model() + result = model._message_to_llm_response(self._message("tool_use")) + assert result.finish_reason_raw == "tool_use" + assert result.finish_reason == GenaiFinishReason.STOP + + def test_unknown_stop_reason_maps_to_unspecified(self): + model = _model() + message = MagicMock() + message.content = [anthropic_types.TextBlock(text="hi", type="text")] + message.usage = anthropic_types.Usage(input_tokens=1, output_tokens=2) + message.stop_reason = "something_else" + message.model_dump_json = MagicMock(return_value="{}") + result = model._message_to_llm_response(message) + assert result.finish_reason_raw == "something_else" + assert result.finish_reason == GenaiFinishReason.FINISH_REASON_UNSPECIFIED + + # --------------------------------------------------------------------------- # _update_type_string # --------------------------------------------------------------------------- class TestUpdateTypeString: + def test_lowercases_type_field(self): """Type field is lowered.""" model = _model() @@ -233,7 +291,9 @@ def test_recursive_properties_in_items(self): "items": { "type": "OBJECT", "properties": { - "name": {"type": "STRING"}, + "name": { + "type": "STRING" + }, }, }, } @@ -247,6 +307,7 @@ def test_recursive_properties_in_items(self): class TestFunctionDeclarationToToolParam: + def test_with_parameters_json_schema(self): """Uses parameters_json_schema when available.""" model = _model() @@ -290,6 +351,7 @@ def test_no_parameters(self): class TestContentBlockToPart: + def test_text_block(self): """TextBlock converts to text Part.""" model = _model() @@ -336,6 +398,7 @@ def test_unsupported_block_raises(self): class TestFormatMessagesEdgeCases: + def test_image_in_assistant_turn_skipped(self): """Image data in assistant turn is skipped with warning.""" model = _model() @@ -359,6 +422,7 @@ def test_empty_message_blocks_filtered(self): class TestCreateStreamingToolCallResponse: + def test_returns_none_when_no_tool_uses(self): """Returns None with empty accumulated_tool_uses.""" model = _model() @@ -368,7 +432,11 @@ def test_returns_none_when_tool_name_empty(self): """Returns None when current tool has no name.""" model = _model() result = model._create_streaming_tool_call_response( - [{"id": "id1", "name": "", "accumulated_input": ""}], + [{ + "id": "id1", + "name": "", + "accumulated_input": "" + }], '{"x":1}', ) assert result is None @@ -377,7 +445,11 @@ def test_returns_none_when_tool_not_in_streaming_set(self): """Returns None when tool name is not in streaming_tool_names.""" model = _model() result = model._create_streaming_tool_call_response( - [{"id": "id1", "name": "my_tool", "accumulated_input": "{}"}], + [{ + "id": "id1", + "name": "my_tool", + "accumulated_input": "{}" + }], '{"x":1}', streaming_tool_names={"other_tool"}, ) @@ -387,7 +459,11 @@ def test_returns_response_when_tool_in_streaming_set(self): """Returns LlmResponse when tool name is in streaming_tool_names.""" model = _model() result = model._create_streaming_tool_call_response( - [{"id": "id1", "name": "my_tool", "accumulated_input": '{"x":'}], + [{ + "id": "id1", + "name": "my_tool", + "accumulated_input": '{"x":' + }], '1}', streaming_tool_names={"my_tool"}, ) @@ -399,7 +475,11 @@ def test_returns_response_when_streaming_names_none(self): """Returns response for any tool when streaming_tool_names is None.""" model = _model() result = model._create_streaming_tool_call_response( - [{"id": "id1", "name": "any_tool", "accumulated_input": ""}], + [{ + "id": "id1", + "name": "any_tool", + "accumulated_input": "" + }], '{}', streaming_tool_names=None, ) @@ -412,6 +492,7 @@ def test_returns_response_when_streaming_names_none(self): class TestMergeConfigsExtended: + def test_no_request_no_default_returns_empty(self): """No request config and no default yields empty GenerateContentConfig.""" model = _model() @@ -435,6 +516,7 @@ def test_default_fields_applied_to_request(self): class TestLogUnsupportedConfigOptions: + def test_logs_warning_for_unsupported_options(self): """Logs warnings for every unsupported config option.""" model = _model() @@ -466,6 +548,7 @@ def test_no_warning_for_supported_options(self): class TestGenerateSingleError: + @pytest.mark.asyncio async def test_api_error_raises_and_closes_client(self): """API error during single generation is raised to the retry layer.""" @@ -485,6 +568,7 @@ async def test_api_error_raises_and_closes_client(self): class TestCreateAsyncClient: + def test_client_created_with_api_key(self): """Client is created with the model's api_key.""" model = _model(api_key="sk-test", base_url="https://example.com") @@ -511,6 +595,7 @@ def test_client_without_base_url(self): class TestValidateRequestExtended: + def test_part_with_inline_data_is_valid(self): """Part with inline_data passes validation.""" model = _model() @@ -518,5 +603,6 @@ def test_part_with_inline_data_is_valid(self): request = LlmRequest(contents=[Content(parts=[part], role="user")]) model.validate_request(request) + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/models/test_litellm_model.py b/tests/models/test_litellm_model.py index f201d53a6..4965e3808 100644 --- a/tests/models/test_litellm_model.py +++ b/tests/models/test_litellm_model.py @@ -22,6 +22,7 @@ from trpc_agent_sdk.models._litellm_model import _build_response_format_for_litellm from trpc_agent_sdk.models._litellm_model import _is_litellm_gemini_model from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import GenerateContentConfig from trpc_agent_sdk.types import Part @@ -136,6 +137,8 @@ def test_normal_text_and_usage_no_error_code_when_stop(self): assert r.content.parts[0].text == "Hi there" assert getattr(r.content.parts[0], "thought", None) is False assert r.usage_metadata.prompt_token_count == 10 + assert r.finish_reason == GenaiFinishReason.STOP + assert r.finish_reason_raw == "stop" def test_finish_reason_not_stop_sets_error_code(self): model = LiteLLMModel(model_name="openai/gpt-4") @@ -157,6 +160,8 @@ def test_finish_reason_not_stop_sets_error_code(self): assert r.error_code == "length" assert r.content is not None assert r.content.parts[0].text == "x" + assert r.finish_reason == GenaiFinishReason.MAX_TOKENS + assert r.finish_reason_raw == "length" def test_tool_calls_from_message(self): model = LiteLLMModel(model_name="openai/gpt-4") @@ -490,7 +495,9 @@ async def test_generate_async_stream_yields_partial_then_final(self): } chunk3 = Mock() chunk3.model_dump.return_value = { - CHOICES: [], + CHOICES: [{ + FINISH_REASON: "length" + }], USAGE: { "prompt_tokens": 1, "completion_tokens": 2, @@ -526,6 +533,8 @@ def fake_import(name, *a, **k): assert "Hello " in finals[0].content.parts[0].text or "world" in finals[0].content.parts[0].text assert finals[0].usage_metadata is not None assert finals[0].usage_metadata.total_token_count == 3 + assert finals[0].finish_reason == GenaiFinishReason.MAX_TOKENS + assert finals[0].finish_reason_raw == "length" @pytest.mark.asyncio async def test_generate_async_stream_includes_stream_options_in_api_params(self): @@ -947,7 +956,8 @@ def _model(self): def test_x_should_retry_header_has_priority(self): model = self._model() assert model._get_model_retry_info(_LiteLLMRetryTestError(400, {"x-should-retry": "true"})).should_retry is True - assert model._get_model_retry_info(_LiteLLMRetryTestError(500, {"x-should-retry": "false"})).should_retry is False + assert model._get_model_retry_info(_LiteLLMRetryTestError(500, + {"x-should-retry": "false"})).should_retry is False @pytest.mark.parametrize("status_code", [408, 409, 429, 500, 503]) def test_retryable_status_codes(self, status_code): diff --git a/tests/models/test_openai_model_ext.py b/tests/models/test_openai_model_ext.py index 4b650a039..748c956a6 100644 --- a/tests/models/test_openai_model_ext.py +++ b/tests/models/test_openai_model_ext.py @@ -21,6 +21,7 @@ ) from trpc_agent_sdk.types import ( Content, + FinishReason as GenaiFinishReason, FunctionCall, FunctionDeclaration, FunctionResponse, @@ -524,6 +525,8 @@ def test_without_content_stop(self): result = model._create_response_without_content(resp) assert result.content is None assert result.error_code is None + assert result.finish_reason == GenaiFinishReason.STOP + assert result.finish_reason_raw == "stop" assert result.response_id == "resp_1" def test_without_content_length(self): @@ -532,6 +535,8 @@ def test_without_content_length(self): resp = {"id": "r2", "choices": [{"finish_reason": "length"}]} result = model._create_response_without_content(resp) assert result.error_code == "length" + assert result.finish_reason == GenaiFinishReason.MAX_TOKENS + assert result.finish_reason_raw == "length" def test_with_content_text_only(self): """Text content response has text part.""" @@ -576,6 +581,53 @@ def test_with_content_tool_calls(self): fc_parts = [p for p in result.content.parts if p.function_call] assert len(fc_parts) == 1 assert fc_parts[0].function_call.name == "search" + assert result.finish_reason == GenaiFinishReason.STOP + assert result.finish_reason_raw == "tool_calls" + + def test_with_content_content_filter(self): + """content_filter maps to SAFETY while the raw value is preserved.""" + model = _model() + resp = { + "choices": [{ + "message": { + "content": "partial", + "role": "assistant" + }, + "finish_reason": "content_filter" + }], + } + result = model._create_response_with_content(resp) + assert result.error_code == "content_filter" + assert result.finish_reason == GenaiFinishReason.SAFETY + assert result.finish_reason_raw == "content_filter" + + def test_responses_finish_reason_completed(self): + """Responses API 'completed' maps to STOP.""" + raw, normalized = OpenAIModel._responses_finish_reason({"status": "completed"}) + assert raw == "completed" + assert normalized == GenaiFinishReason.STOP + + def test_responses_finish_reason_incomplete_max_tokens(self): + """Responses API 'incomplete' + max_output_tokens maps to MAX_TOKENS.""" + raw, normalized = OpenAIModel._responses_finish_reason({ + "status": "incomplete", + "incomplete_details": { + "reason": "max_output_tokens" + }, + }) + assert raw == "max_output_tokens" + assert normalized == GenaiFinishReason.MAX_TOKENS + + def test_responses_response_carries_finish_reason(self): + """_create_responses_response writes both fields.""" + model = _model() + result = model._create_responses_response({ + "id": "r1", + "status": "completed", + "output": [], + }) + assert result.finish_reason == GenaiFinishReason.STOP + assert result.finish_reason_raw == "completed" def test_with_content_empty_fallback(self): """Response with no text and no tool calls produces empty text part.""" @@ -1559,7 +1611,8 @@ def extract_metadata(response_data): }, "finish_reason": "stop", }], - "usage": None, + "usage": + None, } with patch.object(model, "_create_async_client") as mock_factory: @@ -1647,6 +1700,55 @@ async def mock_stream(): thought_partials = [r for r in partial_responses if r.content and r.content.parts[0].thought] assert len(thought_partials) >= 1 + @pytest.mark.asyncio + async def test_streaming_reports_finish_reason(self): + """The terminal chunk's finish_reason reaches the final response.""" + model = _model() + request = _request([Content(parts=[Part.from_text(text="hi")], role="user")]) + + content_chunk = Mock() + content_chunk.model_dump.return_value = { + "id": "resp_1", + "choices": [{ + "delta": { + "content": "hello" + }, + "finish_reason": None + }], + "usage": None, + } + finish_chunk = Mock() + finish_chunk.model_dump.return_value = { + "id": "resp_1", + "choices": [{ + "delta": {}, + "finish_reason": "length" + }], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3 + }, + } + + async def mock_stream(): + yield content_chunk + yield finish_chunk + + with patch.object(model, "_create_async_client") as mock_factory: + mock_client = AsyncMock() + mock_client.chat.completions.create = AsyncMock(return_value=mock_stream()) + mock_client.close = AsyncMock() + mock_factory.return_value = mock_client + + responses = [] + async for r in model.generate_async(request, stream=True): + responses.append(r) + + final = [r for r in responses if not r.partial][-1] + assert final.finish_reason == GenaiFinishReason.MAX_TOKENS + assert final.finish_reason_raw == "length" + @pytest.mark.asyncio async def test_streaming_extracts_provider_metadata_from_first_chunk(self): """Provider metadata is extracted once from the first stream chunk.""" @@ -1801,8 +1903,10 @@ async def stream_events(): yield { "type": "response.completed", "response": { - "id": "resp_meta", - "status": "completed", + "id": + "resp_meta", + "status": + "completed", "output": [{ "type": "message", "role": "assistant", diff --git a/tests/telemetry/test_custom_trace.py b/tests/telemetry/test_custom_trace.py index f5e0c961b..149cb6a15 100644 --- a/tests/telemetry/test_custom_trace.py +++ b/tests/telemetry/test_custom_trace.py @@ -29,6 +29,7 @@ CustomTraceReporter, _SyntheticTool, ) +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason # --------------------------------------------------------------------------- # Helpers @@ -78,6 +79,8 @@ def _make_event( error_code=None, error_message=None, custom_metadata=None, + finish_reason=None, + finish_reason_raw=None, ): event = MagicMock() event.partial = partial @@ -86,6 +89,8 @@ def _make_event( event.error_code = error_code event.error_message = error_message event.custom_metadata = custom_metadata + event.finish_reason = finish_reason + event.finish_reason_raw = finish_reason_raw event.is_error = MagicMock(return_value=error_code is not None) event.get_function_calls = MagicMock(return_value=function_calls or []) event.get_function_responses = MagicMock(return_value=function_responses or []) @@ -215,6 +220,8 @@ def test_with_event(self, MockLlmResponse): error_code="REMOTE_ERROR", error_message="err", custom_metadata={"error_type": "RemoteError"}, + finish_reason=None, + finish_reason_raw=None, ) @patch("trpc_agent_sdk.telemetry._custom_trace.LlmResponse") @@ -229,6 +236,28 @@ def test_with_none_event(self, MockLlmResponse): error_code=None, error_message=None, custom_metadata=None, + finish_reason=None, + finish_reason_raw=None, + ) + + @patch("trpc_agent_sdk.telemetry._custom_trace.LlmResponse") + def test_forwards_finish_reason(self, MockLlmResponse): + reporter = CustomTraceReporter(agent_name="a") + event = _make_event( + content="content_obj", + finish_reason=GenaiFinishReason.STOP, + finish_reason_raw="end_turn", + ) + + reporter._create_synthetic_llm_response(event) + + MockLlmResponse.assert_called_once_with( + content="content_obj", + error_code=None, + error_message=None, + custom_metadata=None, + finish_reason=GenaiFinishReason.STOP, + finish_reason_raw="end_turn", ) diff --git a/trpc_agent_sdk/agents/core/_llm_processor.py b/trpc_agent_sdk/agents/core/_llm_processor.py index d6e877fd5..cf58c63d9 100644 --- a/trpc_agent_sdk/agents/core/_llm_processor.py +++ b/trpc_agent_sdk/agents/core/_llm_processor.py @@ -271,6 +271,8 @@ def _create_event_from_response(self, ctx: InvocationContext, event_id: str, res custom_metadata=response.custom_metadata, usage_metadata=response.usage_metadata, response_id=response.response_id, + finish_reason=response.finish_reason, + finish_reason_raw=response.finish_reason_raw, # Event-specific fields extracted from context id=event_id, invocation_id=ctx.invocation_id, diff --git a/trpc_agent_sdk/models/_anthropic_model.py b/trpc_agent_sdk/models/_anthropic_model.py index ad4b9a1fa..71dcd7c0b 100644 --- a/trpc_agent_sdk/models/_anthropic_model.py +++ b/trpc_agent_sdk/models/_anthropic_model.py @@ -29,6 +29,7 @@ from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.log import logger from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import GenerateContentConfig from trpc_agent_sdk.types import GenerateContentResponseUsageMetadata from trpc_agent_sdk.types import Part @@ -92,6 +93,26 @@ class _FinishReason(str, Enum): TOOL_USE = "tool_use" +# Mapping of Anthropic stop_reason strings to google-genai FinishReason values. +# Keep this aligned with ADK's `_STOP_REASON_MAPPING` in +# reference/adk-python/src/google/adk/models/anthropic_llm.py. +_ANTHROPIC_STOP_REASON_TO_GENAI = { + "end_turn": GenaiFinishReason.STOP, + "stop_sequence": GenaiFinishReason.STOP, + "tool_use": GenaiFinishReason.STOP, + "pause_turn": GenaiFinishReason.STOP, + "max_tokens": GenaiFinishReason.MAX_TOKENS, + "refusal": GenaiFinishReason.SAFETY, +} + + +def _to_genai_finish_reason(stop_reason: Optional[str]) -> Optional[GenaiFinishReason]: + """Map an Anthropic stop_reason to google-genai's enum.""" + if not stop_reason: + return None + return _ANTHROPIC_STOP_REASON_TO_GENAI.get(stop_reason, GenaiFinishReason.FINISH_REASON_UNSPECIFIED) + + class _ApiParamsKey(str, Enum): """API parameter keys for Anthropic API calls.""" @@ -488,6 +509,8 @@ def _message_to_llm_response(self, message: anthropic_types.Message) -> LlmRespo parts=[self._content_block_to_part(cb) for cb in message.content], ), usage_metadata=self._build_usage_metadata(message.usage), + finish_reason=_to_genai_finish_reason(getattr(message, "stop_reason", None)), + finish_reason_raw=getattr(message, "stop_reason", None), ) def _merge_configs(self, request_config: Optional[GenerateContentConfig]) -> GenerateContentConfig: @@ -690,6 +713,8 @@ async def _generate_stream( yield LlmResponse(content=final_content, usage_metadata=final_usage, partial=False, + finish_reason=_to_genai_finish_reason(getattr(final_message, "stop_reason", None)), + finish_reason_raw=getattr(final_message, "stop_reason", None), custom_metadata={"stream_complete": True}) except Exception: diff --git a/trpc_agent_sdk/models/_litellm_model.py b/trpc_agent_sdk/models/_litellm_model.py index fa3bc404d..fac5db555 100644 --- a/trpc_agent_sdk/models/_litellm_model.py +++ b/trpc_agent_sdk/models/_litellm_model.py @@ -20,6 +20,7 @@ from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.log import logger from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import GenerateContentConfig from trpc_agent_sdk.types import GenerateContentResponseUsageMetadata from trpc_agent_sdk.types import Part @@ -155,6 +156,26 @@ class _LiteLLMApiParamsKey(str, Enum): CACHE_CONTROL_INJECTION_POINTS = "cache_control_injection_points" +# Mapping of LiteLLM finish_reason strings to google-genai FinishReason values. +# Keep this aligned with ADK's `_FINISH_REASON_MAPPING` in +# reference/adk-python/src/google/adk/models/lite_llm.py, where unknown values +# fall back to OTHER (unlike the OpenAI adapter, which uses UNSPECIFIED). +_LITELLM_FINISH_REASON_TO_GENAI = { + "stop": GenaiFinishReason.STOP, + "tool_calls": GenaiFinishReason.STOP, + "function_call": GenaiFinishReason.STOP, + "length": GenaiFinishReason.MAX_TOKENS, + "content_filter": GenaiFinishReason.SAFETY, +} + + +def _to_genai_finish_reason(raw: Optional[str]) -> Optional[GenaiFinishReason]: + """Map a LiteLLM finish reason to google-genai's enum.""" + if not raw: + return None + return _LITELLM_FINISH_REASON_TO_GENAI.get(str(raw).lower(), GenaiFinishReason.OTHER) + + @register_model(model_name="LiteLLMModel", supported_models=_LITELLM_SUPPORTED_MODELS) class LiteLLMModel(OpenAIModel): """model_name must be provider/model (e.g. openai/gpt-4o). kwargs: api_key, base_url→api_base.""" @@ -393,7 +414,14 @@ def _create_response_with_content(self, response_dict: Dict[str, Any], partial: finish_reason = first_choice.get(const.FINISH_REASON) if finish_reason and finish_reason != FinishReason.STOP.value: error_code = finish_reason - return LlmResponse(content=content, usage_metadata=usage_meta, partial=partial, error_code=error_code) + return LlmResponse( + content=content, + usage_metadata=usage_meta, + partial=partial, + error_code=error_code, + finish_reason=_to_genai_finish_reason(finish_reason), + finish_reason_raw=finish_reason, + ) @override def _log_unsupported_config_options(self, config: GenerateContentConfig) -> None: @@ -524,6 +552,7 @@ async def _generate_stream( accumulated_content = "" accumulated_tool_calls: List[Dict[str, Any]] = [] usage_meta: Optional[GenerateContentResponseUsageMetadata] = None + final_finish_reason_raw: Optional[str] = None async for chunk in stream_handle: if chunk is None: @@ -532,6 +561,10 @@ async def _generate_stream( choices = chunk_dict.get(const.CHOICES) or [] if choices: choice = choices[0] or {} + chunk_finish_reason = (choice.get(const.FINISH_REASON) if isinstance(choice, dict) else getattr( + choice, "finish_reason", None)) + if chunk_finish_reason: + final_finish_reason_raw = chunk_finish_reason delta = (choice.get(const.DELTA) if isinstance(choice, dict) else getattr(choice, "delta", None)) or {} content_delta = (delta.get(const.CONTENT) if isinstance(delta, dict) else getattr( @@ -573,6 +606,8 @@ async def _generate_stream( content=final_content, usage_metadata=usage_meta, partial=False, + finish_reason=_to_genai_finish_reason(final_finish_reason_raw), + finish_reason_raw=final_finish_reason_raw, custom_metadata={"stream_complete": True}, ) except Exception: diff --git a/trpc_agent_sdk/models/_llm_response.py b/trpc_agent_sdk/models/_llm_response.py index aa2fbc932..9c0615057 100644 --- a/trpc_agent_sdk/models/_llm_response.py +++ b/trpc_agent_sdk/models/_llm_response.py @@ -22,9 +22,12 @@ # limitations under the License. """LLM response class for TRPC Agent framework.""" +from typing import Optional + from typing_extensions import override from trpc_agent_sdk.abc import ResponseABC +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import GenerateContentResponse @@ -47,8 +50,16 @@ class LlmResponse(ResponseABC): interrupted: Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming. custom_metadata: The custom metadata of the LlmResponse. + finish_reason: The normalized finish reason, aligned with google-genai. + finish_reason_raw: The raw finish reason returned by the provider. """ + finish_reason: Optional[GenaiFinishReason] = None + """The normalized finish reason, aligned with google-genai.""" + + finish_reason_raw: Optional[str] = None + """The raw finish reason returned by the provider.""" + def has_content(self) -> bool: """Returns whether the response carries user-visible content (text or function call). diff --git a/trpc_agent_sdk/models/_openai_model.py b/trpc_agent_sdk/models/_openai_model.py index cbeaf6923..03102b185 100644 --- a/trpc_agent_sdk/models/_openai_model.py +++ b/trpc_agent_sdk/models/_openai_model.py @@ -35,6 +35,7 @@ from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.log import logger from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FinishReason as GenaiFinishReason from trpc_agent_sdk.types import FunctionResponse from trpc_agent_sdk.types import GenerateContentConfig from trpc_agent_sdk.types import GenerateContentResponseUsageMetadata @@ -288,6 +289,25 @@ class FinishReason(str, Enum): TOOL_CALLS = const.TOOL_CALLS +# Mapping of OpenAI finish_reason strings to google-genai FinishReason values. +# Keep this aligned with ADK's `map_finish_reason` in +# reference/adk-python/src/google/adk/labs/openai/_openai_common.py. +_OPENAI_FINISH_REASON_TO_GENAI = { + "stop": GenaiFinishReason.STOP, + "tool_calls": GenaiFinishReason.STOP, + "function_call": GenaiFinishReason.STOP, + "length": GenaiFinishReason.MAX_TOKENS, + "content_filter": GenaiFinishReason.SAFETY, +} + + +def _to_genai_finish_reason(raw: Optional[str]) -> Optional[GenaiFinishReason]: + """Map an OpenAI-style finish reason to google-genai's enum.""" + if not raw: + return None + return _OPENAI_FINISH_REASON_TO_GENAI.get(str(raw).lower(), GenaiFinishReason.FINISH_REASON_UNSPECIFIED) + + class ToolKey(str, Enum): """Tool keys for tool calls.""" @@ -1427,16 +1447,22 @@ def _create_response_without_content(self, response_dict: dict) -> LlmResponse: # Get finish reason choices: list[dict] = response_dict.get(const.CHOICES, [{}]) + finish_reason_raw = choices[0].get(const.FINISH_REASON) if choices else None error_code = None - if choices and choices[0].get(const.FINISH_REASON): - finish_reason_value = choices[0][const.FINISH_REASON] - if finish_reason_value != FinishReason.STOP.value: - error_code = finish_reason_value + if finish_reason_raw and finish_reason_raw != FinishReason.STOP.value: + error_code = finish_reason_raw # Get response ID response_id = response_dict.get("id") - return LlmResponse(content=None, usage_metadata=usage, error_code=error_code, response_id=response_id) + return LlmResponse( + content=None, + usage_metadata=usage, + error_code=error_code, + finish_reason=_to_genai_finish_reason(finish_reason_raw), + finish_reason_raw=finish_reason_raw, + response_id=response_id, + ) def _create_response_with_content(self, response_dict: dict) -> LlmResponse: """Create a LlmResponse with text content.""" @@ -1506,14 +1532,22 @@ def _create_response_with_content(self, response_dict: dict) -> LlmResponse: usage = self._process_usage_from_response(response_dict) # Get finish reason for error handling + finish_reason_raw = choice.get(const.FINISH_REASON) error_code = None - if choice.get(const.FINISH_REASON) and choice[const.FINISH_REASON] != FinishReason.STOP.value: - error_code = choice[const.FINISH_REASON] + if finish_reason_raw and finish_reason_raw != FinishReason.STOP.value: + error_code = finish_reason_raw # Get response ID response_id = response_dict.get("id") - return LlmResponse(content=content, usage_metadata=usage, error_code=error_code, response_id=response_id) + return LlmResponse( + content=content, + usage_metadata=usage, + error_code=error_code, + finish_reason=_to_genai_finish_reason(finish_reason_raw), + finish_reason_raw=finish_reason_raw, + response_id=response_id, + ) @staticmethod def _model_dump(value: Any) -> dict: @@ -1743,6 +1777,26 @@ def _responses_error(response_dict: dict) -> tuple[Optional[str], Optional[str]] return str(error.get("code") or status), str(error.get("message") or error.get("reason") or status) return str(status), str(error) + @staticmethod + def _responses_finish_reason(response_dict: dict) -> tuple[Optional[str], Optional[GenaiFinishReason]]: + """Derive the raw and normalized finish reason from a Responses payload.""" + status = response_dict.get("status") + details = response_dict.get("incomplete_details") or {} + reason = details.get("reason") if isinstance(details, dict) else None + + if status == "completed": + return "completed", GenaiFinishReason.STOP + if status == "incomplete": + raw = str(reason) if reason else "incomplete" + if reason == "max_output_tokens": + return raw, GenaiFinishReason.MAX_TOKENS + if reason == "content_filter": + return raw, GenaiFinishReason.SAFETY + return raw, GenaiFinishReason.OTHER + if status in {"failed", "cancelled"}: + return str(status), GenaiFinishReason.OTHER + return None, None + def _create_responses_response(self, response_dict: dict) -> LlmResponse: """Convert a completed Responses payload into ``LlmResponse``.""" parts = [] @@ -1782,11 +1836,14 @@ def _create_responses_response(self, response_dict: dict) -> LlmResponse: content = Content(parts=parts, role=const.MODEL) if parts else None usage = self._process_responses_usage(response_dict) error_code, error_message = self._responses_error(response_dict) + finish_reason_raw, finish_reason = self._responses_finish_reason(response_dict) return LlmResponse( content=content, usage_metadata=usage, error_code=error_code, error_message=error_message, + finish_reason=finish_reason, + finish_reason_raw=finish_reason_raw, response_id=response_dict.get("id"), ) @@ -2465,6 +2522,7 @@ async def _generate_stream( accumulated_tool_calls: list[dict] = [] is_thinking = False # Track whether we're currently in thinking mode response_id: str | None = None # Track response ID from API + final_finish_reason_raw: str | None = None # Terminal chunk's finish_reason # For streaming tool call arguments - get the set of tool names that should stream streaming_tool_names = getattr(request, "streaming_tool_names", None) or set() @@ -2521,6 +2579,14 @@ async def _generate_stream( # Verify if the chunk contains valid text content if not self._verify_text_content_in_delta_response(chunk_dict): + # Read the raw value directly so unknown provider reasons are + # not collapsed to FinishReason.ERROR by the parser below. + chunk_choices = chunk_dict.get(const.CHOICES) or [] + if chunk_choices and isinstance(chunk_choices[0], dict): + chunk_finish_reason = chunk_choices[0].get(const.FINISH_REASON) + if chunk_finish_reason: + final_finish_reason_raw = chunk_finish_reason + # Process chunk without content (this is where tool calls are streamed) _, usage, delta_arguments = self._process_chunk_without_content(chunk_dict, accumulated_tool_calls) @@ -2545,6 +2611,9 @@ async def _generate_stream( if not choices: continue # Skip if no choices available choice: dict[str, dict] = choices[0] + chunk_finish_reason = choice.get(const.FINISH_REASON) + if chunk_finish_reason: + final_finish_reason_raw = chunk_finish_reason delta = choice[const.DELTA] # Handle reasoning content (thinking content) first @@ -2703,6 +2772,8 @@ async def _generate_stream( content=final_content, usage_metadata=final_usage, partial=False, + finish_reason=_to_genai_finish_reason(final_finish_reason_raw), + finish_reason_raw=final_finish_reason_raw, response_id=response_id, custom_metadata={"stream_complete": True}, ) diff --git a/trpc_agent_sdk/telemetry/_custom_trace.py b/trpc_agent_sdk/telemetry/_custom_trace.py index 04653c3ee..c270b160d 100644 --- a/trpc_agent_sdk/telemetry/_custom_trace.py +++ b/trpc_agent_sdk/telemetry/_custom_trace.py @@ -144,11 +144,17 @@ def _create_synthetic_llm_response(self, event: Event) -> LlmResponse: Returns: A synthetic LlmResponse for tracing purposes. """ + finish_reason = event.finish_reason if event else None + finish_reason_raw = event.finish_reason_raw if event else None return LlmResponse( content=event.content if event else None, error_code=event.error_code if event else None, error_message=event.error_message if event else None, custom_metadata=event.custom_metadata if event else None, + # Guard against MagicMock events used in tests: only pass real + # string-like values through to the pydantic model. + finish_reason=finish_reason if isinstance(finish_reason, str) else None, + finish_reason_raw=finish_reason_raw if isinstance(finish_reason_raw, str) else None, ) def _trace_function_call(self, event: Event) -> None: