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'