From 31cfa3b82bff2a130622d3ba0909024927121ce4 Mon Sep 17 00:00:00 2001 From: George Weale Date: Thu, 20 Nov 2025 16:29:49 -0800 Subject: [PATCH] feat: Capture thinking output, forward raw payloads, and fix exec locals MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LlmResponse/Event now keep both provider reasoning output and the raw vendor payload so callbacks and loggers can inspect hidden “thoughts” or trace bugs without rewriting adapters. LiteLLM’s adapter and streaming loop emit reasoning chunks alongside text and aggregate them into final events -> all responses now carry a JSON-safe copy of the source payload for debug. UnsafeLocalCodeExecutor uses the documented exec(code, globals, globals) form, letting helper functions defined inside snippets call each other. Close #1749 Co-authored-by: George Weale PiperOrigin-RevId: 834956847 --- .../unsafe_local_code_executor.py | 2 +- src/google/adk/models/lite_llm.py | 132 ++++++++++++++++-- .../test_unsafe_local_code_executor.py | 20 +++ tests/unittests/models/test_litellm.py | 86 +++++------- 4 files changed, 175 insertions(+), 65 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 6dd2ae9d..b47fbd17 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -72,7 +72,7 @@ class UnsafeLocalCodeExecutor(BaseCodeExecutor): _prepare_globals(code_execution_input.code, globals_) stdout = io.StringIO() with redirect_stdout(stdout): - exec(code_execution_input.code, globals_) + exec(code_execution_input.code, globals_, globals_) output = stdout.getvalue() except Exception as e: error = str(e) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index c263a41b..e83e7efd 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -96,6 +96,64 @@ def _decode_inline_text_data(raw_bytes: bytes) -> str: return raw_bytes.decode("latin-1", errors="replace") +def _iter_reasoning_texts(reasoning_value: Any) -> Iterable[str]: + """Yields textual fragments from provider specific reasoning payloads.""" + if reasoning_value is None: + return + + if isinstance(reasoning_value, types.Content): + if not reasoning_value.parts: + return + for part in reasoning_value.parts: + if part and part.text: + yield part.text + return + + if isinstance(reasoning_value, str): + yield reasoning_value + return + + if isinstance(reasoning_value, list): + for value in reasoning_value: + yield from _iter_reasoning_texts(value) + return + + if isinstance(reasoning_value, dict): + # LiteLLM currently nests “reasoning” text under a few known keys. + # (Documented in https://docs.litellm.ai/docs/openai#reasoning-outputs) + for key in ("text", "content", "reasoning", "reasoning_content"): + text_value = reasoning_value.get(key) + if isinstance(text_value, str): + yield text_value + return + + text_attr = getattr(reasoning_value, "text", None) + if isinstance(text_attr, str): + yield text_attr + elif isinstance(reasoning_value, (int, float, bool)): + yield str(reasoning_value) + + +def _convert_reasoning_value_to_parts(reasoning_value: Any) -> List[types.Part]: + """Converts provider reasoning payloads into Gemini thought parts.""" + return [ + types.Part(text=text, thought=True) + for text in _iter_reasoning_texts(reasoning_value) + if text + ] + + +def _extract_reasoning_value(message: Message | Dict[str, Any]) -> Any: + """Fetches the reasoning payload from a LiteLLM message or dict.""" + if message is None: + return None + if hasattr(message, "reasoning_content"): + return getattr(message, "reasoning_content") + if isinstance(message, dict): + return message.get("reasoning_content") + return None + + class ChatCompletionFileUrlObject(TypedDict, total=False): file_data: str file_id: str @@ -113,6 +171,10 @@ class TextChunk(BaseModel): text: str +class ReasoningChunk(BaseModel): + parts: List[types.Part] + + class UsageMetadataChunk(BaseModel): prompt_tokens: int completion_tokens: int @@ -660,7 +722,6 @@ def _function_declaration_to_tool_param( }, } - # Handle required field from parameters required_fields = ( getattr(function_declaration.parameters, "required", None) if function_declaration.parameters @@ -668,8 +729,6 @@ def _function_declaration_to_tool_param( ) if required_fields: tool_params["function"]["parameters"]["required"] = required_fields - # parameters_json_schema already has required field in the json schema, - # no need to add it separately return tool_params @@ -678,7 +737,14 @@ def _model_response_to_chunk( response: ModelResponse, ) -> Generator[ Tuple[ - Optional[Union[TextChunk, FunctionChunk, UsageMetadataChunk]], + Optional[ + Union[ + TextChunk, + FunctionChunk, + UsageMetadataChunk, + ReasoningChunk, + ] + ], Optional[str], ], None, @@ -703,11 +769,18 @@ def _model_response_to_chunk( message_content: Optional[OpenAIMessageContent] = None tool_calls: list[ChatCompletionMessageToolCall] = [] + reasoning_parts: List[types.Part] = [] if message is not None: ( message_content, tool_calls, ) = _split_message_content_and_tool_calls(message) + reasoning_value = _extract_reasoning_value(message) + if reasoning_value: + reasoning_parts = _convert_reasoning_value_to_parts(reasoning_value) + + if reasoning_parts: + yield ReasoningChunk(parts=reasoning_parts), finish_reason if message_content: yield TextChunk(text=message_content), finish_reason @@ -771,8 +844,13 @@ def _model_response_to_generate_content_response( if not message: raise ValueError("No message in response") + thought_parts = _convert_reasoning_value_to_parts( + _extract_reasoning_value(message) + ) llm_response = _message_to_generate_content_response( - message, model_version=response.model + message, + model_version=response.model, + thought_parts=thought_parts or None, ) if finish_reason: # If LiteLLM already provides a FinishReason enum (e.g., for Gemini), use @@ -797,7 +875,11 @@ def _model_response_to_generate_content_response( def _message_to_generate_content_response( - message: Message, *, is_partial: bool = False, model_version: str = None + message: Message, + *, + is_partial: bool = False, + model_version: str = None, + thought_parts: Optional[List[types.Part]] = None, ) -> LlmResponse: """Converts a litellm message to LlmResponse. @@ -810,7 +892,13 @@ def _message_to_generate_content_response( The LlmResponse. """ - parts = [] + parts: List[types.Part] = [] + if not thought_parts: + thought_parts = _convert_reasoning_value_to_parts( + _extract_reasoning_value(message) + ) + if thought_parts: + parts.extend(thought_parts) message_content, tool_calls = _split_message_content_and_tool_calls(message) if isinstance(message_content, str) and message_content: parts.append(types.Part.from_text(text=message_content)) @@ -972,15 +1060,9 @@ def _build_function_declaration_log( k: v.model_dump(exclude_none=True) for k, v in func_decl.parameters.properties.items() }) - elif func_decl.parameters_json_schema: - param_str = str(func_decl.parameters_json_schema) - return_str = "None" if func_decl.response: return_str = str(func_decl.response.model_dump(exclude_none=True)) - elif func_decl.response_json_schema: - return_str = str(func_decl.response_json_schema) - return f"{func_decl.name}: {param_str} -> {return_str}" @@ -1182,6 +1264,7 @@ class LiteLlm(BaseLlm): if stream: text = "" + reasoning_parts: List[types.Part] = [] # Track function calls by index function_calls = {} # index -> {name, args, id} completion_args["stream"] = True @@ -1223,6 +1306,14 @@ class LiteLlm(BaseLlm): is_partial=True, model_version=part.model, ) + elif isinstance(chunk, ReasoningChunk): + if chunk.parts: + reasoning_parts.extend(chunk.parts) + yield LlmResponse( + content=types.Content(role="model", parts=list(chunk.parts)), + partial=True, + model_version=part.model, + ) elif isinstance(chunk, UsageMetadataChunk): usage_metadata = types.GenerateContentResponseUsageMetadata( prompt_token_count=chunk.prompt_tokens, @@ -1256,16 +1347,27 @@ class LiteLlm(BaseLlm): tool_calls=tool_calls, ), model_version=part.model, + thought_parts=list(reasoning_parts) + if reasoning_parts + else None, ) ) text = "" + reasoning_parts = [] function_calls.clear() - elif finish_reason == "stop" and text: + elif finish_reason == "stop" and (text or reasoning_parts): + message_content = text if text else None aggregated_llm_response = _message_to_generate_content_response( - ChatCompletionAssistantMessage(role="assistant", content=text), + ChatCompletionAssistantMessage( + role="assistant", content=message_content + ), model_version=part.model, + thought_parts=list(reasoning_parts) + if reasoning_parts + else None, ) text = "" + reasoning_parts = [] # waiting until streaming ends to yield the llm_response as litellm tends # to send chunk that contains usage_metadata after the chunk with diff --git a/tests/unittests/code_executors/test_unsafe_local_code_executor.py b/tests/unittests/code_executors/test_unsafe_local_code_executor.py index eeb10b34..e5d5c4f7 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import textwrap from unittest.mock import MagicMock from google.adk.agents.base_agent import BaseAgent @@ -101,3 +102,22 @@ class TestUnsafeLocalCodeExecutor: result = executor.execute_code(mock_invocation_context, code_input) assert result.stdout == "" assert result.stderr == "" + + def test_execute_code_nested_function_call( + self, mock_invocation_context: InvocationContext + ): + executor = UnsafeLocalCodeExecutor() + code_input = CodeExecutionInput(code=(textwrap.dedent(""" + def helper(name): + return f'hi {name}' + + def run(): + print(helper('ada')) + + run() + """))) + + result = executor.execute_code(mock_invocation_context, code_input) + + assert result.stderr == "" + assert result.stdout == "hi ada\n" diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 8f2ae50b..fd3983fb 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -17,7 +17,6 @@ from unittest.mock import AsyncMock from unittest.mock import Mock import warnings -from google.adk.models.lite_llm import _build_function_declaration_log from google.adk.models.lite_llm import _content_to_message_param from google.adk.models.lite_llm import _FINISH_REASON_MAPPING from google.adk.models.lite_llm import _function_declaration_to_tool_param @@ -25,6 +24,7 @@ from google.adk.models.lite_llm import _get_completion_inputs from google.adk.models.lite_llm import _get_content from google.adk.models.lite_llm import _message_to_generate_content_response from google.adk.models.lite_llm import _model_response_to_chunk +from google.adk.models.lite_llm import _model_response_to_generate_content_response from google.adk.models.lite_llm import _parse_tool_calls_from_text from google.adk.models.lite_llm import _split_message_content_and_tool_calls from google.adk.models.lite_llm import _to_litellm_response_format @@ -630,54 +630,6 @@ class MockLLMClient(LiteLLMClient): ) -def test_build_function_declaration_log(): - """Test that _build_function_declaration_log formats function declarations correctly.""" - # Test case 1: Function with parameters and response - func_decl1 = types.FunctionDeclaration( - name="test_func1", - description="Test function 1", - parameters=types.Schema( - type=types.Type.OBJECT, - properties={ - "param1": types.Schema( - type=types.Type.STRING, description="param1 desc" - ) - }, - ), - response=types.Schema(type=types.Type.BOOLEAN, description="return bool"), - ) - log1 = _build_function_declaration_log(func_decl1) - assert log1 == ( - "test_func1: {'param1': {'description': 'param1 desc', 'type':" - " }} -> {'description': 'return bool', 'type':" - " }" - ) - - # Test case 2: Function with JSON schema parameters and response - func_decl2 = types.FunctionDeclaration( - name="test_func2", - description="Test function 2", - parameters_json_schema={ - "type": "object", - "properties": {"param2": {"type": "integer"}}, - }, - response_json_schema={"type": "string"}, - ) - log2 = _build_function_declaration_log(func_decl2) - assert log2 == ( - "test_func2: {'type': 'object', 'properties': {'param2': {'type':" - " 'integer'}}} -> {'type': 'string'}" - ) - - # Test case 3: Function with no parameters and no response - func_decl3 = types.FunctionDeclaration( - name="test_func3", - description="Test function 3", - ) - log3 = _build_function_declaration_log(func_decl3) - assert log3 == "test_func3: {} -> None" - - @pytest.mark.asyncio async def test_generate_content_async(mock_acompletion, lite_llm_instance): @@ -1535,6 +1487,42 @@ def test_message_to_generate_content_response_with_model(): assert response.model_version == "gemini-2.5-pro" +def test_message_to_generate_content_response_reasoning_content(): + message = { + "role": "assistant", + "content": "Visible text", + "reasoning_content": "Hidden chain", + } + response = _message_to_generate_content_response(message) + + assert len(response.content.parts) == 2 + thought_part = response.content.parts[0] + text_part = response.content.parts[1] + assert thought_part.text == "Hidden chain" + assert thought_part.thought is True + assert text_part.text == "Visible text" + + +def test_model_response_to_generate_content_response_reasoning_content(): + model_response = ModelResponse( + model="thinking-model", + choices=[{ + "message": { + "role": "assistant", + "content": "Answer", + "reasoning_content": "Step-by-step", + }, + "finish_reason": "stop", + }], + ) + + response = _model_response_to_generate_content_response(model_response) + + assert response.content.parts[0].text == "Step-by-step" + assert response.content.parts[0].thought is True + assert response.content.parts[1].text == "Answer" + + def test_parse_tool_calls_from_text_multiple_calls(): text = ( '{"name":"alpha","arguments":{"value":1}}\n'