wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
6339 lines
200 KiB
Python
6339 lines
200 KiB
Python
# Copyright 2026 Google LLC
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License
|
||
|
||
import base64
|
||
import contextlib
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from unittest.mock import ANY
|
||
from unittest.mock import AsyncMock
|
||
from unittest.mock import MagicMock
|
||
from unittest.mock import Mock
|
||
from unittest.mock import patch
|
||
import warnings
|
||
|
||
from google.adk.models.lite_llm import _aggregate_streaming_thought_parts
|
||
from google.adk.models.lite_llm import _append_fallback_user_content_if_missing
|
||
from google.adk.models.lite_llm import _BraceDepthTracker
|
||
from google.adk.models.lite_llm import _content_to_message_param
|
||
from google.adk.models.lite_llm import _convert_reasoning_value_to_parts
|
||
from google.adk.models.lite_llm import _enforce_strict_openai_schema
|
||
from google.adk.models.lite_llm import _extract_json_from_deepseek_args
|
||
from google.adk.models.lite_llm import _extract_reasoning_value
|
||
from google.adk.models.lite_llm import _extract_thought_signature_from_tool_call
|
||
from google.adk.models.lite_llm import _FILE_ID_REQUIRED_PROVIDERS
|
||
from google.adk.models.lite_llm import _FINISH_REASON_MAPPING
|
||
from google.adk.models.lite_llm import _function_declaration_to_tool_param
|
||
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 _get_provider_from_model
|
||
from google.adk.models.lite_llm import _is_anthropic_model
|
||
from google.adk.models.lite_llm import _is_anthropic_provider
|
||
from google.adk.models.lite_llm import _is_anthropic_route
|
||
from google.adk.models.lite_llm import _looks_like_openai_file_id
|
||
from google.adk.models.lite_llm import _message_to_generate_content_response
|
||
from google.adk.models.lite_llm import _MISSING_TOOL_RESULT_MESSAGE
|
||
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_deepseek_tool_calls_from_text
|
||
from google.adk.models.lite_llm import _parse_tool_calls_from_text
|
||
from google.adk.models.lite_llm import _redact_file_uri_for_log
|
||
from google.adk.models.lite_llm import _redirect_litellm_loggers_to_stdout
|
||
from google.adk.models.lite_llm import _safe_json_serialize
|
||
from google.adk.models.lite_llm import _schema_to_dict
|
||
from google.adk.models.lite_llm import _split_message_content_and_tool_calls
|
||
from google.adk.models.lite_llm import _THOUGHT_SIGNATURE_SEPARATOR
|
||
from google.adk.models.lite_llm import _to_litellm_response_format
|
||
from google.adk.models.lite_llm import _to_litellm_role
|
||
from google.adk.models.lite_llm import FunctionChunk
|
||
from google.adk.models.lite_llm import LiteLlm
|
||
from google.adk.models.lite_llm import LiteLLMClient
|
||
from google.adk.models.lite_llm import ReasoningChunk
|
||
from google.adk.models.lite_llm import TextChunk
|
||
from google.adk.models.lite_llm import UsageMetadataChunk
|
||
from google.adk.models.llm_request import LlmRequest
|
||
from google.genai import types
|
||
import litellm
|
||
from litellm import ChatCompletionAssistantMessage
|
||
from litellm import ChatCompletionMessageToolCall
|
||
from litellm import Function
|
||
from litellm.types.utils import ChatCompletionDeltaToolCall
|
||
from litellm.types.utils import Choices
|
||
from litellm.types.utils import Delta
|
||
from litellm.types.utils import ModelResponse
|
||
from litellm.types.utils import ModelResponseStream
|
||
from litellm.types.utils import StreamingChoices
|
||
from pydantic import BaseModel
|
||
from pydantic import Field
|
||
import pytest
|
||
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
tools=[
|
||
types.Tool(
|
||
function_declarations=[
|
||
types.FunctionDeclaration(
|
||
name="test_function",
|
||
description="Test function description",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"test_arg": types.Schema(
|
||
type=types.Type.STRING
|
||
),
|
||
"array_arg": types.Schema(
|
||
type=types.Type.ARRAY,
|
||
items={
|
||
"type": types.Type.STRING,
|
||
},
|
||
),
|
||
"nested_arg": types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"nested_key1": types.Schema(
|
||
type=types.Type.STRING
|
||
),
|
||
"nested_key2": types.Schema(
|
||
type=types.Type.STRING
|
||
),
|
||
},
|
||
),
|
||
},
|
||
),
|
||
)
|
||
]
|
||
)
|
||
],
|
||
),
|
||
)
|
||
|
||
FILE_URI_TEST_CASES = [
|
||
pytest.param("gs://bucket/document.pdf", "application/pdf", id="pdf"),
|
||
pytest.param("gs://bucket/data.json", "application/json", id="json"),
|
||
pytest.param("gs://bucket/data.txt", "text/plain", id="txt"),
|
||
]
|
||
|
||
FILE_BYTES_TEST_CASES = [
|
||
pytest.param(
|
||
b"test_pdf_data",
|
||
"application/pdf",
|
||
"data:application/pdf;base64,dGVzdF9wZGZfZGF0YQ==",
|
||
id="pdf",
|
||
),
|
||
pytest.param(
|
||
b'{"hello":"world"}',
|
||
"application/json",
|
||
"data:application/json;base64,eyJoZWxsbyI6IndvcmxkIn0=",
|
||
id="json",
|
||
),
|
||
]
|
||
|
||
STREAMING_MODEL_RESPONSE = [
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
content="zero, ",
|
||
),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
content="one, ",
|
||
),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
content="two:",
|
||
),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="test_tool_call_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg": "test_',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments='value"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason="tool_use",
|
||
)
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
class _StructuredOutput(BaseModel):
|
||
value: int = Field(description="Value to emit")
|
||
|
||
|
||
class _ModelDumpOnly:
|
||
"""Test helper that mimics objects exposing only model_dump."""
|
||
|
||
def __init__(self):
|
||
self._schema = {
|
||
"type": "object",
|
||
"properties": {"foo": {"type": "string"}},
|
||
}
|
||
|
||
def model_dump(self, *, exclude_none=True, mode="json"):
|
||
# The method signature matches pydantic BaseModel.model_dump to simulate
|
||
# google.genai schema-like objects.
|
||
del exclude_none
|
||
del mode
|
||
return self._schema
|
||
|
||
|
||
async def test_get_completion_inputs_formats_pydantic_schema_for_litellm():
|
||
llm_request = LlmRequest(
|
||
config=types.GenerateContentConfig(response_schema=_StructuredOutput)
|
||
)
|
||
|
||
_, _, response_format, _ = await _get_completion_inputs(
|
||
llm_request, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert response_format == {
|
||
"type": "json_object",
|
||
"response_schema": _StructuredOutput.model_json_schema(),
|
||
}
|
||
|
||
|
||
def test_to_litellm_response_format_passes_preformatted_dict():
|
||
response_format = {
|
||
"type": "json_object",
|
||
"response_schema": {
|
||
"type": "object",
|
||
"properties": {"foo": {"type": "string"}},
|
||
},
|
||
}
|
||
|
||
assert (
|
||
_to_litellm_response_format(
|
||
response_format, model="gemini/gemini-2.5-flash"
|
||
)
|
||
== response_format
|
||
)
|
||
|
||
|
||
def test_to_litellm_response_format_wraps_json_schema_dict():
|
||
schema = {
|
||
"type": "object",
|
||
"properties": {"foo": {"type": "string"}},
|
||
}
|
||
|
||
formatted = _to_litellm_response_format(
|
||
schema, model="gemini/gemini-2.5-flash"
|
||
)
|
||
assert formatted["type"] == "json_object"
|
||
assert formatted["response_schema"] == schema
|
||
|
||
|
||
def test_to_litellm_response_format_handles_model_dump_object():
|
||
schema_obj = _ModelDumpOnly()
|
||
|
||
formatted = _to_litellm_response_format(
|
||
schema_obj, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert formatted["type"] == "json_object"
|
||
assert formatted["response_schema"] == schema_obj.model_dump()
|
||
|
||
|
||
def test_to_litellm_response_format_handles_genai_schema_instance():
|
||
schema_instance = types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={"foo": types.Schema(type=types.Type.STRING)},
|
||
required=["foo"],
|
||
)
|
||
|
||
formatted = _to_litellm_response_format(
|
||
schema_instance, model="gemini/gemini-2.5-flash"
|
||
)
|
||
assert formatted["type"] == "json_object"
|
||
assert formatted["response_schema"] == schema_instance.model_dump(
|
||
exclude_none=True, mode="json"
|
||
)
|
||
|
||
|
||
def test_to_litellm_response_format_uses_json_schema_for_openai_model():
|
||
"""Test that OpenAI models use json_schema format instead of response_schema."""
|
||
formatted = _to_litellm_response_format(
|
||
_StructuredOutput, model="gpt-4o-mini"
|
||
)
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
assert "json_schema" in formatted
|
||
assert formatted["json_schema"]["name"] == "_StructuredOutput"
|
||
assert formatted["json_schema"]["strict"] is True
|
||
assert formatted["json_schema"]["schema"]["additionalProperties"] is False
|
||
assert "additionalProperties" in formatted["json_schema"]["schema"]
|
||
|
||
|
||
def test_to_litellm_response_format_uses_response_schema_for_gemini_model():
|
||
"""Test that Gemini models continue to use response_schema format."""
|
||
formatted = _to_litellm_response_format(
|
||
_StructuredOutput, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert formatted["type"] == "json_object"
|
||
assert "response_schema" in formatted
|
||
assert formatted["response_schema"] == _StructuredOutput.model_json_schema()
|
||
|
||
|
||
def test_to_litellm_response_format_uses_response_schema_for_vertex_gemini():
|
||
"""Test that Vertex AI Gemini models use response_schema format."""
|
||
formatted = _to_litellm_response_format(
|
||
_StructuredOutput, model="vertex_ai/gemini-2.5-flash"
|
||
)
|
||
|
||
assert formatted["type"] == "json_object"
|
||
assert "response_schema" in formatted
|
||
assert formatted["response_schema"] == _StructuredOutput.model_json_schema()
|
||
|
||
|
||
def test_to_litellm_response_format_uses_json_schema_for_azure_openai():
|
||
"""Test that Azure OpenAI models use json_schema format."""
|
||
formatted = _to_litellm_response_format(
|
||
_StructuredOutput, model="azure/gpt-4o"
|
||
)
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
assert "json_schema" in formatted
|
||
assert formatted["json_schema"]["name"] == "_StructuredOutput"
|
||
assert formatted["json_schema"]["strict"] is True
|
||
assert formatted["json_schema"]["schema"]["additionalProperties"] is False
|
||
assert "additionalProperties" in formatted["json_schema"]["schema"]
|
||
|
||
|
||
def test_to_litellm_response_format_uses_json_schema_for_anthropic():
|
||
"""Test that Anthropic models use json_schema format."""
|
||
formatted = _to_litellm_response_format(
|
||
_StructuredOutput, model="anthropic/claude-3-5-sonnet"
|
||
)
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
assert "json_schema" in formatted
|
||
assert formatted["json_schema"]["name"] == "_StructuredOutput"
|
||
assert formatted["json_schema"]["strict"] is True
|
||
assert formatted["json_schema"]["schema"]["additionalProperties"] is False
|
||
assert "additionalProperties" in formatted["json_schema"]["schema"]
|
||
|
||
|
||
def test_to_litellm_response_format_with_dict_schema_for_openai():
|
||
"""Test dict schema with OpenAI model uses json_schema format."""
|
||
schema = {
|
||
"type": "object",
|
||
"properties": {"foo": {"type": "string"}},
|
||
}
|
||
|
||
formatted = _to_litellm_response_format(schema, model="gpt-4o")
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
assert formatted["json_schema"]["name"] == "response"
|
||
assert formatted["json_schema"]["strict"] is True
|
||
assert formatted["json_schema"]["schema"]["additionalProperties"] is False
|
||
|
||
|
||
class _InnerModel(BaseModel):
|
||
value: str = Field(description="A value")
|
||
optional_field: str | None = Field(default=None, description="Optional")
|
||
|
||
|
||
class _OuterModel(BaseModel):
|
||
inner: _InnerModel = Field(description="Nested model")
|
||
name: str
|
||
|
||
|
||
class _WithList(BaseModel):
|
||
items: list[_InnerModel] = Field(description="List of items")
|
||
label: str
|
||
|
||
|
||
def test_enforce_strict_openai_schema_adds_additional_properties_recursively():
|
||
"""additionalProperties: false must appear on all object schemas."""
|
||
schema = _OuterModel.model_json_schema()
|
||
|
||
_enforce_strict_openai_schema(schema)
|
||
|
||
# Root level
|
||
assert schema["additionalProperties"] is False
|
||
# Nested model in $defs
|
||
inner_def = schema["$defs"]["_InnerModel"]
|
||
assert inner_def["additionalProperties"] is False
|
||
|
||
|
||
def test_enforce_strict_openai_schema_marks_all_properties_required():
|
||
"""All properties must appear in 'required', including optional fields."""
|
||
schema = _InnerModel.model_json_schema()
|
||
|
||
_enforce_strict_openai_schema(schema)
|
||
|
||
assert sorted(schema["required"]) == ["optional_field", "value"]
|
||
|
||
|
||
def test_enforce_strict_openai_schema_strips_ref_sibling_keywords():
|
||
"""$ref nodes must have no sibling keywords like 'description'."""
|
||
schema = _OuterModel.model_json_schema()
|
||
# Pydantic v2 generates {"$ref": "...", "description": "..."} for nested models
|
||
inner_prop = schema["properties"]["inner"]
|
||
assert "$ref" in inner_prop, "Expected Pydantic to generate a $ref property"
|
||
assert len(inner_prop) > 1, "Expected sibling keywords alongside $ref"
|
||
|
||
_enforce_strict_openai_schema(schema)
|
||
|
||
inner_prop = schema["properties"]["inner"]
|
||
assert list(inner_prop.keys()) == ["$ref"]
|
||
|
||
|
||
def test_enforce_strict_openai_schema_handles_array_items():
|
||
"""Array item schemas should also be recursively transformed."""
|
||
schema = _WithList.model_json_schema()
|
||
|
||
_enforce_strict_openai_schema(schema)
|
||
|
||
assert schema["additionalProperties"] is False
|
||
inner_def = schema["$defs"]["_InnerModel"]
|
||
assert inner_def["additionalProperties"] is False
|
||
assert sorted(inner_def["required"]) == ["optional_field", "value"]
|
||
|
||
|
||
def test_enforce_strict_openai_schema_preserves_anyof_and_default():
|
||
"""anyOf structure and default value for Optional fields must be preserved."""
|
||
schema = _InnerModel.model_json_schema()
|
||
|
||
_enforce_strict_openai_schema(schema)
|
||
|
||
opt_prop = schema["properties"]["optional_field"]
|
||
assert opt_prop["anyOf"] == [{"type": "string"}, {"type": "null"}]
|
||
assert opt_prop["default"] is None
|
||
|
||
|
||
def test_to_litellm_response_format_dict_input_not_mutated():
|
||
"""Passing a raw dict should not mutate the caller's original dict."""
|
||
schema = {
|
||
"type": "object",
|
||
"properties": {
|
||
"nested": {
|
||
"type": "object",
|
||
"properties": {"x": {"type": "string"}},
|
||
}
|
||
},
|
||
}
|
||
import copy
|
||
|
||
original = copy.deepcopy(schema)
|
||
|
||
_to_litellm_response_format(schema, model="gpt-4o")
|
||
|
||
assert schema == original, "Caller's input dict was mutated"
|
||
|
||
|
||
def test_to_litellm_response_format_instance_input_for_openai():
|
||
"""Passing a BaseModel instance should produce a valid strict schema."""
|
||
instance = _OuterModel(
|
||
inner=_InnerModel(value="test", optional_field=None), name="foo"
|
||
)
|
||
|
||
formatted = _to_litellm_response_format(instance, model="gpt-4o")
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
schema = formatted["json_schema"]["schema"]
|
||
assert schema["additionalProperties"] is False
|
||
inner_def = schema["$defs"]["_InnerModel"]
|
||
assert inner_def["additionalProperties"] is False
|
||
assert sorted(inner_def["required"]) == ["optional_field", "value"]
|
||
|
||
|
||
def test_to_litellm_response_format_nested_pydantic_for_openai():
|
||
"""Nested Pydantic model should produce a valid OpenAI strict schema."""
|
||
formatted = _to_litellm_response_format(_OuterModel, model="gpt-4o")
|
||
|
||
assert formatted["type"] == "json_schema"
|
||
assert formatted["json_schema"]["strict"] is True
|
||
|
||
schema = formatted["json_schema"]["schema"]
|
||
assert schema["additionalProperties"] is False
|
||
assert sorted(schema["required"]) == ["inner", "name"]
|
||
|
||
# $defs inner model must also be strict
|
||
inner_def = schema["$defs"]["_InnerModel"]
|
||
assert inner_def["additionalProperties"] is False
|
||
assert sorted(inner_def["required"]) == ["optional_field", "value"]
|
||
|
||
|
||
def test_to_litellm_response_format_nested_pydantic_for_gemini_unchanged():
|
||
"""Gemini models should NOT get the strict OpenAI transformations."""
|
||
formatted = _to_litellm_response_format(
|
||
_OuterModel, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert formatted["type"] == "json_object"
|
||
schema = formatted["response_schema"]
|
||
# Gemini path should pass through the raw Pydantic schema untouched
|
||
assert schema == _OuterModel.model_json_schema()
|
||
|
||
|
||
async def test_get_completion_inputs_uses_openai_format_for_openai_model():
|
||
"""Test that _get_completion_inputs produces OpenAI-compatible format."""
|
||
llm_request = LlmRequest(
|
||
model="gpt-4o-mini",
|
||
config=types.GenerateContentConfig(response_schema=_StructuredOutput),
|
||
)
|
||
|
||
_, _, response_format, _ = await _get_completion_inputs(
|
||
llm_request, model="gpt-4o-mini"
|
||
)
|
||
|
||
assert response_format["type"] == "json_schema"
|
||
assert "json_schema" in response_format
|
||
assert response_format["json_schema"]["name"] == "_StructuredOutput"
|
||
assert response_format["json_schema"]["strict"] is True
|
||
assert (
|
||
response_format["json_schema"]["schema"]["additionalProperties"] is False
|
||
)
|
||
|
||
|
||
async def test_get_completion_inputs_uses_gemini_format_for_gemini_model():
|
||
"""Test that _get_completion_inputs produces Gemini-compatible format."""
|
||
llm_request = LlmRequest(
|
||
model="gemini/gemini-2.5-flash",
|
||
config=types.GenerateContentConfig(response_schema=_StructuredOutput),
|
||
)
|
||
|
||
_, _, response_format, _ = await _get_completion_inputs(
|
||
llm_request, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert response_format["type"] == "json_object"
|
||
assert "response_schema" in response_format
|
||
|
||
|
||
async def test_get_completion_inputs_uses_passed_model_for_response_format():
|
||
"""Test that _get_completion_inputs uses the passed model parameter for response format.
|
||
|
||
This verifies that when llm_request.model is None, the explicit model parameter
|
||
is used to determine the correct response format (Gemini vs OpenAI).
|
||
"""
|
||
llm_request = LlmRequest(
|
||
model=None, # No model in request
|
||
config=types.GenerateContentConfig(response_schema=_StructuredOutput),
|
||
)
|
||
|
||
# Pass OpenAI model explicitly - should use json_schema format
|
||
_, _, response_format, _ = await _get_completion_inputs(
|
||
llm_request, model="gpt-4o-mini"
|
||
)
|
||
|
||
assert response_format["type"] == "json_schema"
|
||
assert "json_schema" in response_format
|
||
assert response_format["json_schema"]["name"] == "_StructuredOutput"
|
||
assert response_format["json_schema"]["strict"] is True
|
||
assert (
|
||
response_format["json_schema"]["schema"]["additionalProperties"] is False
|
||
)
|
||
|
||
|
||
async def test_get_completion_inputs_uses_passed_model_for_gemini_format():
|
||
"""Test that _get_completion_inputs uses passed model for Gemini response format.
|
||
|
||
This verifies that when self.model is a Gemini model and passed explicitly,
|
||
the response format uses the Gemini-specific format.
|
||
"""
|
||
llm_request = LlmRequest(
|
||
model=None, # No model in request
|
||
config=types.GenerateContentConfig(response_schema=_StructuredOutput),
|
||
)
|
||
|
||
# Pass Gemini model explicitly - should use response_schema format
|
||
_, _, response_format, _ = await _get_completion_inputs(
|
||
llm_request, model="gemini/gemini-2.5-flash"
|
||
)
|
||
|
||
assert response_format["type"] == "json_object"
|
||
assert "response_schema" in response_format
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_inserts_missing_tool_results():
|
||
user_content = types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Hi")]
|
||
)
|
||
assistant_content = types.Content(
|
||
role="assistant",
|
||
parts=[
|
||
types.Part.from_text(text="Calling tool."),
|
||
types.Part.from_function_call(
|
||
name="get_weather", args={"location": "Seoul"}
|
||
),
|
||
],
|
||
)
|
||
assistant_content.parts[1].function_call.id = "tool_call_1"
|
||
followup_user = types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Next question.")]
|
||
)
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[user_content, assistant_content, followup_user]
|
||
)
|
||
messages, _, _, _ = await _get_completion_inputs(
|
||
llm_request, model="openai/gpt-4o"
|
||
)
|
||
|
||
assert [message["role"] for message in messages] == [
|
||
"user",
|
||
"assistant",
|
||
"tool",
|
||
"user",
|
||
]
|
||
tool_message = messages[2]
|
||
assert tool_message["tool_call_id"] == "tool_call_1"
|
||
assert tool_message["content"] == _MISSING_TOOL_RESULT_MESSAGE
|
||
|
||
|
||
def test_schema_to_dict_filters_none_enum_values():
|
||
# Use model_construct to bypass strict enum validation.
|
||
top_level_schema = types.Schema.model_construct(
|
||
type=types.Type.STRING,
|
||
enum=["ACTIVE", None, "INACTIVE"],
|
||
)
|
||
nested_schema = types.Schema.model_construct(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"status": types.Schema.model_construct(
|
||
type=types.Type.STRING, enum=["READY", None, "DONE"]
|
||
),
|
||
},
|
||
)
|
||
|
||
assert _schema_to_dict(top_level_schema)["enum"] == ["ACTIVE", "INACTIVE"]
|
||
assert _schema_to_dict(nested_schema)["properties"]["status"]["enum"] == [
|
||
"READY",
|
||
"DONE",
|
||
]
|
||
|
||
|
||
def test_safe_json_serialize_serializable_object():
|
||
assert _safe_json_serialize({"a": 1, "b": [2, 3]}) == '{"a": 1, "b": [2, 3]}'
|
||
|
||
|
||
def test_safe_json_serialize_non_serializable_object_falls_back_to_str():
|
||
class _NotJsonable:
|
||
|
||
def __repr__(self):
|
||
return "<not jsonable>"
|
||
|
||
assert _safe_json_serialize(_NotJsonable()) == "<not jsonable>"
|
||
|
||
|
||
def test_safe_json_serialize_circular_dict_falls_back_to_str():
|
||
obj = {}
|
||
obj["self"] = obj
|
||
assert isinstance(_safe_json_serialize(obj), str)
|
||
|
||
|
||
def test_safe_json_serialize_circular_list_falls_back_to_str():
|
||
obj = []
|
||
obj.append(obj)
|
||
assert isinstance(_safe_json_serialize(obj), str)
|
||
|
||
|
||
def test_safe_json_serialize_recursion_error_falls_back_to_str():
|
||
with patch(
|
||
"google.adk.models.lite_llm.json.dumps",
|
||
side_effect=RecursionError("maximum recursion depth"),
|
||
):
|
||
assert _safe_json_serialize({"a": 1}) == str({"a": 1})
|
||
|
||
|
||
MULTIPLE_FUNCTION_CALLS_STREAM = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="call_1",
|
||
function=Function(
|
||
name="function_1",
|
||
arguments='{"arg": "val',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments='ue1"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="call_2",
|
||
function=Function(
|
||
name="function_2",
|
||
arguments='{"arg": "val',
|
||
),
|
||
index=1,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments='ue2"}',
|
||
),
|
||
index=1,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason="tool_calls",
|
||
)
|
||
]
|
||
),
|
||
]
|
||
|
||
|
||
STREAM_WITH_EMPTY_CHUNK = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="call_abc",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg":',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments=' "value"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
# This is the problematic empty chunk that should be ignored.
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments="",
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[StreamingChoices(finish_reason="tool_calls", delta=Delta())]
|
||
),
|
||
]
|
||
|
||
|
||
@pytest.fixture
|
||
def mock_response():
|
||
return ModelResponse(
|
||
model="test_model",
|
||
choices=[
|
||
Choices(
|
||
message=ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="test_tool_call_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg": "test_value"}',
|
||
),
|
||
)
|
||
],
|
||
)
|
||
)
|
||
],
|
||
)
|
||
|
||
|
||
# Test case reflecting litellm v1.71.2, ollama v0.9.0 streaming response
|
||
# no tool call ids
|
||
# indices all 0
|
||
# finish_reason stop instead of tool_calls
|
||
NON_COMPLIANT_MULTIPLE_FUNCTION_CALLS_STREAM = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name="function_1",
|
||
arguments='{"arg": "val',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments='ue1"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name="function_2",
|
||
arguments='{"arg": "val',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=None,
|
||
function=Function(
|
||
name=None,
|
||
arguments='ue2"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason="stop",
|
||
)
|
||
]
|
||
),
|
||
]
|
||
|
||
|
||
@pytest.fixture
|
||
def mock_acompletion(mock_response):
|
||
return AsyncMock(return_value=mock_response)
|
||
|
||
|
||
@pytest.fixture
|
||
def mock_completion(mock_response):
|
||
return Mock(return_value=mock_response)
|
||
|
||
|
||
@pytest.fixture
|
||
def mock_client(mock_acompletion, mock_completion):
|
||
return MockLLMClient(mock_acompletion, mock_completion)
|
||
|
||
|
||
@pytest.fixture
|
||
def lite_llm_instance(mock_client):
|
||
return LiteLlm(model="test_model", llm_client=mock_client)
|
||
|
||
|
||
class MockLLMClient(LiteLLMClient):
|
||
|
||
def __init__(self, acompletion_mock, completion_mock):
|
||
self.acompletion_mock = acompletion_mock
|
||
self.completion_mock = completion_mock
|
||
|
||
async def acompletion(self, model, messages, tools, **kwargs):
|
||
if kwargs.get("stream", False):
|
||
kwargs_copy = dict(kwargs)
|
||
kwargs_copy.pop("stream", None)
|
||
|
||
async def stream_generator():
|
||
stream_data = self.completion_mock(
|
||
model=model,
|
||
messages=messages,
|
||
tools=tools,
|
||
stream=True,
|
||
**kwargs_copy,
|
||
)
|
||
for item in stream_data:
|
||
yield item
|
||
|
||
return stream_generator()
|
||
else:
|
||
return await self.acompletion_mock(
|
||
model=model, messages=messages, tools=tools, **kwargs
|
||
)
|
||
|
||
def completion(self, model, messages, tools, stream, **kwargs):
|
||
return self.completion_mock(
|
||
model=model, messages=messages, tools=tools, stream=stream, **kwargs
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async(mock_acompletion, lite_llm_instance):
|
||
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION
|
||
):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
assert response.content.parts[1].function_call.name == "test_function"
|
||
assert response.content.parts[1].function_call.args == {
|
||
"test_arg": "test_value"
|
||
}
|
||
assert response.content.parts[1].function_call.id == "test_tool_call_id"
|
||
assert response.model_version == "test_model"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
assert kwargs["tools"][0]["function"]["name"] == "test_function"
|
||
assert (
|
||
kwargs["tools"][0]["function"]["description"]
|
||
== "Test function description"
|
||
)
|
||
assert (
|
||
kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
|
||
"type"
|
||
]
|
||
== "string"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_with_model_override(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
llm_request = LlmRequest(
|
||
model="overridden_model",
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "overridden_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_without_model_override(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
llm_request = LlmRequest(
|
||
model=None,
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_adds_fallback_user_message(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[],
|
||
)
|
||
]
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
user_messages = [
|
||
message for message in kwargs["messages"] if message["role"] == "user"
|
||
]
|
||
assert any(
|
||
message.get("content")
|
||
== "Handle the requests as specified in the System Instruction."
|
||
for message in user_messages
|
||
)
|
||
assert (
|
||
sum(1 for content in llm_request.contents if content.role == "user") == 1
|
||
)
|
||
assert llm_request.contents[-1].parts[0].text == (
|
||
"Handle the requests as specified in the System Instruction."
|
||
)
|
||
|
||
|
||
def test_append_fallback_user_content_ignores_function_response_parts():
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_function_response(
|
||
name="add", response={"result": 6}
|
||
)
|
||
],
|
||
)
|
||
]
|
||
)
|
||
|
||
_append_fallback_user_content_if_missing(llm_request)
|
||
|
||
assert len(llm_request.contents) == 1
|
||
assert len(llm_request.contents[0].parts) == 1
|
||
assert llm_request.contents[0].parts[0].function_response is not None
|
||
assert llm_request.contents[0].parts[0].text is None
|
||
|
||
|
||
litellm_append_user_content_test_cases = [
|
||
pytest.param(
|
||
LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="developer",
|
||
parts=[types.Part.from_text(text="Test prompt")],
|
||
)
|
||
]
|
||
),
|
||
2,
|
||
id="litellm request without user content",
|
||
),
|
||
pytest.param(
|
||
LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[types.Part.from_text(text="user prompt")],
|
||
)
|
||
]
|
||
),
|
||
1,
|
||
id="litellm request with user content",
|
||
),
|
||
pytest.param(
|
||
LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="model",
|
||
parts=[types.Part.from_text(text="model prompt")],
|
||
),
|
||
types.Content(
|
||
role="user",
|
||
parts=[types.Part.from_text(text="user prompt")],
|
||
),
|
||
types.Content(
|
||
role="model",
|
||
parts=[types.Part.from_text(text="model prompt")],
|
||
),
|
||
]
|
||
),
|
||
4,
|
||
id="user content is not the last message scenario",
|
||
),
|
||
]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"llm_request, expected_output", litellm_append_user_content_test_cases
|
||
)
|
||
def test_maybe_append_user_content(
|
||
lite_llm_instance, llm_request, expected_output
|
||
):
|
||
|
||
lite_llm_instance._maybe_append_user_content(llm_request)
|
||
|
||
assert len(llm_request.contents) == expected_output
|
||
|
||
|
||
function_declaration_test_cases = [
|
||
(
|
||
"simple_function",
|
||
types.FunctionDeclaration(
|
||
name="test_function",
|
||
description="Test function description",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"test_arg": types.Schema(type=types.Type.STRING),
|
||
"array_arg": types.Schema(
|
||
type=types.Type.ARRAY,
|
||
items=types.Schema(
|
||
type=types.Type.STRING,
|
||
),
|
||
),
|
||
"nested_arg": types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"nested_key1": types.Schema(type=types.Type.STRING),
|
||
"nested_key2": types.Schema(type=types.Type.STRING),
|
||
},
|
||
required=["nested_key1"],
|
||
),
|
||
},
|
||
required=["nested_arg"],
|
||
),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function",
|
||
"description": "Test function description",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"test_arg": {"type": "string"},
|
||
"array_arg": {
|
||
"items": {"type": "string"},
|
||
"type": "array",
|
||
},
|
||
"nested_arg": {
|
||
"properties": {
|
||
"nested_key1": {"type": "string"},
|
||
"nested_key2": {"type": "string"},
|
||
},
|
||
"type": "object",
|
||
"required": ["nested_key1"],
|
||
},
|
||
},
|
||
"required": ["nested_arg"],
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"no_description",
|
||
types.FunctionDeclaration(
|
||
name="test_function_no_description",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"test_arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_no_description",
|
||
"description": "",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"test_arg": {"type": "string"},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"empty_parameters",
|
||
types.FunctionDeclaration(
|
||
name="test_function_empty_params",
|
||
parameters=types.Schema(type=types.Type.OBJECT, properties={}),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_empty_params",
|
||
"description": "",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"nested_array",
|
||
types.FunctionDeclaration(
|
||
name="test_function_nested_array",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"array_arg": types.Schema(
|
||
type=types.Type.ARRAY,
|
||
items=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"nested_key": types.Schema(
|
||
type=types.Type.STRING
|
||
)
|
||
},
|
||
),
|
||
),
|
||
},
|
||
),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_nested_array",
|
||
"description": "",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"array_arg": {
|
||
"items": {
|
||
"properties": {
|
||
"nested_key": {"type": "string"}
|
||
},
|
||
"type": "object",
|
||
},
|
||
"type": "array",
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"nested_properties",
|
||
types.FunctionDeclaration(
|
||
name="test_function_nested_properties",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"array_arg": types.Schema(
|
||
type=types.Type.ARRAY,
|
||
items=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"nested_key": types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"inner_key": types.Schema(
|
||
type=types.Type.STRING,
|
||
)
|
||
},
|
||
)
|
||
},
|
||
),
|
||
),
|
||
},
|
||
),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_nested_properties",
|
||
"description": "",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"array_arg": {
|
||
"items": {
|
||
"type": "object",
|
||
"properties": {
|
||
"nested_key": {
|
||
"type": "object",
|
||
"properties": {
|
||
"inner_key": {"type": "string"},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
"type": "array",
|
||
},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"no_parameters",
|
||
types.FunctionDeclaration(
|
||
name="test_function_no_params",
|
||
description="Test function with no parameters",
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_no_params",
|
||
"description": "Test function with no parameters",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
(
|
||
"parameters_without_required",
|
||
types.FunctionDeclaration(
|
||
name="test_function_no_required",
|
||
description="Test function with parameters but no required field",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"optional_arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_no_required",
|
||
"description": (
|
||
"Test function with parameters but no required field"
|
||
),
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"optional_arg": {"type": "string"},
|
||
},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"_, function_declaration, expected_output",
|
||
function_declaration_test_cases,
|
||
ids=[case[0] for case in function_declaration_test_cases],
|
||
)
|
||
def test_function_declaration_to_tool_param(
|
||
_, function_declaration, expected_output
|
||
):
|
||
assert (
|
||
_function_declaration_to_tool_param(function_declaration)
|
||
== expected_output
|
||
)
|
||
|
||
|
||
def test_function_declaration_to_tool_param_without_required_attribute():
|
||
"""Ensure tools without a required field attribute don't raise errors."""
|
||
|
||
class SchemaWithoutRequired:
|
||
"""Mimics a Schema object that lacks the required attribute."""
|
||
|
||
def __init__(self):
|
||
self.properties = {
|
||
"optional_arg": types.Schema(type=types.Type.STRING),
|
||
}
|
||
|
||
func_decl = types.FunctionDeclaration(
|
||
name="function_without_required_attr",
|
||
description="Function missing required attribute",
|
||
)
|
||
func_decl.parameters = SchemaWithoutRequired()
|
||
|
||
expected = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "function_without_required_attr",
|
||
"description": "Function missing required attribute",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"optional_arg": {"type": "string"},
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
assert _function_declaration_to_tool_param(func_decl) == expected
|
||
|
||
|
||
def test_function_declaration_to_tool_param_with_parameters_json_schema():
|
||
"""Ensure function declarations using parameters_json_schema are handled.
|
||
|
||
This verifies that when a FunctionDeclaration includes a raw
|
||
`parameters_json_schema` dict, it is used directly as the function
|
||
parameters in the resulting tool param.
|
||
"""
|
||
|
||
func_decl = types.FunctionDeclaration(
|
||
name="fn_with_json",
|
||
description="desc",
|
||
parameters_json_schema={
|
||
"type": "object",
|
||
"properties": {
|
||
"a": {"type": "string"},
|
||
"b": {"type": "array", "items": {"type": "string"}},
|
||
},
|
||
"required": ["a"],
|
||
},
|
||
)
|
||
|
||
expected = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "fn_with_json",
|
||
"description": "desc",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"a": {"type": "string"},
|
||
"b": {"type": "array", "items": {"type": "string"}},
|
||
},
|
||
"required": ["a"],
|
||
},
|
||
},
|
||
}
|
||
|
||
assert _function_declaration_to_tool_param(func_decl) == expected
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_with_system_instruction(
|
||
lite_llm_instance, mock_acompletion
|
||
):
|
||
mock_response_with_system_instruction = ModelResponse(
|
||
choices=[
|
||
Choices(
|
||
message=ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
)
|
||
)
|
||
]
|
||
)
|
||
mock_acompletion.return_value = mock_response_with_system_instruction
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
system_instruction="Test system instruction"
|
||
),
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
assert kwargs["messages"][0]["role"] == "system"
|
||
assert kwargs["messages"][0]["content"] == "Test system instruction"
|
||
assert kwargs["messages"][1]["role"] == "user"
|
||
assert kwargs["messages"][1]["content"] == "Test prompt"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_with_tool_response(
|
||
lite_llm_instance, mock_acompletion
|
||
):
|
||
mock_response_with_tool_response = ModelResponse(
|
||
choices=[
|
||
Choices(
|
||
message=ChatCompletionAssistantMessage(
|
||
role="tool",
|
||
content='{"result": "test_result"}',
|
||
tool_call_id="test_tool_call_id",
|
||
)
|
||
)
|
||
]
|
||
)
|
||
mock_acompletion.return_value = mock_response_with_tool_response
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
),
|
||
types.Content(
|
||
role="tool",
|
||
parts=[
|
||
types.Part.from_function_response(
|
||
name="test_function",
|
||
response={"result": "test_result"},
|
||
)
|
||
],
|
||
),
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
system_instruction="test instruction",
|
||
),
|
||
)
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == '{"result": "test_result"}'
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
|
||
assert kwargs["messages"][2]["role"] == "tool"
|
||
assert kwargs["messages"][2]["content"] == '{"result": "test_result"}'
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_with_usage_metadata(
|
||
lite_llm_instance, mock_acompletion
|
||
):
|
||
mock_response_with_usage_metadata = ModelResponse(
|
||
choices=[
|
||
Choices(
|
||
message=ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
)
|
||
)
|
||
],
|
||
usage={
|
||
"prompt_tokens": 10,
|
||
"completion_tokens": 5,
|
||
"total_tokens": 15,
|
||
"cached_tokens": 8,
|
||
"completion_tokens_details": {"reasoning_tokens": 5},
|
||
},
|
||
)
|
||
mock_acompletion.return_value = mock_response_with_usage_metadata
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
),
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
system_instruction="test instruction",
|
||
),
|
||
)
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
assert response.usage_metadata.prompt_token_count == 10
|
||
assert response.usage_metadata.candidates_token_count == 5
|
||
assert response.usage_metadata.total_token_count == 15
|
||
assert response.usage_metadata.cached_content_token_count == 8
|
||
assert response.usage_metadata.thoughts_token_count == 5
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_ollama_chat_preserves_multimodal_content(
|
||
mock_acompletion, mock_completion
|
||
):
|
||
llm_client = MockLLMClient(mock_acompletion, mock_completion)
|
||
lite_llm_instance = LiteLlm(
|
||
model="ollama_chat/qwen2.5:7b", llm_client=llm_client
|
||
)
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Describe this image."),
|
||
types.Part.from_bytes(
|
||
data=b"test_image", mime_type="image/png"
|
||
),
|
||
],
|
||
)
|
||
]
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once_with(
|
||
model="ollama_chat/qwen2.5:7b",
|
||
messages=ANY,
|
||
tools=ANY,
|
||
response_format=ANY,
|
||
)
|
||
_, kwargs = mock_acompletion.call_args
|
||
message_content = kwargs["messages"][0]["content"]
|
||
# Multimodal content (text + image) should be kept as a list so LiteLLM
|
||
# can convert it to Ollama's native images field.
|
||
assert isinstance(message_content, list)
|
||
text_blocks = [
|
||
b
|
||
for b in message_content
|
||
if isinstance(b, dict) and b.get("type") == "text"
|
||
]
|
||
image_blocks = [
|
||
b
|
||
for b in message_content
|
||
if isinstance(b, dict) and b.get("type") == "image_url"
|
||
]
|
||
assert len(text_blocks) >= 1
|
||
assert "Describe this image." in text_blocks[0].get("text", "")
|
||
assert len(image_blocks) >= 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_custom_provider_preserves_multimodal(
|
||
mock_acompletion, mock_completion
|
||
):
|
||
llm_client = MockLLMClient(mock_acompletion, mock_completion)
|
||
lite_llm_instance = LiteLlm(
|
||
model="qwen2.5:7b",
|
||
llm_client=llm_client,
|
||
custom_llm_provider="ollama_chat",
|
||
)
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Describe this image."),
|
||
types.Part.from_bytes(
|
||
data=b"test_image", mime_type="image/png"
|
||
),
|
||
],
|
||
)
|
||
]
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["custom_llm_provider"] == "ollama_chat"
|
||
assert kwargs["model"] == "qwen2.5:7b"
|
||
message_content = kwargs["messages"][0]["content"]
|
||
# Multimodal content should be preserved as a list.
|
||
assert isinstance(message_content, list)
|
||
text_blocks = [
|
||
b
|
||
for b in message_content
|
||
if isinstance(b, dict) and b.get("type") == "text"
|
||
]
|
||
assert any("Describe this image." in b.get("text", "") for b in text_blocks)
|
||
|
||
|
||
def test_flatten_ollama_content_accepts_tuple_blocks():
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
content = (
|
||
{"type": "text", "text": "first"},
|
||
{"type": "text", "text": "second"},
|
||
)
|
||
flattened = _flatten_ollama_content(content)
|
||
assert flattened == "first\nsecond"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"content, expected",
|
||
[
|
||
(None, None),
|
||
("hello", "hello"),
|
||
(
|
||
[
|
||
{"type": "text", "text": "first"},
|
||
{"type": "text", "text": "second"},
|
||
],
|
||
"first\nsecond",
|
||
),
|
||
],
|
||
)
|
||
def test_flatten_ollama_content_returns_str_or_none(content, expected):
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
flattened = _flatten_ollama_content(content)
|
||
assert flattened == expected
|
||
assert flattened is None or isinstance(flattened, str)
|
||
|
||
|
||
def test_flatten_ollama_content_preserves_image_url_blocks():
|
||
"""Media blocks should be kept as a list so LiteLLM can convert them."""
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
blocks = [
|
||
{"type": "image_url", "image_url": {"url": "http://example.com/img.png"}},
|
||
]
|
||
result = _flatten_ollama_content(blocks)
|
||
assert isinstance(result, list)
|
||
assert result == blocks
|
||
|
||
|
||
def test_flatten_ollama_content_preserves_mixed_text_and_image():
|
||
"""Text + image_url should return the full list, not just the text."""
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
blocks = [
|
||
{"type": "text", "text": "Describe this image."},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
|
||
},
|
||
]
|
||
result = _flatten_ollama_content(blocks)
|
||
assert isinstance(result, list)
|
||
assert len(result) == 2
|
||
assert result[0]["type"] == "text"
|
||
assert result[1]["type"] == "image_url"
|
||
|
||
|
||
def test_flatten_ollama_content_preserves_video_url_blocks():
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
blocks = [
|
||
{"type": "text", "text": "What happens in this clip?"},
|
||
{"type": "video_url", "video_url": {"url": "http://example.com/v.mp4"}},
|
||
]
|
||
result = _flatten_ollama_content(blocks)
|
||
assert isinstance(result, list)
|
||
assert len(result) == 2
|
||
|
||
|
||
def test_flatten_ollama_content_serializes_non_media_non_text_blocks_to_json():
|
||
"""Blocks with unknown types and no media should still serialize to JSON."""
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
blocks = [
|
||
{"type": "custom_block", "data": "something"},
|
||
]
|
||
result = _flatten_ollama_content(blocks)
|
||
assert isinstance(result, str)
|
||
assert json.loads(result) == blocks
|
||
|
||
|
||
def test_flatten_ollama_content_serializes_dict_to_json():
|
||
from google.adk.models.lite_llm import _flatten_ollama_content
|
||
|
||
content = {"type": "image_url", "image_url": {"url": "http://example.com"}}
|
||
flattened = _flatten_ollama_content(content)
|
||
assert isinstance(flattened, str)
|
||
assert json.loads(flattened) == content
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_user_message():
|
||
content = types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "user"
|
||
assert message["content"] == "Test prompt"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
|
||
async def test_content_to_message_param_user_message_with_file_uri(
|
||
file_uri, mime_type
|
||
):
|
||
file_part = types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Summarize this file."),
|
||
file_part,
|
||
],
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
assert message == {
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "Summarize this file."},
|
||
{"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
|
||
],
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
|
||
async def test_content_to_message_param_user_message_file_uri_only(
|
||
file_uri, mime_type
|
||
):
|
||
file_part = types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[
|
||
file_part,
|
||
],
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
assert message == {
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
|
||
],
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_user_message_file_uri_without_mime_type():
|
||
"""Test that file_data without an inferable mime_type raises ValueError.
|
||
|
||
When using GcsArtifactService, artifacts may have file_uri (gs://...) but
|
||
without mime_type set. When the MIME type cannot be determined from the URI
|
||
extension or display_name, ADK raises a clear ValueError rather than
|
||
forwarding an unsupported 'application/octet-stream' to LiteLLM.
|
||
"""
|
||
file_part = types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://agent-artifact-bucket/app/user/session/artifact/0"
|
||
)
|
||
)
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Analyze this file."),
|
||
file_part,
|
||
],
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="Cannot process file_uri"):
|
||
await _content_to_message_param(content)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_user_message_file_uri_explicit_octet_stream():
|
||
"""Test that an explicit application/octet-stream MIME type raises ValueError.
|
||
|
||
Upstream callers may explicitly set mime_type to 'application/octet-stream'
|
||
when the true type is unknown. ADK treats this identically to a missing MIME
|
||
type and raises early rather than forwarding the unsupported type to LiteLLM.
|
||
"""
|
||
file_part = types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://agent-artifact-bucket/app/user/session/artifact/0",
|
||
mime_type="application/octet-stream",
|
||
)
|
||
)
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Analyze this file."),
|
||
file_part,
|
||
],
|
||
)
|
||
|
||
with pytest.raises(ValueError, match="application/octet-stream"):
|
||
await _content_to_message_param(content)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_user_message_file_uri_infer_mime_type():
|
||
"""Test MIME type inference from file_uri extension.
|
||
|
||
When file_data has a file_uri with a recognizable extension but no explicit
|
||
mime_type, the MIME type should be inferred from the extension.
|
||
"""
|
||
file_part = types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf",
|
||
)
|
||
)
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[file_part],
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
assert message == {
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "file",
|
||
"file": {
|
||
"file_id": "gs://bucket/path/to/document.pdf",
|
||
"format": "application/pdf",
|
||
},
|
||
},
|
||
],
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_multi_part_function_response():
|
||
part1 = types.Part.from_function_response(
|
||
name="function_one",
|
||
response={"result": "result_one"},
|
||
)
|
||
part1.function_response.id = "tool_call_1"
|
||
|
||
part2 = types.Part.from_function_response(
|
||
name="function_two",
|
||
response={"value": 123},
|
||
)
|
||
part2.function_response.id = "tool_call_2"
|
||
|
||
content = types.Content(
|
||
role="tool",
|
||
parts=[part1, part2],
|
||
)
|
||
messages = await _content_to_message_param(content)
|
||
assert isinstance(messages, list)
|
||
assert len(messages) == 2
|
||
|
||
assert messages[0]["role"] == "tool"
|
||
assert messages[0]["tool_call_id"] == "tool_call_1"
|
||
assert messages[0]["content"] == '{"result": "result_one"}'
|
||
|
||
assert messages[1]["role"] == "tool"
|
||
assert messages[1]["tool_call_id"] == "tool_call_2"
|
||
assert messages[1]["content"] == '{"value": 123}'
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_function_response_with_extra_parts():
|
||
tool_part = types.Part.from_function_response(
|
||
name="load_image",
|
||
response={"status": "success"},
|
||
)
|
||
tool_part.function_response.id = "tool_call_1"
|
||
|
||
text_part = types.Part.from_text(text="[Image: img_123.png]")
|
||
image_bytes = b"test_image_data"
|
||
image_part = types.Part.from_bytes(data=image_bytes, mime_type="image/png")
|
||
|
||
content = types.Content(
|
||
role="user",
|
||
parts=[tool_part, text_part, image_part],
|
||
)
|
||
|
||
messages = await _content_to_message_param(content)
|
||
assert isinstance(messages, list)
|
||
assert messages == [
|
||
{
|
||
"role": "tool",
|
||
"tool_call_id": "tool_call_1",
|
||
"content": '{"status": "success"}',
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": "[Image: img_123.png]"},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": "data:image/png;base64,dGVzdF9pbWFnZV9kYXRh"
|
||
},
|
||
},
|
||
],
|
||
},
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_function_response_preserves_string():
|
||
"""Tests that string responses are used directly without double-serialization.
|
||
|
||
The google.genai FunctionResponse.response field is typed as dict, but
|
||
_content_to_message_param defensively handles string responses to avoid
|
||
double-serialization. This test verifies that behavior by mocking a
|
||
function_response with a string response attribute.
|
||
"""
|
||
response_payload = '{"type": "files", "count": 2}'
|
||
|
||
# Create a Part with a dict response, then mock the response to be a string
|
||
# to simulate edge cases where response might be set directly as a string
|
||
part = types.Part.from_function_response(
|
||
name="list_files",
|
||
response={"placeholder": "will be mocked"},
|
||
)
|
||
|
||
# Mock the response attribute to return a string
|
||
# Using Mock without spec_set to allow setting response to a string,
|
||
# which simulates the edge case we're testing
|
||
mock_function_response = Mock(spec=types.FunctionResponse)
|
||
mock_function_response.response = response_payload
|
||
mock_function_response.id = "tool_call_1"
|
||
part.function_response = mock_function_response
|
||
|
||
content = types.Content(
|
||
role="tool",
|
||
parts=[part],
|
||
)
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "tool"
|
||
assert message["tool_call_id"] == "tool_call_1"
|
||
assert message["content"] == response_payload
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_assistant_message():
|
||
content = types.Content(
|
||
role="assistant", parts=[types.Part.from_text(text="Test response")]
|
||
)
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] == "Test response"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_user_filters_thought_parts():
|
||
thought_part = types.Part.from_text(text="internal reasoning")
|
||
thought_part.thought = True
|
||
content_part = types.Part.from_text(text="visible content")
|
||
content = types.Content(role="user", parts=[thought_part, content_part])
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "user"
|
||
assert message["content"] == "visible content"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_assistant_thought_message():
|
||
part = types.Part.from_text(text="internal reasoning")
|
||
part.thought = True
|
||
content = types.Content(role="assistant", parts=[part])
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] is None
|
||
assert message["reasoning_content"] == "internal reasoning"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_merges_reasoning_chunks_without_separator():
|
||
first_part = types.Part.from_text(text="Let")
|
||
first_part.thought = True
|
||
second_part = types.Part.from_text(text=" me think")
|
||
second_part.thought = True
|
||
third_part = types.Part.from_text(text=" this through.")
|
||
third_part.thought = True
|
||
content = types.Content(
|
||
role="assistant", parts=[first_part, second_part, third_part]
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] is None
|
||
assert message["reasoning_content"] == "Let me think this through."
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_model_thought_message():
|
||
part = types.Part.from_text(text="internal reasoning")
|
||
part.thought = True
|
||
content = types.Content(role="model", parts=[part])
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] is None
|
||
assert message["reasoning_content"] == "internal reasoning"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_assistant_thought_and_content_message():
|
||
thought_part = types.Part.from_text(text="internal reasoning")
|
||
thought_part.thought = True
|
||
content_part = types.Part.from_text(text="visible content")
|
||
content = types.Content(role="assistant", parts=[thought_part, content_part])
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] == "visible content"
|
||
assert message["reasoning_content"] == "internal reasoning"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_preserves_chunked_reasoning_deltas():
|
||
thought_part_1 = types.Part.from_text(text="Hel")
|
||
thought_part_1.thought = True
|
||
thought_part_2 = types.Part.from_text(text="lo")
|
||
thought_part_2.thought = True
|
||
content = types.Content(
|
||
role="assistant", parts=[thought_part_1, thought_part_2]
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] is None
|
||
assert message["reasoning_content"] == "Hello"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_preserves_reasoning_newlines():
|
||
thought_part_1 = types.Part.from_text(text="line 1\n")
|
||
thought_part_1.thought = True
|
||
thought_part_2 = types.Part.from_text(text="line 2")
|
||
thought_part_2.thought = True
|
||
content = types.Content(
|
||
role="assistant", parts=[thought_part_1, thought_part_2]
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
|
||
assert message["reasoning_content"] == "line 1\nline 2"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_function_call():
|
||
content = types.Content(
|
||
role="assistant",
|
||
parts=[
|
||
types.Part.from_text(text="test response"),
|
||
types.Part.from_function_call(
|
||
name="test_function", args={"test_arg": "test_value"}
|
||
),
|
||
],
|
||
)
|
||
content.parts[1].function_call.id = "test_tool_call_id"
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "assistant"
|
||
assert message["content"] == "test response"
|
||
|
||
tool_call = message["tool_calls"][0]
|
||
assert tool_call["type"] == "function"
|
||
assert tool_call["id"] == "test_tool_call_id"
|
||
assert tool_call["function"]["name"] == "test_function"
|
||
assert tool_call["function"]["arguments"] == '{"test_arg": "test_value"}'
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_multipart_content():
|
||
"""Test handling of multipart content where final_content is a list with text objects."""
|
||
content = types.Content(
|
||
role="assistant",
|
||
parts=[
|
||
types.Part.from_text(text="text part"),
|
||
types.Part.from_bytes(data=b"test_image_data", mime_type="image/png"),
|
||
],
|
||
)
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "assistant"
|
||
# When content is a list and the first element is a text object with type "text",
|
||
# it should extract the text (for providers like ollama_chat that don't handle lists well)
|
||
# This is the behavior implemented in the fix
|
||
assert message["content"] == "text part"
|
||
assert message["tool_calls"] is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_single_text_object_in_list(mocker):
|
||
"""Test extraction of text from single text object in list (for ollama_chat compatibility)."""
|
||
from google.adk.models import lite_llm
|
||
|
||
# Mock _get_content to return a list with single text object
|
||
async def mock_get_content(*args, **kwargs):
|
||
return [{"type": "text", "text": "single text"}]
|
||
|
||
mocker.patch.object(lite_llm, "_get_content", side_effect=mock_get_content)
|
||
|
||
content = types.Content(
|
||
role="assistant",
|
||
parts=[types.Part.from_text(text="single text")],
|
||
)
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "assistant"
|
||
# Should extract the text from the single text object
|
||
assert message["content"] == "single text"
|
||
assert message["tool_calls"] is None
|
||
|
||
|
||
def test_message_to_generate_content_response_text():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
)
|
||
response = _message_to_generate_content_response(message)
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
|
||
|
||
def test_message_to_generate_content_response_tool_call():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="test_tool_call_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg": "test_value"}',
|
||
),
|
||
)
|
||
],
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].function_call.name == "test_function"
|
||
assert response.content.parts[0].function_call.args == {
|
||
"test_arg": "test_value"
|
||
}
|
||
assert response.content.parts[0].function_call.id == "test_tool_call_id"
|
||
|
||
|
||
def test_message_to_generate_content_response_tool_call_accepts_python_literal_arguments():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="test_tool_call_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments="{'query': 'MATCH (n) RETURN n'}",
|
||
),
|
||
)
|
||
],
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].function_call.name == "test_function"
|
||
assert response.content.parts[0].function_call.args == {
|
||
"query": "MATCH (n) RETURN n"
|
||
}
|
||
|
||
|
||
def test_message_to_generate_content_response_tool_call_accepts_unquoted_json_keys():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="test_tool_call_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{query: "MATCH (n) RETURN n", limit: 5}',
|
||
),
|
||
)
|
||
],
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].function_call.name == "test_function"
|
||
assert response.content.parts[0].function_call.args == {
|
||
"query": "MATCH (n) RETURN n",
|
||
"limit": 5,
|
||
}
|
||
|
||
|
||
def test_message_to_generate_content_response_inline_tool_call_text():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=(
|
||
'{"id":"inline_call","name":"get_current_time",'
|
||
'"arguments":{"timezone_str":"Asia/Taipei"}} <|im_end|>system'
|
||
),
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
assert len(response.content.parts) == 2
|
||
text_part = response.content.parts[0]
|
||
tool_part = response.content.parts[1]
|
||
assert text_part.text == "<|im_end|>system"
|
||
assert tool_part.function_call.name == "get_current_time"
|
||
assert tool_part.function_call.args == {"timezone_str": "Asia/Taipei"}
|
||
assert tool_part.function_call.id == "inline_call"
|
||
|
||
|
||
def test_message_to_generate_content_response_with_model():
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
)
|
||
response = _message_to_generate_content_response(
|
||
message, model_version="gemini-2.5-pro"
|
||
)
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
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_message_to_generate_content_response_reasoning_field():
|
||
"""Test that the 'reasoning' field is supported (LM Studio, vLLM)."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Final answer",
|
||
"reasoning": "Thinking process",
|
||
}
|
||
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 == "Thinking process"
|
||
assert thought_part.thought is True
|
||
assert text_part.text == "Final answer"
|
||
|
||
|
||
def test_model_response_to_generate_content_response_reasoning_field():
|
||
"""Test that 'reasoning' field is supported in ModelResponse."""
|
||
model_response = ModelResponse(
|
||
model="test-model",
|
||
choices=[{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": "Result",
|
||
"reasoning": "Chain of thought",
|
||
},
|
||
"finish_reason": "stop",
|
||
}],
|
||
)
|
||
|
||
response = _model_response_to_generate_content_response(model_response)
|
||
|
||
assert response.content.parts[0].text == "Chain of thought"
|
||
assert response.content.parts[0].thought is True
|
||
assert response.content.parts[1].text == "Result"
|
||
|
||
|
||
def test_model_response_to_generate_content_response_grounding_metadata_dict():
|
||
"""vertex_ai_grounding_metadata as a dict is propagated to the LlmResponse."""
|
||
model_response = ModelResponse(
|
||
model="gemini/gemini-2.5-flash",
|
||
choices=[{
|
||
"message": {"role": "assistant", "content": "Answer"},
|
||
"finish_reason": "stop",
|
||
}],
|
||
)
|
||
model_response.vertex_ai_grounding_metadata = {
|
||
"grounding_chunks": [
|
||
{"web": {"uri": "https://example.com", "title": "Example"}}
|
||
],
|
||
}
|
||
|
||
response = _model_response_to_generate_content_response(model_response)
|
||
|
||
assert response.grounding_metadata is not None
|
||
assert (
|
||
response.grounding_metadata.grounding_chunks[0].web.uri
|
||
== "https://example.com"
|
||
)
|
||
|
||
|
||
def test_model_response_to_generate_content_response_grounding_metadata_list():
|
||
"""LiteLLM may emit a list (per candidate); the first entry is used."""
|
||
model_response = ModelResponse(
|
||
model="gemini/gemini-2.5-flash",
|
||
choices=[{
|
||
"message": {"role": "assistant", "content": "Answer"},
|
||
"finish_reason": "stop",
|
||
}],
|
||
)
|
||
model_response.vertex_ai_grounding_metadata = [
|
||
{"grounding_chunks": [{"web": {"uri": "https://a.test", "title": "A"}}]},
|
||
{"grounding_chunks": [{"web": {"uri": "https://b.test", "title": "B"}}]},
|
||
]
|
||
|
||
response = _model_response_to_generate_content_response(model_response)
|
||
|
||
assert response.grounding_metadata is not None
|
||
assert (
|
||
response.grounding_metadata.grounding_chunks[0].web.uri
|
||
== "https://a.test"
|
||
)
|
||
|
||
|
||
def test_model_response_to_generate_content_response_no_grounding_metadata():
|
||
"""Without vertex_ai_grounding_metadata, grounding_metadata stays None."""
|
||
model_response = ModelResponse(
|
||
model="gemini/gemini-2.5-flash",
|
||
choices=[{
|
||
"message": {"role": "assistant", "content": "Answer"},
|
||
"finish_reason": "stop",
|
||
}],
|
||
)
|
||
|
||
response = _model_response_to_generate_content_response(model_response)
|
||
|
||
assert response.grounding_metadata is None
|
||
|
||
|
||
def test_reasoning_content_takes_precedence_over_reasoning():
|
||
"""Test that 'reasoning_content' is prioritized over 'reasoning'."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"reasoning_content": "LiteLLM standard reasoning",
|
||
"reasoning": "Alternative reasoning",
|
||
}
|
||
response = _message_to_generate_content_response(message)
|
||
|
||
assert len(response.content.parts) == 2
|
||
thought_part = response.content.parts[0]
|
||
assert thought_part.text == "LiteLLM standard reasoning"
|
||
assert thought_part.thought is True
|
||
|
||
|
||
def test_extract_reasoning_value_from_reasoning_content():
|
||
"""Test extraction from reasoning_content (LiteLLM standard)."""
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Answer",
|
||
reasoning_content="LiteLLM reasoning",
|
||
)
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "LiteLLM reasoning"
|
||
|
||
|
||
def test_extract_reasoning_value_from_reasoning():
|
||
"""Test extraction from reasoning (LM Studio, vLLM)."""
|
||
|
||
class MockMessage:
|
||
|
||
def __init__(self):
|
||
self.role = "assistant"
|
||
self.content = "Answer"
|
||
self.reasoning = "Alternative reasoning"
|
||
|
||
def get(self, key, default=None):
|
||
return getattr(self, key, default)
|
||
|
||
message = MockMessage()
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "Alternative reasoning"
|
||
|
||
|
||
def test_extract_reasoning_value_dict_reasoning_content():
|
||
"""Test extraction from dict with reasoning_content field."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"reasoning_content": "Dict reasoning content",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "Dict reasoning content"
|
||
|
||
|
||
def test_extract_reasoning_value_dict_reasoning():
|
||
"""Test extraction from dict with reasoning field."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"reasoning": "Dict reasoning",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "Dict reasoning"
|
||
|
||
|
||
def test_extract_reasoning_value_dict_prefers_reasoning_content():
|
||
"""Test that reasoning_content takes precedence over reasoning in dicts."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"reasoning_content": "Primary",
|
||
"reasoning": "Secondary",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "Primary"
|
||
|
||
|
||
def test_extract_reasoning_value_none_message():
|
||
"""Test that None message returns None."""
|
||
result = _extract_reasoning_value(None)
|
||
assert result is None
|
||
|
||
|
||
def test_extract_reasoning_value_no_reasoning_fields():
|
||
"""Test that None is returned when no reasoning fields exist."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer only",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result is None
|
||
|
||
|
||
def test_extract_thought_signature_from_extra_content():
|
||
"""Extracts thought_signature from extra_content (OpenAI-compatible path)."""
|
||
sig_b64 = base64.b64encode(b"test_signature").decode("utf-8")
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_123",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
extra_content={"google": {"thought_signature": sig_b64}},
|
||
)
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result == b"test_signature"
|
||
|
||
|
||
def test_extract_thought_signature_from_provider_specific_fields():
|
||
"""Extracts thought_signature from provider_specific_fields (Vertex path)."""
|
||
sig_b64 = base64.b64encode(b"vertex_sig").decode("utf-8")
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_456",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
provider_specific_fields={"thought_signature": sig_b64},
|
||
)
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result == b"vertex_sig"
|
||
|
||
|
||
def test_extract_thought_signature_from_function_provider_fields():
|
||
"""Extracts thought_signature from function's provider_specific_fields.
|
||
|
||
When provider_specific_fields is set directly on the function object
|
||
(e.g. by litellm internals), the extraction should find it.
|
||
"""
|
||
sig_b64 = base64.b64encode(b"func_sig").decode("utf-8")
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_func",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
)
|
||
# Simulate litellm setting provider_specific_fields on the function
|
||
tc.function.provider_specific_fields = {
|
||
"thought_signature": sig_b64,
|
||
}
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result == b"func_sig"
|
||
|
||
|
||
def test_extract_thought_signature_from_id():
|
||
"""Extracts thought_signature from tool call ID (__thought__ separator)."""
|
||
sig_b64 = base64.b64encode(b"id_sig").decode("utf-8")
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id=f"call_789{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
)
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result == b"id_sig"
|
||
|
||
|
||
def test_extract_thought_signature_returns_none_when_absent():
|
||
"""Returns None when no thought_signature is present."""
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_plain",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
)
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result is None
|
||
|
||
|
||
def test_extract_thought_signature_corrupted_base64_returns_none():
|
||
"""Returns None gracefully for corrupted base64 signatures."""
|
||
tc = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_bad",
|
||
function=Function(name="test_fn", arguments="{}"),
|
||
extra_content={"google": {"thought_signature": "!!!not_valid_base64!!!"}},
|
||
)
|
||
result = _extract_thought_signature_from_tool_call(tc)
|
||
assert result is None
|
||
|
||
|
||
def test_message_to_generate_content_response_preserves_thought_signature():
|
||
"""thought_signature from tool call is preserved on the output Part."""
|
||
sig_b64 = base64.b64encode(b"round_trip_sig").decode("utf-8")
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_ts_1",
|
||
function=Function(
|
||
name="load_skill",
|
||
arguments='{"skill": "my_skill"}',
|
||
),
|
||
extra_content={"google": {"thought_signature": sig_b64}},
|
||
)
|
||
],
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
fc_part = response.content.parts[0]
|
||
assert fc_part.function_call.name == "load_skill"
|
||
assert fc_part.function_call.id == "call_ts_1"
|
||
assert fc_part.thought_signature == b"round_trip_sig"
|
||
|
||
|
||
def test_message_to_generate_content_response_no_thought_signature():
|
||
"""Parts without thought_signature have thought_signature=None."""
|
||
message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_no_ts",
|
||
function=Function(
|
||
name="plain_tool",
|
||
arguments="{}",
|
||
),
|
||
)
|
||
],
|
||
)
|
||
|
||
response = _message_to_generate_content_response(message)
|
||
fc_part = response.content.parts[0]
|
||
assert fc_part.function_call.name == "plain_tool"
|
||
assert fc_part.thought_signature is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_preserves_thought_signature():
|
||
"""thought_signature on Part is emitted on both tool call metadata paths."""
|
||
sig_bytes = b"preserved_sig"
|
||
sig_b64 = base64.b64encode(sig_bytes).decode("utf-8")
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(
|
||
function_call=types.FunctionCall(
|
||
name="load_skill",
|
||
args={"skill": "my_skill"},
|
||
id="call_rt",
|
||
),
|
||
thought_signature=sig_bytes,
|
||
),
|
||
],
|
||
)
|
||
|
||
message = await _content_to_message_param(content)
|
||
assert message["role"] == "assistant"
|
||
tc = message["tool_calls"][0]
|
||
assert tc["function"]["name"] == "load_skill"
|
||
assert tc["id"] == "call_rt"
|
||
assert tc["provider_specific_fields"] == {"thought_signature": sig_b64}
|
||
assert tc["extra_content"] == {"google": {"thought_signature": sig_b64}}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_no_thought_signature():
|
||
"""Tool calls without thought_signature have no signature metadata."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part.from_function_call(name="plain_tool", args={"key": "val"}),
|
||
],
|
||
)
|
||
content.parts[0].function_call.id = "call_plain"
|
||
|
||
message = await _content_to_message_param(content)
|
||
tc = message["tool_calls"][0]
|
||
assert tc["id"] == "call_plain"
|
||
assert "provider_specific_fields" not in tc
|
||
assert "extra_content" not in tc
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_thought_signature_round_trip():
|
||
"""thought_signature survives a full round trip through ADK conversions.
|
||
|
||
Simulates the flow: litellm response → types.Part → litellm request.
|
||
"""
|
||
sig_b64 = base64.b64encode(b"full_round_trip").decode("utf-8")
|
||
|
||
# Step 1: Incoming litellm message with thought_signature
|
||
incoming_message = ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=None,
|
||
tool_calls=[
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="call_round",
|
||
function=Function(
|
||
name="load_skill",
|
||
arguments='{"skill_name": "test"}',
|
||
),
|
||
extra_content={"google": {"thought_signature": sig_b64}},
|
||
)
|
||
],
|
||
)
|
||
|
||
# Step 2: Convert to ADK internal format (types.Content)
|
||
llm_response = _message_to_generate_content_response(incoming_message)
|
||
fc_part = llm_response.content.parts[0]
|
||
assert fc_part.thought_signature == b"full_round_trip"
|
||
|
||
# Step 3: Convert back to litellm format
|
||
outgoing_message = await _content_to_message_param(llm_response.content)
|
||
out_tc = outgoing_message["tool_calls"][0]
|
||
assert out_tc["provider_specific_fields"] == {"thought_signature": sig_b64}
|
||
assert out_tc["extra_content"] == {"google": {"thought_signature": sig_b64}}
|
||
|
||
|
||
def test_parse_tool_calls_from_text_multiple_calls():
|
||
text = (
|
||
'{"name":"alpha","arguments":{"value":1}}\n'
|
||
"Some filler text "
|
||
'{"id":"custom","name":"beta","arguments":{"timezone":"Asia/Taipei"}} '
|
||
"ignored suffix"
|
||
)
|
||
tool_calls, remainder = _parse_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 2
|
||
assert tool_calls[0].function.name == "alpha"
|
||
assert json.loads(tool_calls[0].function.arguments) == {"value": 1}
|
||
assert tool_calls[1].id == "custom"
|
||
assert tool_calls[1].function.name == "beta"
|
||
assert json.loads(tool_calls[1].function.arguments) == {
|
||
"timezone": "Asia/Taipei"
|
||
}
|
||
assert remainder == "Some filler text ignored suffix"
|
||
|
||
|
||
def test_parse_tool_calls_from_text_invalid_json_returns_remainder():
|
||
text = 'Leading {"unused": "payload"} trailing text'
|
||
tool_calls, remainder = _parse_tool_calls_from_text(text)
|
||
assert tool_calls == []
|
||
assert remainder == 'Leading {"unused": "payload"} trailing text'
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DeepSeek proprietary inline tool-call format tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_DS_BEGIN_CALLS = "\u003c\uff5ctool\u2581calls\u2581begin\uff5c\u003e"
|
||
_DS_END_CALLS = "\u003c\uff5ctool\u2581calls\u2581end\uff5c\u003e"
|
||
_DS_BEGIN_CALL = "\u003c\uff5ctool\u2581call\u2581begin\uff5c\u003e"
|
||
_DS_END_CALL = "\u003c\uff5ctool\u2581call\u2581end\uff5c\u003e"
|
||
_DS_SEP = "\u003c\uff5ctool\u2581sep\uff5c\u003e"
|
||
|
||
|
||
def _ds_tool_call(name: str, args_json: str) -> str:
|
||
"""Build a single DeepSeek-style tool-call block."""
|
||
return (
|
||
f"{_DS_BEGIN_CALL}function{_DS_SEP}{name}\n"
|
||
f"```json\n{args_json}\n```"
|
||
f"{_DS_END_CALL}"
|
||
)
|
||
|
||
|
||
def _ds_wrapped(inner: str) -> str:
|
||
"""Wrap content in <｜tool▁calls▁begin｜>...<｜tool▁calls▁end｜>."""
|
||
return f"{_DS_BEGIN_CALLS}{inner}{_DS_END_CALLS}"
|
||
|
||
|
||
def test_parse_deepseek_single_tool_call():
|
||
"""Single DeepSeek tool call with code-fenced JSON args."""
|
||
text = _ds_wrapped(
|
||
_ds_tool_call("get_weather", '{"city": "Beijing", "unit": "celsius"}')
|
||
)
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 1
|
||
assert tool_calls[0].function.name == "get_weather"
|
||
assert json.loads(tool_calls[0].function.arguments) == {
|
||
"city": "Beijing",
|
||
"unit": "celsius",
|
||
}
|
||
assert remainder is None
|
||
|
||
|
||
def test_parse_deepseek_multi_tool_call():
|
||
"""Multiple DeepSeek tool calls in a single wrapped block."""
|
||
inner = _ds_tool_call("func_a", '{"x": 1}') + _ds_tool_call(
|
||
"func_b", '{"y": 2}'
|
||
)
|
||
text = _ds_wrapped(inner)
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 2
|
||
assert tool_calls[0].function.name == "func_a"
|
||
assert json.loads(tool_calls[0].function.arguments) == {"x": 1}
|
||
assert tool_calls[1].function.name == "func_b"
|
||
assert json.loads(tool_calls[1].function.arguments) == {"y": 2}
|
||
assert remainder is None
|
||
|
||
|
||
def test_parse_deepseek_plain_json_args():
|
||
"""DeepSeek tool call without Markdown code fences around args."""
|
||
inner = (
|
||
f"{_DS_BEGIN_CALL}function{_DS_SEP}search\n"
|
||
f'{{"query": "天气"}}'
|
||
f"{_DS_END_CALL}"
|
||
)
|
||
text = _ds_wrapped(inner)
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 1
|
||
assert tool_calls[0].function.name == "search"
|
||
assert json.loads(tool_calls[0].function.arguments) == {"query": "天气"}
|
||
|
||
|
||
def test_parse_deepseek_with_surrounding_text():
|
||
"""DeepSeek tool call embedded in surrounding non-tool text."""
|
||
prefix = "Let me think about this.\n"
|
||
suffix = "\nI'll proceed now."
|
||
inner = _ds_tool_call("calculate", '{"expr": "2+2"}')
|
||
text = prefix + _ds_wrapped(inner) + suffix
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 1
|
||
assert tool_calls[0].function.name == "calculate"
|
||
assert remainder == "Let me think about this.\n\nI'll proceed now."
|
||
|
||
|
||
def test_parse_deepseek_no_tokens_returns_empty():
|
||
"""Text without DeepSeek tokens returns no tool calls and None remainder."""
|
||
text = "Just a regular response, no special tokens here."
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert tool_calls == []
|
||
assert remainder is None
|
||
|
||
|
||
def test_parse_tool_calls_from_text_handles_deepseek_format():
|
||
"""Integration: the generic parser delegates to the DeepSeek parser."""
|
||
text = _ds_wrapped(
|
||
_ds_tool_call("fetch_page", '{"url": "https://example.com"}')
|
||
)
|
||
tool_calls, remainder = _parse_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 1
|
||
assert tool_calls[0].function.name == "fetch_page"
|
||
assert json.loads(tool_calls[0].function.arguments) == {
|
||
"url": "https://example.com"
|
||
}
|
||
assert remainder is None
|
||
|
||
|
||
def test_parse_tool_calls_from_text_mixed_formats():
|
||
"""DeepSeek tokens + standard inline JSON in the same text."""
|
||
ds_part = _ds_wrapped(_ds_tool_call("ds_func", '{"a": 1}'))
|
||
standard_part = '{"name": "std_func", "arguments": {"b": 2}}'
|
||
text = ds_part + " some text " + standard_part
|
||
tool_calls, remainder = _parse_tool_calls_from_text(text)
|
||
assert len(tool_calls) == 2
|
||
assert tool_calls[0].function.name == "ds_func"
|
||
assert tool_calls[1].function.name == "std_func"
|
||
assert remainder == "some text"
|
||
|
||
|
||
def test_parse_deepseek_empty_text():
|
||
"""Empty or whitespace-only text returns no tool calls."""
|
||
for text in ("", " ", "\n\n"):
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
|
||
assert tool_calls == []
|
||
assert remainder is None
|
||
|
||
|
||
def test_parse_deepseek_unwrapped_call_before_wrapped_block():
|
||
"""Unwrapped call preceding a wrapped block is not dropped."""
|
||
unwrapped = _ds_tool_call("first", '{"x": 1}')
|
||
wrapped = _ds_wrapped(_ds_tool_call("second", '{"y": 2}'))
|
||
tool_calls, remainder = _parse_deepseek_tool_calls_from_text(
|
||
unwrapped + wrapped
|
||
)
|
||
assert [tc.function.name for tc in tool_calls] == ["first", "second"]
|
||
assert remainder is None
|
||
|
||
|
||
def test_extract_json_from_deepseek_args_invalid_fence_returns_none():
|
||
"""Invalid JSON inside a code fence is rejected rather than returned."""
|
||
assert _extract_json_from_deepseek_args('```json\n{"a": 1,}\n```') is None
|
||
|
||
|
||
def test_split_message_content_and_tool_calls_inline_text():
|
||
message = {
|
||
"role": "assistant",
|
||
"content": (
|
||
'Intro {"name":"alpha","arguments":{"value":1}} trailing content'
|
||
),
|
||
}
|
||
content, tool_calls = _split_message_content_and_tool_calls(message)
|
||
assert content == "Intro trailing content"
|
||
assert len(tool_calls) == 1
|
||
assert tool_calls[0].function.name == "alpha"
|
||
assert json.loads(tool_calls[0].function.arguments) == {"value": 1}
|
||
|
||
|
||
def test_split_message_content_prefers_existing_structured_calls():
|
||
tool_call = ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="existing",
|
||
function=Function(
|
||
name="existing_call",
|
||
arguments='{"arg": "value"}',
|
||
),
|
||
)
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "ignored",
|
||
"tool_calls": [tool_call],
|
||
}
|
||
content, tool_calls = _split_message_content_and_tool_calls(message)
|
||
assert content == "ignored"
|
||
assert tool_calls == [tool_call]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_does_not_filter_thought_parts():
|
||
"""Test that _get_content does not drop thought parts.
|
||
|
||
Thought filtering is handled by the caller (e.g., _content_to_message_param)
|
||
to avoid duplicating logic across helpers.
|
||
"""
|
||
thought_part = types.Part(text="Internal reasoning...", thought=True)
|
||
regular_part = types.Part.from_text(text="Visible response")
|
||
|
||
content = await _get_content([thought_part, regular_part])
|
||
|
||
assert content == [
|
||
{"type": "text", "text": "Internal reasoning..."},
|
||
{"type": "text", "text": "Visible response"},
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_all_thought_parts():
|
||
"""Test that thought parts convert like regular text parts."""
|
||
thought_part1 = types.Part(text="First reasoning...", thought=True)
|
||
thought_part2 = types.Part(text="Second reasoning...", thought=True)
|
||
|
||
content = await _get_content([thought_part1, thought_part2])
|
||
|
||
assert content == [
|
||
{"type": "text", "text": "First reasoning..."},
|
||
{"type": "text", "text": "Second reasoning..."},
|
||
]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_text():
|
||
parts = [types.Part.from_text(text="Test text")]
|
||
content = await _get_content(parts)
|
||
assert content == "Test text"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_text_inline_data_single_part():
|
||
parts = [
|
||
types.Part.from_bytes(
|
||
data="Inline text".encode("utf-8"), mime_type="text/plain"
|
||
)
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content == "Inline text"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_text_inline_data_multiple_parts():
|
||
parts = [
|
||
types.Part.from_bytes(
|
||
data="First part".encode("utf-8"), mime_type="text/plain"
|
||
),
|
||
types.Part.from_text(text="Second part"),
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0]["type"] == "text"
|
||
assert content[0]["text"] == "First part"
|
||
assert content[1]["type"] == "text"
|
||
assert content[1]["text"] == "Second part"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_text_inline_data_fallback_decoding():
|
||
parts = [
|
||
types.Part.from_bytes(data=b"\xff", mime_type="text/plain"),
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content == "ÿ"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_image():
|
||
parts = [
|
||
types.Part.from_bytes(data=b"test_image_data", mime_type="image/png")
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0]["type"] == "image_url"
|
||
assert (
|
||
content[0]["image_url"]["url"]
|
||
== "data:image/png;base64,dGVzdF9pbWFnZV9kYXRh"
|
||
)
|
||
assert "format" not in content[0]["image_url"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_video():
|
||
parts = [
|
||
types.Part.from_bytes(data=b"test_video_data", mime_type="video/mp4")
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0]["type"] == "video_url"
|
||
assert (
|
||
content[0]["video_url"]["url"]
|
||
== "data:video/mp4;base64,dGVzdF92aWRlb19kYXRh"
|
||
)
|
||
assert "format" not in content[0]["video_url"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"file_data,mime_type,expected_base64", FILE_BYTES_TEST_CASES
|
||
)
|
||
async def test_get_content_file_bytes(file_data, mime_type, expected_base64):
|
||
parts = [types.Part.from_bytes(data=file_data, mime_type=mime_type)]
|
||
content = await _get_content(parts)
|
||
assert content[0]["type"] == "file"
|
||
assert content[0]["file"]["file_data"] == expected_base64
|
||
assert "format" not in content[0]["file"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
|
||
async def test_get_content_file_uri(file_uri, mime_type):
|
||
parts = [types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)]
|
||
content = await _get_content(parts)
|
||
assert content[0] == {
|
||
"type": "file",
|
||
"file": {"file_id": file_uri, "format": mime_type},
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"provider,model",
|
||
[
|
||
("openai", "openai/gpt-4o"),
|
||
("azure", "azure/gpt-4"),
|
||
],
|
||
)
|
||
async def test_get_content_file_uri_file_id_required_raises_error(
|
||
provider, model
|
||
):
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf",
|
||
mime_type="application/pdf",
|
||
display_name="document.pdf",
|
||
)
|
||
)
|
||
]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match=f"File URI `document.pdf` not supported for provider: {provider}",
|
||
):
|
||
_ = await _get_content(parts, provider=provider, model=model)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"provider,model",
|
||
[
|
||
("openai", "openai/gpt-4o"),
|
||
("azure", "azure/gpt-4"),
|
||
],
|
||
)
|
||
@pytest.mark.parametrize(
|
||
"file_uri,mime_type,expected_type",
|
||
[
|
||
pytest.param(
|
||
"https://example.com/image.png",
|
||
"image/png",
|
||
"image_url",
|
||
id="image",
|
||
),
|
||
pytest.param(
|
||
"https://example.com/video.mp4",
|
||
"video/mp4",
|
||
"video_url",
|
||
id="video",
|
||
),
|
||
],
|
||
)
|
||
async def test_get_content_file_uri_media_url_file_id_required_uses_url_type(
|
||
provider, model, file_uri, mime_type, expected_type
|
||
):
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri=file_uri,
|
||
mime_type=mime_type,
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(parts, provider=provider, model=model)
|
||
assert content == [{
|
||
"type": expected_type,
|
||
expected_type: {"url": file_uri},
|
||
}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"provider,model",
|
||
[
|
||
("openai", "openai/gpt-4o"),
|
||
("azure", "azure/gpt-4"),
|
||
],
|
||
)
|
||
async def test_get_content_file_uri_file_id_required_preserves_file_id(
|
||
provider, model
|
||
):
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="file-abc123",
|
||
mime_type="application/pdf",
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(parts, provider=provider, model=model)
|
||
assert content == [{"type": "file", "file": {"file_id": "file-abc123"}}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_azure_preserves_assistant_file_id():
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="assistant-abc123",
|
||
mime_type="application/pdf",
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(parts, provider="azure", model="azure/gpt-4.1")
|
||
assert content == [{"type": "file", "file": {"file_id": "assistant-abc123"}}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"provider,model",
|
||
[
|
||
("openai", "openai/gpt-4o"),
|
||
("azure", "azure/gpt-4"),
|
||
],
|
||
)
|
||
async def test_get_content_file_uri_http_pdf_file_id_required_raises_error(
|
||
provider, model
|
||
):
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="https://example.com/document.pdf",
|
||
mime_type="application/pdf",
|
||
display_name="document.pdf",
|
||
)
|
||
)
|
||
]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match=f"File URI `document.pdf` not supported for provider: {provider}",
|
||
):
|
||
_ = await _get_content(parts, provider=provider, model=model)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_http_pdf_non_file_id_provider_uses_file():
|
||
file_uri = "https://example.com/document.pdf"
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri=file_uri,
|
||
mime_type="application/pdf",
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(
|
||
parts, provider="vertex_ai", model="vertex_ai/gemini-2.5-flash"
|
||
)
|
||
assert content == [{
|
||
"type": "file",
|
||
"file": {"file_id": file_uri, "format": "application/pdf"},
|
||
}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_anthropic_raises_error():
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf",
|
||
mime_type="application/pdf",
|
||
display_name="document.pdf",
|
||
)
|
||
)
|
||
]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match="File URI `document.pdf` not supported for provider: anthropic",
|
||
):
|
||
_ = await _get_content(
|
||
parts, provider="anthropic", model="anthropic/claude-3-5"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_anthropic_openai_file_id_raises_error():
|
||
parts = [types.Part(file_data=types.FileData(file_uri="file-abc123"))]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match="File URI `file-<redacted>` not supported for provider: anthropic",
|
||
):
|
||
_ = await _get_content(
|
||
parts, provider="anthropic", model="anthropic/claude-3-5"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_vertex_ai_non_gemini_raises_error():
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf",
|
||
mime_type="application/pdf",
|
||
display_name="document.pdf",
|
||
)
|
||
)
|
||
]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match="File URI `document.pdf` not supported for provider: vertex_ai",
|
||
):
|
||
_ = await _get_content(
|
||
parts, provider="vertex_ai", model="vertex_ai/claude-3-5"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_vertex_ai_gemini_keeps_file_block():
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf",
|
||
mime_type="application/pdf",
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(
|
||
parts, provider="vertex_ai", model="vertex_ai/gemini-2.5-flash"
|
||
)
|
||
assert content == [{
|
||
"type": "file",
|
||
"file": {
|
||
"file_id": "gs://bucket/path/to/document.pdf",
|
||
"format": "application/pdf",
|
||
},
|
||
}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_infer_mime_type():
|
||
"""Test MIME type inference from file_uri extension.
|
||
|
||
When file_data has a file_uri with a recognizable extension but no explicit
|
||
mime_type, the MIME type should be inferred from the extension.
|
||
"""
|
||
# Use Part constructor directly to test MIME type inference in _get_content
|
||
# (types.Part.from_uri does its own inference, so we bypass it)
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(file_uri="gs://bucket/path/to/document.pdf")
|
||
)
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0] == {
|
||
"type": "file",
|
||
"file": {
|
||
"file_id": "gs://bucket/path/to/document.pdf",
|
||
"format": "application/pdf",
|
||
},
|
||
}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_versioned_infer_mime_type():
|
||
"""Test MIME type inference from versioned artifact URIs."""
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/path/to/document.pdf/0"
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0]["file"]["format"] == "application/pdf"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_infers_from_display_name():
|
||
"""Test MIME type inference from display_name when URI lacks extension."""
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/artifact/0",
|
||
display_name="document.pdf",
|
||
)
|
||
)
|
||
]
|
||
content = await _get_content(parts)
|
||
assert content[0]["file"]["format"] == "application/pdf"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_default_mime_type():
|
||
"""Test that file_uri without an inferable extension raises ValueError.
|
||
|
||
When file_data has a file_uri without a recognizable extension and no explicit
|
||
mime_type, ADK raises a clear ValueError instead of forwarding the unsupported
|
||
'application/octet-stream' MIME type to LiteLLM.
|
||
"""
|
||
parts = [
|
||
types.Part(file_data=types.FileData(file_uri="gs://bucket/artifact/0"))
|
||
]
|
||
with pytest.raises(ValueError, match="Cannot process file_uri"):
|
||
await _get_content(parts)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_file_uri_explicit_octet_stream_raises():
|
||
"""Test that an explicit application/octet-stream MIME type raises ValueError.
|
||
|
||
'application/octet-stream' is semantically equivalent to an unknown type and
|
||
causes the same downstream ValueError from LiteLLM whether it arrives as a
|
||
default fallback or is set explicitly by the caller. ADK raises early with
|
||
an actionable message in both cases.
|
||
"""
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(
|
||
file_uri="gs://bucket/artifact/0",
|
||
mime_type="application/octet-stream",
|
||
)
|
||
)
|
||
]
|
||
with pytest.raises(ValueError, match="application/octet-stream"):
|
||
await _get_content(parts)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"uri,expected_mime_type",
|
||
[
|
||
("gs://bucket/file.pdf", "application/pdf"),
|
||
("gs://bucket/path/to/document.json", "application/json"),
|
||
("gs://bucket/image.png", "image/png"),
|
||
("gs://bucket/image.jpg", "image/jpeg"),
|
||
("gs://bucket/audio.mp3", "audio/mpeg"),
|
||
("gs://bucket/video.mp4", "video/mp4"),
|
||
],
|
||
)
|
||
async def test_get_content_file_uri_mime_type_inference(
|
||
uri, expected_mime_type
|
||
):
|
||
"""Test MIME type inference from various file extensions."""
|
||
# Use Part constructor directly to test MIME type inference in _get_content
|
||
parts = [types.Part(file_data=types.FileData(file_uri=uri))]
|
||
content = await _get_content(parts)
|
||
assert content[0]["file"]["format"] == expected_mime_type
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"mime_type,expected_format",
|
||
[
|
||
("audio/mpeg", "mp3"),
|
||
("audio/mp3", "mp3"),
|
||
("audio/wav", "wav"),
|
||
("audio/x-wav", "wav"),
|
||
("audio/wave", "wav"),
|
||
("audio/flac", "flac"),
|
||
("audio/ogg", "ogg"),
|
||
("audio/mp4", "mp4"),
|
||
],
|
||
)
|
||
async def test_get_content_audio_inline_data_emits_input_audio(
|
||
mime_type, expected_format
|
||
):
|
||
"""Audio inline_data is serialised as `input_audio` with raw base64 + format."""
|
||
parts = [types.Part.from_bytes(data=b"test_audio_data", mime_type=mime_type)]
|
||
content = await _get_content(parts)
|
||
assert content == [{
|
||
"type": "input_audio",
|
||
"input_audio": {
|
||
"data": "dGVzdF9hdWRpb19kYXRh",
|
||
"format": expected_format,
|
||
},
|
||
}]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"provider,model",
|
||
[
|
||
("openai", "openai/gpt-4o"),
|
||
("azure", "azure/gpt-4"),
|
||
],
|
||
)
|
||
async def test_get_content_audio_file_uri_http_raises_error(provider, model):
|
||
"""Audio HTTP file_uri raises an error for openai/azure."""
|
||
file_uri = "https://example.com/audio.mp3"
|
||
parts = [
|
||
types.Part(
|
||
file_data=types.FileData(file_uri=file_uri, mime_type="audio/mpeg")
|
||
)
|
||
]
|
||
with pytest.raises(
|
||
ValueError,
|
||
match=(
|
||
"File URI `https://<redacted>/audio.mp3` not supported for provider:"
|
||
f" {provider}"
|
||
),
|
||
):
|
||
_ = await _get_content(parts, provider=provider, model=model)
|
||
|
||
|
||
def test_to_litellm_role():
|
||
assert _to_litellm_role("model") == "assistant"
|
||
assert _to_litellm_role("assistant") == "assistant"
|
||
assert _to_litellm_role("user") == "user"
|
||
assert _to_litellm_role(None) == "user"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"response, expected_chunks, expected_usage_chunk, expected_finished",
|
||
[
|
||
(
|
||
ModelResponse(
|
||
choices=[
|
||
{
|
||
"message": {
|
||
"content": "this is a test",
|
||
}
|
||
}
|
||
],
|
||
usage={
|
||
"prompt_tokens": 0,
|
||
"completion_tokens": 0,
|
||
"total_tokens": 0,
|
||
},
|
||
),
|
||
[TextChunk(text="this is a test")],
|
||
UsageMetadataChunk(
|
||
prompt_tokens=0, completion_tokens=0, total_tokens=0
|
||
),
|
||
"stop",
|
||
),
|
||
(
|
||
ModelResponse(
|
||
choices=[
|
||
{
|
||
"message": {
|
||
"content": "this is a test",
|
||
}
|
||
}
|
||
],
|
||
usage={
|
||
"prompt_tokens": 3,
|
||
"completion_tokens": 5,
|
||
"total_tokens": 8,
|
||
},
|
||
),
|
||
[TextChunk(text="this is a test")],
|
||
UsageMetadataChunk(
|
||
prompt_tokens=3, completion_tokens=5, total_tokens=8
|
||
),
|
||
"stop",
|
||
),
|
||
(
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="1",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"key": "va',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
[FunctionChunk(id="1", name="test_function", args='{"key": "va')],
|
||
None,
|
||
# LiteLLM 1.81+ defaults finish_reason to "stop" for partial chunks,
|
||
# older versions return None. Both are valid for streaming chunks.
|
||
(None, "stop"),
|
||
),
|
||
(
|
||
ModelResponse(choices=[{"finish_reason": "tool_calls"}]),
|
||
[None],
|
||
(
|
||
None,
|
||
UsageMetadataChunk(
|
||
prompt_tokens=0, completion_tokens=0, total_tokens=0
|
||
),
|
||
),
|
||
"tool_calls",
|
||
),
|
||
(
|
||
ModelResponse(choices=[{}]),
|
||
[None],
|
||
(
|
||
None,
|
||
UsageMetadataChunk(
|
||
prompt_tokens=0, completion_tokens=0, total_tokens=0
|
||
),
|
||
),
|
||
"stop",
|
||
),
|
||
(
|
||
ModelResponse(
|
||
choices=[{
|
||
"finish_reason": "tool_calls",
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": (
|
||
'{"id":"call_1","name":"get_current_time",'
|
||
'"arguments":{"timezone_str":"Asia/Taipei"}}'
|
||
),
|
||
},
|
||
}],
|
||
usage={
|
||
"prompt_tokens": 7,
|
||
"completion_tokens": 9,
|
||
"total_tokens": 16,
|
||
},
|
||
),
|
||
[
|
||
FunctionChunk(
|
||
id="call_1",
|
||
name="get_current_time",
|
||
args='{"timezone_str": "Asia/Taipei"}',
|
||
index=0,
|
||
),
|
||
],
|
||
UsageMetadataChunk(
|
||
prompt_tokens=7, completion_tokens=9, total_tokens=16
|
||
),
|
||
"tool_calls",
|
||
),
|
||
(
|
||
ModelResponse(
|
||
choices=[{
|
||
"finish_reason": "tool_calls",
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": (
|
||
'Intro {"id":"call_2","name":"alpha",'
|
||
'"arguments":{"foo":"bar"}} wrap'
|
||
),
|
||
},
|
||
}],
|
||
usage={
|
||
"prompt_tokens": 11,
|
||
"completion_tokens": 13,
|
||
"total_tokens": 24,
|
||
},
|
||
),
|
||
[
|
||
TextChunk(text="Intro wrap"),
|
||
FunctionChunk(
|
||
id="call_2",
|
||
name="alpha",
|
||
args='{"foo": "bar"}',
|
||
index=0,
|
||
),
|
||
],
|
||
UsageMetadataChunk(
|
||
prompt_tokens=11, completion_tokens=13, total_tokens=24
|
||
),
|
||
"tool_calls",
|
||
),
|
||
(
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(role="assistant", content="Hello"),
|
||
)
|
||
],
|
||
usage=None,
|
||
),
|
||
[TextChunk(text="Hello")],
|
||
None,
|
||
(None, "stop"),
|
||
),
|
||
(
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason="stop",
|
||
delta=Delta(
|
||
role="assistant", reasoning_content="thinking..."
|
||
),
|
||
)
|
||
],
|
||
usage=None,
|
||
),
|
||
[
|
||
ReasoningChunk(
|
||
parts=[types.Part(text="thinking...", thought=True)]
|
||
)
|
||
],
|
||
None,
|
||
"stop",
|
||
),
|
||
],
|
||
)
|
||
def test_model_response_to_chunk(
|
||
response, expected_chunks, expected_usage_chunk, expected_finished
|
||
):
|
||
result = list(_model_response_to_chunk(response))
|
||
observed_chunks = []
|
||
usage_chunk = None
|
||
for chunk, finished in result:
|
||
if isinstance(chunk, UsageMetadataChunk):
|
||
usage_chunk = chunk
|
||
continue
|
||
observed_chunks.append((chunk, finished))
|
||
|
||
assert len(observed_chunks) == len(expected_chunks)
|
||
for (chunk, finished), expected_chunk in zip(
|
||
observed_chunks, expected_chunks
|
||
):
|
||
if expected_chunk is None:
|
||
assert chunk is None
|
||
else:
|
||
assert isinstance(chunk, type(expected_chunk))
|
||
assert chunk == expected_chunk
|
||
if isinstance(expected_finished, tuple):
|
||
assert finished in expected_finished
|
||
else:
|
||
assert finished == expected_finished
|
||
|
||
if isinstance(expected_usage_chunk, tuple):
|
||
assert usage_chunk in expected_usage_chunk
|
||
elif expected_usage_chunk is None:
|
||
assert usage_chunk is None
|
||
else:
|
||
assert usage_chunk is not None
|
||
assert usage_chunk == expected_usage_chunk
|
||
|
||
|
||
def test_model_response_to_chunk_does_not_mutate_delta_object():
|
||
"""Verify that _model_response_to_chunk doesn't mutate the Delta object.
|
||
|
||
In real streaming responses, LiteLLM's StreamingChoices only has 'delta'
|
||
(message is explicitly popped in StreamingChoices constructor). The delta
|
||
object itself carries reasoning_content when present.
|
||
"""
|
||
delta = Delta(
|
||
role="assistant", content="Hello", reasoning_content="thinking..."
|
||
)
|
||
response = ModelResponseStream(
|
||
choices=[StreamingChoices(delta=delta, finish_reason=None)]
|
||
)
|
||
|
||
chunks = [chunk for chunk, _ in _model_response_to_chunk(response) if chunk]
|
||
|
||
assert (
|
||
ReasoningChunk(parts=[types.Part(text="thinking...", thought=True)])
|
||
in chunks
|
||
)
|
||
assert TextChunk(text="Hello") in chunks
|
||
|
||
# Verify we don't accidentally mutate the original delta object.
|
||
assert delta.content == "Hello"
|
||
assert delta.reasoning_content == "thinking..."
|
||
|
||
|
||
def test_model_response_to_chunk_rejects_dict_response():
|
||
with pytest.raises(TypeError):
|
||
list(_model_response_to_chunk({"choices": []}))
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_additional_args(mock_acompletion, mock_client):
|
||
lite_llm_instance = LiteLlm(
|
||
# valid args
|
||
model="vertex_ai/test_model",
|
||
llm_client=mock_client,
|
||
api_key="test_key",
|
||
api_base="some://url",
|
||
api_version="2024-09-12",
|
||
headers={"custom": "header"}, # Add custom header to test merge
|
||
# invalid args (ignored)
|
||
stream=True,
|
||
messages=[{"role": "invalid", "content": "invalid"}],
|
||
tools=[{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "invalid",
|
||
},
|
||
}],
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION
|
||
):
|
||
assert response.content.role == "model"
|
||
assert response.content.parts[0].text == "Test response"
|
||
assert response.content.parts[1].function_call.name == "test_function"
|
||
assert response.content.parts[1].function_call.args == {
|
||
"test_arg": "test_value"
|
||
}
|
||
assert response.content.parts[1].function_call.id == "test_tool_call_id"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
|
||
assert kwargs["model"] == "vertex_ai/test_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
assert kwargs["tools"][0]["function"]["name"] == "test_function"
|
||
assert "stream" not in kwargs
|
||
assert "llm_client" not in kwargs
|
||
assert kwargs["api_base"] == "some://url"
|
||
assert "headers" in kwargs
|
||
assert kwargs["headers"]["custom"] == "header"
|
||
assert "x-goog-api-client" in kwargs["headers"]
|
||
assert "user-agent" in kwargs["headers"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_additional_args_non_vertex(
|
||
mock_acompletion, mock_client
|
||
):
|
||
"""Test that tracking headers are not added for non-Vertex AI models."""
|
||
lite_llm_instance = LiteLlm(
|
||
model="openai/gpt-4o",
|
||
llm_client=mock_client,
|
||
api_key="test_key",
|
||
headers={"custom": "header"},
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION
|
||
):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["model"] == "openai/gpt-4o"
|
||
assert "headers" in kwargs
|
||
assert kwargs["headers"]["custom"] == "header"
|
||
assert "x-goog-api-client" not in kwargs["headers"]
|
||
assert "user-agent" not in kwargs["headers"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_acompletion_with_drop_params(mock_acompletion, mock_client):
|
||
lite_llm_instance = LiteLlm(
|
||
model="test_model", llm_client=mock_client, drop_params=True
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION
|
||
):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert kwargs["drop_params"] is True
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_completion_additional_args(mock_completion, mock_client):
|
||
lite_llm_instance = LiteLlm(
|
||
# valid args
|
||
model="test_model",
|
||
llm_client=mock_client,
|
||
api_key="test_key",
|
||
api_base="some://url",
|
||
api_version="2024-09-12",
|
||
# invalid args (ignored)
|
||
stream=False,
|
||
messages=[{"role": "invalid", "content": "invalid"}],
|
||
tools=[{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "invalid",
|
||
},
|
||
}],
|
||
)
|
||
|
||
mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
assert len(responses) == 4
|
||
mock_completion.assert_called_once()
|
||
|
||
_, kwargs = mock_completion.call_args
|
||
|
||
assert kwargs["model"] == "test_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
assert kwargs["tools"][0]["function"]["name"] == "test_function"
|
||
assert kwargs["stream"]
|
||
assert "llm_client" not in kwargs
|
||
assert kwargs["api_base"] == "some://url"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_completion_with_drop_params(mock_completion, mock_client):
|
||
lite_llm_instance = LiteLlm(
|
||
model="test_model", llm_client=mock_client, drop_params=True
|
||
)
|
||
|
||
mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
assert len(responses) == 4
|
||
|
||
mock_completion.assert_called_once()
|
||
|
||
_, kwargs = mock_completion.call_args
|
||
assert kwargs["drop_params"] is True
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_grounding_metadata(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
final_chunk = ModelResponseStream(
|
||
model="test_model",
|
||
choices=[StreamingChoices(finish_reason="stop", delta=Delta())],
|
||
)
|
||
final_chunk.vertex_ai_grounding_metadata = {
|
||
"grounding_chunks": [
|
||
{"web": {"uri": "https://example.com", "title": "Example"}}
|
||
],
|
||
}
|
||
mock_completion.return_value = iter([
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(role="assistant", content="Grounded answer"),
|
||
)
|
||
],
|
||
),
|
||
final_chunk,
|
||
])
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
llm_request, stream=True
|
||
)
|
||
]
|
||
|
||
assert responses[-1].partial is False
|
||
assert responses[-1].grounding_metadata is not None
|
||
assert (
|
||
responses[-1].grounding_metadata.grounding_chunks[0].web.uri
|
||
== "https://example.com"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_with_usage_metadata(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
|
||
mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
assert len(responses) == 4
|
||
assert responses[0].content.role == "model"
|
||
assert responses[0].content.parts[0].text == "zero, "
|
||
assert responses[0].model_version == "test_model"
|
||
assert responses[1].content.role == "model"
|
||
assert responses[1].content.parts[0].text == "one, "
|
||
assert responses[1].model_version == "test_model"
|
||
assert responses[2].content.role == "model"
|
||
assert responses[2].content.parts[0].text == "two:"
|
||
assert responses[2].model_version == "test_model"
|
||
assert responses[3].content.role == "model"
|
||
assert responses[3].content.parts[-1].function_call.name == "test_function"
|
||
assert responses[3].content.parts[-1].function_call.args == {
|
||
"test_arg": "test_value"
|
||
}
|
||
assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id"
|
||
assert responses[3].finish_reason == types.FinishReason.STOP
|
||
assert responses[3].model_version == "test_model"
|
||
mock_completion.assert_called_once()
|
||
|
||
_, kwargs = mock_completion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
assert kwargs["tools"][0]["function"]["name"] == "test_function"
|
||
assert (
|
||
kwargs["tools"][0]["function"]["description"]
|
||
== "Test function description"
|
||
)
|
||
assert (
|
||
kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
|
||
"type"
|
||
]
|
||
== "string"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_sets_finish_reason(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
mock_completion.return_value = iter([
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(role="assistant", content="Hello "),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(role="assistant", content="world"),
|
||
)
|
||
],
|
||
),
|
||
ModelResponseStream(
|
||
model="test_model",
|
||
choices=[StreamingChoices(finish_reason="stop", delta=Delta())],
|
||
),
|
||
])
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
llm_request, stream=True
|
||
)
|
||
]
|
||
|
||
assert responses[-1].partial is False
|
||
assert responses[-1].finish_reason == types.FinishReason.STOP
|
||
assert responses[-1].content.parts[0].text == "Hello world"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_with_usage_metadata(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
|
||
streaming_model_response_with_usage_metadata = [
|
||
*STREAMING_MODEL_RESPONSE,
|
||
ModelResponseStream(
|
||
usage={
|
||
"prompt_tokens": 10,
|
||
"completion_tokens": 5,
|
||
"total_tokens": 15,
|
||
"completion_tokens_details": {"reasoning_tokens": 5},
|
||
},
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
)
|
||
],
|
||
),
|
||
]
|
||
|
||
mock_completion.return_value = iter(
|
||
streaming_model_response_with_usage_metadata
|
||
)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
assert len(responses) == 4
|
||
assert responses[0].content.role == "model"
|
||
assert responses[0].content.parts[0].text == "zero, "
|
||
assert responses[1].content.role == "model"
|
||
assert responses[1].content.parts[0].text == "one, "
|
||
assert responses[2].content.role == "model"
|
||
assert responses[2].content.parts[0].text == "two:"
|
||
assert responses[3].content.role == "model"
|
||
assert responses[3].content.parts[-1].function_call.name == "test_function"
|
||
assert responses[3].content.parts[-1].function_call.args == {
|
||
"test_arg": "test_value"
|
||
}
|
||
assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id"
|
||
assert responses[3].finish_reason == types.FinishReason.STOP
|
||
|
||
assert responses[3].usage_metadata.prompt_token_count == 10
|
||
assert responses[3].usage_metadata.candidates_token_count == 5
|
||
assert responses[3].usage_metadata.total_token_count == 15
|
||
assert responses[3].usage_metadata.thoughts_token_count == 5
|
||
|
||
mock_completion.assert_called_once()
|
||
|
||
_, kwargs = mock_completion.call_args
|
||
assert kwargs["model"] == "test_model"
|
||
assert kwargs["messages"][0]["role"] == "user"
|
||
assert kwargs["messages"][0]["content"] == "Test prompt"
|
||
assert kwargs["tools"][0]["function"]["name"] == "test_function"
|
||
assert (
|
||
kwargs["tools"][0]["function"]["description"]
|
||
== "Test function description"
|
||
)
|
||
assert (
|
||
kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
|
||
"type"
|
||
]
|
||
== "string"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_with_usage_metadata(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Tests that cached prompt tokens are propagated in streaming mode."""
|
||
streaming_model_response_with_usage_metadata = [
|
||
*STREAMING_MODEL_RESPONSE,
|
||
ModelResponseStream(
|
||
usage={
|
||
"prompt_tokens": 10,
|
||
"completion_tokens": 5,
|
||
"total_tokens": 15,
|
||
"cached_tokens": 8,
|
||
"completion_tokens_details": {"reasoning_tokens": 5},
|
||
},
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
)
|
||
],
|
||
),
|
||
]
|
||
|
||
mock_completion.return_value = iter(
|
||
streaming_model_response_with_usage_metadata
|
||
)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
assert len(responses) == 4
|
||
assert responses[3].usage_metadata.prompt_token_count == 10
|
||
assert responses[3].usage_metadata.candidates_token_count == 5
|
||
assert responses[3].usage_metadata.total_token_count == 15
|
||
assert responses[3].usage_metadata.cached_content_token_count == 8
|
||
assert responses[3].usage_metadata.thoughts_token_count == 5
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_multiple_function_calls(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Test handling of multiple function calls with different indices in streaming mode.
|
||
|
||
This test verifies that:
|
||
1. Multiple function calls with different indices are handled correctly
|
||
2. Arguments and names are properly accumulated for each function call
|
||
3. The final response contains all function calls with correct indices
|
||
"""
|
||
mock_completion.return_value = MULTIPLE_FUNCTION_CALLS_STREAM
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[types.Part.from_text(text="Test multiple function calls")],
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
tools=[
|
||
types.Tool(
|
||
function_declarations=[
|
||
types.FunctionDeclaration(
|
||
name="function_1",
|
||
description="First test function",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
types.FunctionDeclaration(
|
||
name="function_2",
|
||
description="Second test function",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
]
|
||
)
|
||
],
|
||
),
|
||
)
|
||
|
||
responses = []
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
llm_request, stream=True
|
||
):
|
||
responses.append(response)
|
||
|
||
# Verify we got the final response with both function calls
|
||
assert len(responses) > 0
|
||
final_response = responses[-1]
|
||
assert final_response.content.role == "model"
|
||
assert len(final_response.content.parts) == 2
|
||
|
||
# Verify first function call
|
||
assert final_response.content.parts[0].function_call.name == "function_1"
|
||
assert final_response.content.parts[0].function_call.id == "call_1"
|
||
assert final_response.content.parts[0].function_call.args == {"arg": "value1"}
|
||
|
||
# Verify second function call
|
||
assert final_response.content.parts[1].function_call.name == "function_2"
|
||
assert final_response.content.parts[1].function_call.id == "call_2"
|
||
assert final_response.content.parts[1].function_call.args == {"arg": "value2"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_non_compliant_multiple_function_calls(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Test handling of multiple function calls with same 0 indices in streaming mode.
|
||
|
||
This test verifies that:
|
||
1. Multiple function calls with same indices (0) are handled correctly
|
||
2. Arguments and names are properly accumulated for each function call
|
||
3. The final response contains all function calls with correct incremented
|
||
indices
|
||
"""
|
||
mock_completion.return_value = NON_COMPLIANT_MULTIPLE_FUNCTION_CALLS_STREAM
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[types.Part.from_text(text="Test multiple function calls")],
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
tools=[
|
||
types.Tool(
|
||
function_declarations=[
|
||
types.FunctionDeclaration(
|
||
name="function_1",
|
||
description="First test function",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
types.FunctionDeclaration(
|
||
name="function_2",
|
||
description="Second test function",
|
||
parameters=types.Schema(
|
||
type=types.Type.OBJECT,
|
||
properties={
|
||
"arg": types.Schema(type=types.Type.STRING),
|
||
},
|
||
),
|
||
),
|
||
]
|
||
)
|
||
],
|
||
),
|
||
)
|
||
|
||
responses = []
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
llm_request, stream=True
|
||
):
|
||
responses.append(response)
|
||
|
||
# Verify we got the final response with both function calls
|
||
assert len(responses) > 0
|
||
final_response = responses[-1]
|
||
assert final_response.content.role == "model"
|
||
assert len(final_response.content.parts) == 2
|
||
|
||
# Verify first function call
|
||
assert final_response.content.parts[0].function_call.name == "function_1"
|
||
assert final_response.content.parts[0].function_call.id == "0"
|
||
assert final_response.content.parts[0].function_call.args == {"arg": "value1"}
|
||
|
||
# Verify second function call
|
||
assert final_response.content.parts[1].function_call.name == "function_2"
|
||
assert final_response.content.parts[1].function_call.id == "1"
|
||
assert final_response.content.parts[1].function_call.args == {"arg": "value2"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_stream_with_empty_chunk(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Tests that empty tool call chunks in a stream are ignored."""
|
||
mock_completion.return_value = iter(STREAM_WITH_EMPTY_CHUNK)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
|
||
assert len(responses) == 1
|
||
final_response = responses[0]
|
||
assert final_response.content.role == "model"
|
||
|
||
# Crucially, assert that only ONE tool call was generated,
|
||
# proving the empty chunk was ignored.
|
||
assert len(final_response.content.parts) == 1
|
||
|
||
function_call = final_response.content.parts[0].function_call
|
||
assert function_call.name == "test_function"
|
||
assert function_call.id == "call_abc"
|
||
assert function_call.args == {"test_arg": "value"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_call_truncated_by_max_tokens(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Tests that truncated tool calls with finish_reason='length' yield an error LlmResponse."""
|
||
stream_chunks = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="call_123",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg":',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[StreamingChoices(finish_reason="length", delta=Delta())]
|
||
),
|
||
]
|
||
mock_completion.return_value = iter(stream_chunks)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
|
||
assert len(responses) == 1
|
||
error_response = responses[0]
|
||
assert error_response.error_code == types.FinishReason.MAX_TOKENS
|
||
assert error_response.finish_reason == types.FinishReason.MAX_TOKENS
|
||
assert "truncated" in error_response.error_message
|
||
assert "max_output_tokens" in error_response.error_message
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_call_complete_with_length_finish_reason(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Tests that complete tool calls with finish_reason='length' are yielded normally."""
|
||
stream_chunks = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
tool_calls=[
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id="call_456",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"test_arg": "value"}',
|
||
),
|
||
index=0,
|
||
)
|
||
],
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[StreamingChoices(finish_reason="length", delta=Delta())]
|
||
),
|
||
]
|
||
mock_completion.return_value = iter(stream_chunks)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
|
||
assert len(responses) == 1
|
||
final_response = responses[0]
|
||
assert final_response.content.role == "model"
|
||
assert len(final_response.content.parts) == 1
|
||
|
||
function_call = final_response.content.parts[0].function_call
|
||
assert function_call.name == "test_function"
|
||
assert function_call.id == "call_456"
|
||
assert function_call.args == {"test_arg": "value"}
|
||
assert final_response.finish_reason == types.FinishReason.MAX_TOKENS
|
||
assert final_response.error_code == types.FinishReason.MAX_TOKENS
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_text_truncated_by_max_tokens(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
"""Tests that text responses with finish_reason='length' set MAX_TOKENS error."""
|
||
stream_chunks = [
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None,
|
||
delta=Delta(
|
||
role="assistant",
|
||
content="Hello, I am",
|
||
),
|
||
)
|
||
]
|
||
),
|
||
ModelResponseStream(
|
||
choices=[StreamingChoices(finish_reason="length", delta=Delta())]
|
||
),
|
||
]
|
||
mock_completion.return_value = iter(stream_chunks)
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Say hello")]
|
||
)
|
||
],
|
||
)
|
||
|
||
responses = [
|
||
response
|
||
async for response in lite_llm_instance.generate_content_async(
|
||
llm_request, stream=True
|
||
)
|
||
]
|
||
|
||
partial_responses = [r for r in responses if r.partial]
|
||
aggregated_responses = [r for r in responses if not r.partial]
|
||
|
||
assert len(partial_responses) == 1
|
||
assert len(aggregated_responses) == 1
|
||
aggregated = aggregated_responses[0]
|
||
assert aggregated.finish_reason == types.FinishReason.MAX_TOKENS
|
||
assert aggregated.error_code == types.FinishReason.MAX_TOKENS
|
||
assert "Maximum tokens reached" in aggregated.error_message
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_generation_params():
|
||
# Test that generation_params are extracted and mapped correctly
|
||
req = LlmRequest(
|
||
contents=[
|
||
types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
temperature=0.33,
|
||
max_output_tokens=123,
|
||
top_p=0.88,
|
||
top_k=7,
|
||
stop_sequences=["foo", "bar"],
|
||
presence_penalty=0.1,
|
||
frequency_penalty=0.2,
|
||
),
|
||
)
|
||
|
||
_, _, _, generation_params = await _get_completion_inputs(
|
||
req, model="gpt-4o-mini"
|
||
)
|
||
assert generation_params["temperature"] == 0.33
|
||
assert generation_params["max_completion_tokens"] == 123
|
||
assert generation_params["top_p"] == 0.88
|
||
assert generation_params["top_k"] == 7
|
||
assert generation_params["stop"] == ["foo", "bar"]
|
||
assert generation_params["presence_penalty"] == 0.1
|
||
assert generation_params["frequency_penalty"] == 0.2
|
||
# Should not include max_output_tokens
|
||
assert "max_output_tokens" not in generation_params
|
||
assert "stop_sequences" not in generation_params
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_empty_generation_params():
|
||
# Test that generation_params is None when no generation parameters are set
|
||
req = LlmRequest(
|
||
contents=[
|
||
types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
|
||
],
|
||
config=types.GenerateContentConfig(),
|
||
)
|
||
|
||
_, _, _, generation_params = await _get_completion_inputs(
|
||
req, model="gpt-4o-mini"
|
||
)
|
||
assert generation_params is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_minimal_config():
|
||
# Test that generation_params is None when config has no generation parameters
|
||
req = LlmRequest(
|
||
contents=[
|
||
types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
system_instruction="test instruction" # Non-generation parameter
|
||
),
|
||
)
|
||
|
||
_, _, _, generation_params = await _get_completion_inputs(
|
||
req, model="gpt-4o-mini"
|
||
)
|
||
assert generation_params is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_partial_generation_params():
|
||
# Test that generation_params is correctly built even with only some parameters
|
||
req = LlmRequest(
|
||
contents=[
|
||
types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
temperature=0.7,
|
||
# Only temperature is set, others are None/default
|
||
),
|
||
)
|
||
|
||
_, _, _, generation_params = await _get_completion_inputs(
|
||
req, model="gpt-4o-mini"
|
||
)
|
||
assert generation_params is not None
|
||
assert generation_params["temperature"] == 0.7
|
||
# Should only contain the temperature parameter
|
||
assert len(generation_params) == 1
|
||
|
||
|
||
def test_function_declaration_to_tool_param_edge_cases():
|
||
"""Test edge cases for function declaration conversion that caused the original bug."""
|
||
from google.adk.models.lite_llm import _function_declaration_to_tool_param
|
||
|
||
# Test function with None parameters (the original bug scenario)
|
||
func_decl = types.FunctionDeclaration(
|
||
name="test_function_none_params",
|
||
description="Function with None parameters",
|
||
parameters=None,
|
||
)
|
||
result = _function_declaration_to_tool_param(func_decl)
|
||
expected = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "test_function_none_params",
|
||
"description": "Function with None parameters",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {},
|
||
},
|
||
},
|
||
}
|
||
assert result == expected
|
||
|
||
# Verify no 'required' field is added when parameters is None
|
||
assert "required" not in result["function"]["parameters"]
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"usage, expected_tokens",
|
||
[
|
||
({"prompt_tokens_details": {"cached_tokens": 123}}, 123),
|
||
(
|
||
{
|
||
"prompt_tokens_details": [
|
||
{"cached_tokens": 50},
|
||
{"cached_tokens": 25},
|
||
]
|
||
},
|
||
75,
|
||
),
|
||
({"cached_prompt_tokens": 45}, 45),
|
||
({"cached_tokens": 67}, 67),
|
||
({"prompt_tokens": 100}, 0),
|
||
({}, 0),
|
||
("not a dict", 0),
|
||
(None, 0),
|
||
({"prompt_tokens_details": {"cached_tokens": "not a number"}}, 0),
|
||
(json.dumps({"cached_tokens": 89}), 89),
|
||
(json.dumps({"some_key": "some_value"}), 0),
|
||
],
|
||
)
|
||
def test_extract_cached_prompt_tokens(usage, expected_tokens):
|
||
from google.adk.models.lite_llm import _extract_cached_prompt_tokens
|
||
|
||
assert _extract_cached_prompt_tokens(usage) == expected_tokens
|
||
|
||
|
||
def test_gemini_via_litellm_warning(monkeypatch):
|
||
"""Test that Gemini via LiteLLM shows warning."""
|
||
# Ensure environment variable is not set
|
||
monkeypatch.delenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", raising=False)
|
||
with warnings.catch_warnings(record=True) as w:
|
||
warnings.simplefilter("always")
|
||
# Test with Google AI Studio Gemini via LiteLLM
|
||
LiteLlm(model="gemini/gemini-2.5-pro-exp-03-25")
|
||
assert len(w) == 1
|
||
assert issubclass(w[0].category, UserWarning)
|
||
assert "[GEMINI_VIA_LITELLM]" in str(w[0].message)
|
||
assert "better performance" in str(w[0].message)
|
||
assert "gemini-2.5-pro-exp-03-25" in str(w[0].message)
|
||
assert "ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS" in str(w[0].message)
|
||
|
||
|
||
def test_gemini_via_litellm_warning_vertex_ai(monkeypatch):
|
||
"""Test that Vertex AI Gemini via LiteLLM shows warning."""
|
||
# Ensure environment variable is not set
|
||
monkeypatch.delenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", raising=False)
|
||
with warnings.catch_warnings(record=True) as w:
|
||
warnings.simplefilter("always")
|
||
# Test with Vertex AI Gemini via LiteLLM
|
||
LiteLlm(model="vertex_ai/gemini-2.5-flash")
|
||
assert len(w) == 1
|
||
assert issubclass(w[0].category, UserWarning)
|
||
assert "[GEMINI_VIA_LITELLM]" in str(w[0].message)
|
||
assert "vertex_ai/gemini-2.5-flash" in str(w[0].message)
|
||
|
||
|
||
def test_gemini_via_litellm_warning_suppressed(monkeypatch):
|
||
"""Test that Gemini via LiteLLM warning can be suppressed."""
|
||
monkeypatch.setenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", "true")
|
||
with warnings.catch_warnings(record=True) as w:
|
||
warnings.simplefilter("always")
|
||
LiteLlm(model="gemini/gemini-2.5-pro-exp-03-25")
|
||
assert len(w) == 0
|
||
|
||
|
||
def test_non_gemini_litellm_no_warning():
|
||
"""Test that non-Gemini models via LiteLLM don't show warning."""
|
||
with warnings.catch_warnings(record=True) as w:
|
||
warnings.simplefilter("always")
|
||
# Test with non-Gemini model
|
||
LiteLlm(model="openai/gpt-4o")
|
||
assert len(w) == 0
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"finish_reason,response_content,expected_content,has_tool_calls",
|
||
[
|
||
("length", "Test response", "Test response", False),
|
||
("stop", "Complete response", "Complete response", False),
|
||
(
|
||
"tool_calls",
|
||
"",
|
||
"",
|
||
True,
|
||
),
|
||
("content_filter", "", "", False),
|
||
],
|
||
ids=["length", "stop", "tool_calls", "content_filter"],
|
||
)
|
||
@pytest.mark.asyncio
|
||
async def test_finish_reason_propagation(
|
||
mock_acompletion,
|
||
lite_llm_instance,
|
||
finish_reason,
|
||
response_content,
|
||
expected_content,
|
||
has_tool_calls,
|
||
):
|
||
"""Test that finish_reason is properly propagated from LiteLLM response."""
|
||
tool_calls = None
|
||
if has_tool_calls:
|
||
tool_calls = [
|
||
ChatCompletionMessageToolCall(
|
||
type="function",
|
||
id="test_id",
|
||
function=Function(
|
||
name="test_function",
|
||
arguments='{"arg": "value"}',
|
||
),
|
||
)
|
||
]
|
||
|
||
mock_response = ModelResponse(
|
||
choices=[
|
||
Choices(
|
||
message=ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content=response_content,
|
||
tool_calls=tool_calls,
|
||
),
|
||
finish_reason=finish_reason,
|
||
)
|
||
]
|
||
)
|
||
mock_acompletion.return_value = mock_response
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
# Verify finish_reason is mapped to FinishReason enum
|
||
assert isinstance(response.finish_reason, types.FinishReason)
|
||
# Verify correct enum mapping using the actual mapping from lite_llm
|
||
assert response.finish_reason == _FINISH_REASON_MAPPING[finish_reason]
|
||
if expected_content:
|
||
assert response.content.parts[0].text == expected_content
|
||
if has_tool_calls:
|
||
assert len(response.content.parts) > 0
|
||
assert response.content.parts[-1].function_call.name == "test_function"
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
|
||
def test_model_response_to_generate_content_response_no_message_with_finish_reason():
|
||
"""Test response with no message but finish_reason returns empty LlmResponse.
|
||
|
||
This test covers issue #3618: when a turn ends with tool calls and no final
|
||
message, we should return an empty LlmResponse instead of raising ValueError.
|
||
"""
|
||
response = ModelResponse(
|
||
model="test_model",
|
||
choices=[{
|
||
"finish_reason": "tool_calls",
|
||
# message is missing/None
|
||
}],
|
||
usage={
|
||
"prompt_tokens": 10,
|
||
"completion_tokens": 5,
|
||
"total_tokens": 15,
|
||
},
|
||
)
|
||
# Force message to be None to guarantee hitting the else branch
|
||
response.choices[0].message = None
|
||
|
||
llm_response = _model_response_to_generate_content_response(response)
|
||
|
||
# Should return empty LlmResponse, not raise ValueError
|
||
assert llm_response.content is not None
|
||
assert llm_response.content.role == "model"
|
||
assert len(llm_response.content.parts) == 0
|
||
# tool_calls maps to STOP
|
||
assert llm_response.finish_reason == types.FinishReason.STOP
|
||
assert llm_response.usage_metadata is not None
|
||
assert llm_response.usage_metadata.prompt_token_count == 10
|
||
assert llm_response.usage_metadata.candidates_token_count == 5
|
||
assert llm_response.model_version == "test_model"
|
||
|
||
|
||
def test_model_response_to_generate_content_response_no_message_no_finish_reason():
|
||
"""Test response with no message and no finish_reason returns empty LlmResponse."""
|
||
response = ModelResponse(
|
||
model="test_model",
|
||
choices=[{
|
||
# Both message and finish_reason are missing
|
||
}],
|
||
)
|
||
# Force message to be None to guarantee hitting the else branch
|
||
response.choices[0].message = None
|
||
|
||
llm_response = _model_response_to_generate_content_response(response)
|
||
|
||
# Should return empty LlmResponse, not raise ValueError
|
||
assert llm_response.content is not None
|
||
assert llm_response.content.role == "model"
|
||
assert len(llm_response.content.parts) == 0
|
||
# finish_reason may be None or have a default value - the important thing
|
||
# is that we don't raise ValueError
|
||
assert llm_response.model_version == "test_model"
|
||
|
||
|
||
def test_model_response_to_generate_content_response_empty_message_dict():
|
||
"""Test response with empty message dict returns empty LlmResponse."""
|
||
response = ModelResponse(
|
||
model="test_model",
|
||
choices=[{
|
||
"message": {}, # Empty dict is falsy
|
||
"finish_reason": "stop",
|
||
}],
|
||
usage={
|
||
"prompt_tokens": 5,
|
||
"completion_tokens": 3,
|
||
"total_tokens": 8,
|
||
},
|
||
)
|
||
# Ensure we test the parsing of an empty message dictionary rather than None.
|
||
|
||
llm_response = _model_response_to_generate_content_response(response)
|
||
|
||
# Should return empty LlmResponse, not raise ValueError
|
||
assert llm_response.content is not None
|
||
assert llm_response.content.role == "model"
|
||
assert len(llm_response.content.parts) == 0
|
||
assert llm_response.finish_reason == types.FinishReason.STOP
|
||
assert llm_response.usage_metadata is not None
|
||
|
||
|
||
def test_model_response_to_generate_content_response_safety_finish_reason():
|
||
"""Test that SAFETY finish reason sets error_code and error_message."""
|
||
response = ModelResponse(
|
||
model="test_model",
|
||
choices=[{
|
||
"finish_reason": "content_filter",
|
||
}],
|
||
)
|
||
# Force message to be None to guarantee hitting the else branch
|
||
response.choices[0].message = None
|
||
|
||
llm_response = _model_response_to_generate_content_response(response)
|
||
|
||
assert llm_response.finish_reason == types.FinishReason.SAFETY
|
||
assert llm_response.error_code == types.FinishReason.SAFETY
|
||
assert llm_response.error_message == "Finished with SAFETY"
|
||
|
||
|
||
@pytest.mark.skip(reason="LiteLLM finish_reason mapping behaviour changed")
|
||
@pytest.mark.asyncio
|
||
async def test_finish_reason_unknown_maps_to_other(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
"""Test that unmapped finish_reason values map to FinishReason.OTHER."""
|
||
# LiteLLM's Choices model normalizes finish_reason values (e.g., "eos" ->
|
||
# "stop") before ADK processes them. To test ADK's own fallback mapping,
|
||
# construct a mock response that bypasses LiteLLM's normalization and
|
||
# returns a raw unmapped finish_reason string.
|
||
mock_choice = MagicMock()
|
||
mock_choice.get = lambda key, default=None: {
|
||
"message": ChatCompletionAssistantMessage(
|
||
role="assistant",
|
||
content="Test response",
|
||
),
|
||
"finish_reason": "totally_unknown_reason",
|
||
}.get(key, default)
|
||
|
||
mock_response = MagicMock()
|
||
mock_response.get = lambda key, default=None: {
|
||
"choices": [mock_choice],
|
||
}.get(key, default)
|
||
mock_response.model = "test_model"
|
||
|
||
mock_acompletion.return_value = mock_response
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
)
|
||
|
||
async for response in lite_llm_instance.generate_content_async(llm_request):
|
||
assert response.content.role == "model"
|
||
# Unknown finish_reason should map to OTHER
|
||
assert isinstance(response.finish_reason, types.FinishReason)
|
||
assert response.finish_reason == types.FinishReason.OTHER
|
||
|
||
mock_acompletion.assert_called_once()
|
||
|
||
|
||
# Tests for provider detection and file_id support
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"model_string, expected_provider",
|
||
[
|
||
# Standard provider/model format
|
||
("openai/gpt-4o", "openai"),
|
||
("azure/gpt-4", "azure"),
|
||
("groq/llama3-70b", "groq"),
|
||
("anthropic/claude-3", "anthropic"),
|
||
("vertex_ai/gemini-pro", "vertex_ai"),
|
||
# Fallback heuristics
|
||
("gpt-4o", "openai"),
|
||
("o1-preview", "openai"),
|
||
("azure-gpt-4", "azure"),
|
||
# Unknown models
|
||
("custom-model", ""),
|
||
("", ""),
|
||
(None, ""),
|
||
],
|
||
)
|
||
def test_get_provider_from_model(model_string, expected_provider):
|
||
"""Test provider extraction from model strings."""
|
||
assert _get_provider_from_model(model_string) == expected_provider
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"provider, expected_in_list",
|
||
[
|
||
("openai", True),
|
||
("azure", True),
|
||
("anthropic", False),
|
||
("vertex_ai", False),
|
||
],
|
||
)
|
||
def test_file_id_required_providers(provider, expected_in_list):
|
||
"""Test that the correct providers require file_id."""
|
||
assert (provider in _FILE_ID_REQUIRED_PROVIDERS) == expected_in_list
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_pdf_openai_uses_file_id(mocker):
|
||
"""Test that PDF files use file_id for OpenAI provider."""
|
||
mock_file_response = mocker.create_autospec(litellm.FileObject)
|
||
mock_file_response.id = "file-abc123"
|
||
mock_acreate_file = AsyncMock(return_value=mock_file_response)
|
||
mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)
|
||
|
||
parts = [
|
||
types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
|
||
]
|
||
content = await _get_content(parts, provider="openai")
|
||
|
||
assert content[0]["type"] == "file"
|
||
assert content[0]["file"]["file_id"] == "file-abc123"
|
||
assert "file_data" not in content[0]["file"]
|
||
|
||
mock_acreate_file.assert_called_once_with(
|
||
file=b"test_pdf_data",
|
||
purpose="assistants",
|
||
custom_llm_provider="openai",
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_pdf_non_openai_uses_file_data():
|
||
"""Test that PDF files use file_data for non-OpenAI providers."""
|
||
parts = [
|
||
types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
|
||
]
|
||
content = await _get_content(parts, provider="anthropic")
|
||
|
||
assert content[0]["type"] == "file"
|
||
assert "file_data" in content[0]["file"]
|
||
assert content[0]["file"]["file_data"].startswith(
|
||
"data:application/pdf;base64,"
|
||
)
|
||
assert "file_id" not in content[0]["file"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_content_pdf_azure_uses_file_id(mocker):
|
||
"""Test that PDF files use file_id for Azure provider."""
|
||
mock_file_response = mocker.create_autospec(litellm.FileObject)
|
||
mock_file_response.id = "file-xyz789"
|
||
mock_acreate_file = AsyncMock(return_value=mock_file_response)
|
||
mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)
|
||
|
||
parts = [
|
||
types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
|
||
]
|
||
content = await _get_content(parts, provider="azure")
|
||
|
||
assert content[0]["type"] == "file"
|
||
assert content[0]["file"]["file_id"] == "file-xyz789"
|
||
|
||
mock_acreate_file.assert_called_once_with(
|
||
file=b"test_pdf_data",
|
||
purpose="assistants",
|
||
custom_llm_provider="azure",
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_openai_file_upload(mocker):
|
||
"""Test that _get_completion_inputs uploads files for OpenAI models."""
|
||
mock_file_response = mocker.create_autospec(litellm.FileObject)
|
||
mock_file_response.id = "file-uploaded123"
|
||
mock_acreate_file = AsyncMock(return_value=mock_file_response)
|
||
mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)
|
||
|
||
pdf_part = types.Part.from_bytes(
|
||
data=b"test_pdf_content", mime_type="application/pdf"
|
||
)
|
||
llm_request = LlmRequest(
|
||
model="openai/gpt-4o",
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Analyze this PDF"),
|
||
pdf_part,
|
||
],
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(tools=[]),
|
||
)
|
||
|
||
messages, tools, response_format, generation_params = (
|
||
await _get_completion_inputs(llm_request, model="openai/gpt-4o")
|
||
)
|
||
|
||
assert len(messages) == 1
|
||
assert messages[0]["role"] == "user"
|
||
content = messages[0]["content"]
|
||
assert len(content) == 2
|
||
assert content[0]["type"] == "text"
|
||
assert content[0]["text"] == "Analyze this PDF"
|
||
assert content[1]["type"] == "file"
|
||
assert content[1]["file"]["file_id"] == "file-uploaded123"
|
||
|
||
mock_acreate_file.assert_called_once()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_completion_inputs_non_openai_no_file_upload(mocker):
|
||
"""Test that _get_completion_inputs does not upload files for non-OpenAI models."""
|
||
mock_acreate_file = AsyncMock()
|
||
mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)
|
||
|
||
pdf_part = types.Part.from_bytes(
|
||
data=b"test_pdf_content", mime_type="application/pdf"
|
||
)
|
||
llm_request = LlmRequest(
|
||
model="anthropic/claude-3-opus",
|
||
contents=[
|
||
types.Content(
|
||
role="user",
|
||
parts=[
|
||
types.Part.from_text(text="Analyze this PDF"),
|
||
pdf_part,
|
||
],
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(tools=[]),
|
||
)
|
||
|
||
messages, tools, response_format, generation_params = (
|
||
await _get_completion_inputs(llm_request, model="anthropic/claude-3-opus")
|
||
)
|
||
|
||
assert len(messages) == 1
|
||
content = messages[0]["content"]
|
||
assert content[1]["type"] == "file"
|
||
assert "file_data" in content[1]["file"]
|
||
assert "file_id" not in content[1]["file"]
|
||
|
||
mock_acreate_file.assert_not_called()
|
||
|
||
|
||
class TestRedirectLitellmLoggersToStdout(unittest.TestCase):
|
||
"""Tests for _redirect_litellm_loggers_to_stdout function."""
|
||
|
||
def test_redirects_stderr_handler_to_stdout(self):
|
||
"""Test that handlers pointing to stderr are redirected to stdout."""
|
||
test_logger = logging.getLogger("LiteLLM")
|
||
# Create a handler pointing to stderr
|
||
handler = logging.StreamHandler(sys.stderr)
|
||
test_logger.addHandler(handler)
|
||
|
||
try:
|
||
self.assertIs(handler.stream, sys.stderr)
|
||
|
||
_redirect_litellm_loggers_to_stdout()
|
||
|
||
self.assertIs(handler.stream, sys.stdout)
|
||
finally:
|
||
# Clean up
|
||
test_logger.removeHandler(handler)
|
||
|
||
def test_preserves_stdout_handler(self):
|
||
"""Test that handlers already pointing to stdout are not modified."""
|
||
test_logger = logging.getLogger("LiteLLM Proxy")
|
||
# Create a handler already pointing to stdout
|
||
handler = logging.StreamHandler(sys.stdout)
|
||
test_logger.addHandler(handler)
|
||
|
||
try:
|
||
_redirect_litellm_loggers_to_stdout()
|
||
|
||
self.assertIs(handler.stream, sys.stdout)
|
||
finally:
|
||
# Clean up
|
||
test_logger.removeHandler(handler)
|
||
|
||
def test_does_not_affect_non_stream_handlers(self):
|
||
"""Test that non-StreamHandler handlers are not affected."""
|
||
test_logger = logging.getLogger("LiteLLM Router")
|
||
# Create a FileHandler (not a StreamHandler)
|
||
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
|
||
temp_file_name = temp_file.name
|
||
with contextlib.closing(
|
||
logging.FileHandler(temp_file_name)
|
||
) as file_handler:
|
||
test_logger.addHandler(file_handler)
|
||
|
||
try:
|
||
_redirect_litellm_loggers_to_stdout()
|
||
# FileHandler should not be modified (it doesn't point to stderr or stdout)
|
||
self.assertEqual(file_handler.baseFilename, temp_file_name)
|
||
finally:
|
||
# Clean up
|
||
test_logger.removeHandler(file_handler)
|
||
os.unlink(temp_file_name)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"logger_name",
|
||
["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"],
|
||
ids=["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"],
|
||
)
|
||
def test_handles_litellm_logger_names(logger_name):
|
||
"""Test that LiteLLM logger names are processed."""
|
||
test_logger = logging.getLogger(logger_name)
|
||
handler = logging.StreamHandler(sys.stderr)
|
||
test_logger.addHandler(handler)
|
||
|
||
try:
|
||
_redirect_litellm_loggers_to_stdout()
|
||
|
||
assert handler.stream is sys.stdout
|
||
finally:
|
||
# Clean up
|
||
test_logger.removeHandler(handler)
|
||
|
||
|
||
# ── Anthropic thinking_blocks tests ─────────────────────────────
|
||
|
||
|
||
def test_is_anthropic_provider():
|
||
"""Verify _is_anthropic_provider matches known Claude provider prefixes."""
|
||
assert _is_anthropic_provider("anthropic")
|
||
assert _is_anthropic_provider("bedrock")
|
||
assert _is_anthropic_provider("vertex_ai")
|
||
assert _is_anthropic_provider("ANTHROPIC") # case-insensitive
|
||
assert not _is_anthropic_provider("openai")
|
||
assert not _is_anthropic_provider("")
|
||
assert not _is_anthropic_provider(None)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"model_string,expected",
|
||
[
|
||
("anthropic/claude-4-sonnet", True),
|
||
("anthropic/claude-3-5-sonnet-20241022", True),
|
||
("Anthropic/Claude-4-Opus", True),
|
||
("bedrock/anthropic.claude-3-5-sonnet", True),
|
||
("bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0", True),
|
||
("bedrock/claude-3-5-sonnet", True),
|
||
("vertex_ai/claude-3-5-sonnet@20241022", True),
|
||
("openai/gpt-4o", False),
|
||
("gemini/gemini-2.5-pro", False),
|
||
("vertex_ai/gemini-2.5-flash", False),
|
||
("bedrock/amazon.titan-text-express-v1", False),
|
||
],
|
||
ids=[
|
||
"anthropic-prefix",
|
||
"anthropic-versioned",
|
||
"anthropic-uppercase",
|
||
"bedrock-anthropic-dot",
|
||
"bedrock-us-anthropic",
|
||
"bedrock-claude",
|
||
"vertex-claude",
|
||
"openai-no-match",
|
||
"gemini-no-match",
|
||
"vertex-gemini-no-match",
|
||
"bedrock-non-anthropic",
|
||
],
|
||
)
|
||
def test_is_anthropic_model(model_string, expected):
|
||
assert _is_anthropic_model(model_string) is expected
|
||
|
||
|
||
def test_extract_reasoning_value_prefers_thinking_blocks():
|
||
"""thinking_blocks (Anthropic format with signatures) take priority."""
|
||
thinking_blocks = [
|
||
{"type": "thinking", "thinking": "step 1", "signature": "c2lnX2E="},
|
||
{"type": "thinking", "thinking": "step 2", "signature": "c2lnX2I="},
|
||
]
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"thinking_blocks": thinking_blocks,
|
||
"reasoning_content": "flat reasoning",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result is thinking_blocks
|
||
|
||
|
||
def test_extract_reasoning_value_falls_back_without_thinking_blocks():
|
||
"""When thinking_blocks is absent, falls back to reasoning_content."""
|
||
message = {
|
||
"role": "assistant",
|
||
"content": "Answer",
|
||
"reasoning_content": "flat reasoning",
|
||
}
|
||
result = _extract_reasoning_value(message)
|
||
assert result == "flat reasoning"
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_preserves_base64_signature():
|
||
"""Base64 signatures are decoded to raw bytes on thought parts."""
|
||
thinking_blocks = [
|
||
{"type": "thinking", "thinking": "step 1", "signature": "c2lnX2E="},
|
||
{"type": "thinking", "thinking": "step 2", "signature": "c2lnX2I="},
|
||
]
|
||
parts = _convert_reasoning_value_to_parts(thinking_blocks)
|
||
assert len(parts) == 2
|
||
assert parts[0].text == "step 1"
|
||
assert parts[0].thought is True
|
||
assert parts[0].thought_signature == b"sig_a"
|
||
assert parts[1].text == "step 2"
|
||
assert parts[1].thought_signature == b"sig_b"
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_raw_signature_falls_back_to_utf8():
|
||
"""Non-base64 signatures are preserved as utf-8 bytes."""
|
||
thinking_blocks = [
|
||
{"type": "thinking", "thinking": "step 1", "signature": "sig_raw"},
|
||
]
|
||
parts = _convert_reasoning_value_to_parts(thinking_blocks)
|
||
assert len(parts) == 1
|
||
assert parts[0].text == "step 1"
|
||
assert parts[0].thought_signature == b"sig_raw"
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_skips_redacted_blocks():
|
||
"""Redacted thinking blocks are excluded from parts."""
|
||
thinking_blocks = [
|
||
{"type": "thinking", "thinking": "visible", "signature": "c2lnMQ=="},
|
||
{"type": "redacted", "data": "hidden"},
|
||
]
|
||
parts = _convert_reasoning_value_to_parts(thinking_blocks)
|
||
assert len(parts) == 1
|
||
assert parts[0].text == "visible"
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_preserves_signature_only_blocks():
|
||
"""Signature-only blocks (empty text) are preserved for streaming aggregation.
|
||
|
||
Anthropic emits the block_stop signature as a delta with empty thinking text.
|
||
Dropping it would lose the signature, breaking multi-turn thinking continuity.
|
||
Blocks with neither text nor signature are still skipped.
|
||
"""
|
||
thinking_blocks = [
|
||
{"type": "thinking", "thinking": "", "signature": "c2lnMQ=="},
|
||
{"type": "thinking", "thinking": "real thought", "signature": "c2lnMg=="},
|
||
{
|
||
"type": "thinking",
|
||
"thinking": "",
|
||
"signature": "",
|
||
}, # fully empty: drop
|
||
]
|
||
parts = _convert_reasoning_value_to_parts(thinking_blocks)
|
||
assert len(parts) == 2
|
||
assert parts[0].text == ""
|
||
assert parts[0].thought is True
|
||
assert parts[0].thought_signature == b"sig1"
|
||
assert parts[1].text == "real thought"
|
||
assert parts[1].thought_signature == b"sig2"
|
||
|
||
|
||
def test_aggregate_streaming_thought_parts():
|
||
"""Tests aggregating fragmented streaming thought parts and multiple blocks."""
|
||
parts = [
|
||
types.Part(text="First block ", thought=True),
|
||
types.Part(text="text.", thought=True),
|
||
types.Part(text="", thought=True, thought_signature=b"sig1"),
|
||
types.Part(text="Second block", thought=True, thought_signature=b"sig2"),
|
||
types.Part(text="Trailing without sig", thought=True),
|
||
]
|
||
aggregated = _aggregate_streaming_thought_parts(parts)
|
||
assert len(aggregated) == 3
|
||
assert aggregated[0].text == "First block text."
|
||
assert aggregated[0].thought_signature == b"sig1"
|
||
assert aggregated[1].text == "Second block"
|
||
assert aggregated[1].thought_signature == b"sig2"
|
||
assert aggregated[2].text == "Trailing without sig"
|
||
assert aggregated[2].thought_signature is None
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_flat_string_unchanged():
|
||
"""Flat string reasoning still produces thought parts without signature."""
|
||
parts = _convert_reasoning_value_to_parts("simple reasoning text")
|
||
assert len(parts) == 1
|
||
assert parts[0].text == "simple reasoning text"
|
||
assert parts[0].thought is True
|
||
assert parts[0].thought_signature is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_outputs_thinking_blocks():
|
||
"""Anthropic model messages base64-encode thought signatures."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(
|
||
text="deep thought",
|
||
thought=True,
|
||
thought_signature=b"sig_round_trip",
|
||
),
|
||
types.Part(text="Hello!"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content, model="anthropic/claude-4-sonnet"
|
||
)
|
||
assert result["role"] == "assistant"
|
||
assert result["thinking_blocks"] == [{
|
||
"type": "thinking",
|
||
"thinking": "deep thought",
|
||
"signature": "c2lnX3JvdW5kX3RyaXA=",
|
||
}]
|
||
assert result.get("reasoning_content") is None
|
||
assert result["content"] == "Hello!"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_model_round_trip_preserves_signature():
|
||
"""Decoded signatures are re-encoded when rebuilding Anthropic messages."""
|
||
response_message = {
|
||
"role": "assistant",
|
||
"content": "Final answer",
|
||
"thinking_blocks": [{
|
||
"type": "thinking",
|
||
"thinking": "Let me reason...",
|
||
"signature": "c2lnX2E=",
|
||
}],
|
||
}
|
||
|
||
parts = _convert_reasoning_value_to_parts(
|
||
_extract_reasoning_value(response_message)
|
||
)
|
||
content = types.Content(
|
||
role="model",
|
||
parts=parts + [types.Part(text="Final answer")],
|
||
)
|
||
|
||
result = await _content_to_message_param(
|
||
content,
|
||
provider="anthropic",
|
||
model="anthropic/claude-4-sonnet",
|
||
)
|
||
|
||
assert result["thinking_blocks"] == [{
|
||
"type": "thinking",
|
||
"thinking": "Let me reason...",
|
||
"signature": "c2lnX2E=",
|
||
}]
|
||
assert result.get("reasoning_content") is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_split_thinking_and_signature():
|
||
"""Combines separate thinking and signature parts into a single thinking_block."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="deep thought", thought=True),
|
||
types.Part(
|
||
text="", thought=True, thought_signature=b"sig_round_trip"
|
||
),
|
||
types.Part(text="Hello!"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content, model="anthropic/claude-4-sonnet"
|
||
)
|
||
assert result["role"] == "assistant"
|
||
assert "thinking_blocks" in result
|
||
assert result.get("reasoning_content") is None
|
||
blocks = result["thinking_blocks"]
|
||
assert len(blocks) == 1
|
||
assert blocks[0]["type"] == "thinking"
|
||
assert blocks[0]["thinking"] == "deep thought"
|
||
assert blocks[0]["signature"] == "c2lnX3JvdW5kX3RyaXA="
|
||
assert result["content"] == "Hello!"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_non_anthropic_uses_reasoning_content():
|
||
"""For non-Anthropic models, reasoning_content is used as before."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking text", thought=True),
|
||
types.Part(text="Answer"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(content, model="openai/gpt-4o")
|
||
assert result["role"] == "assistant"
|
||
assert result.get("reasoning_content") == "thinking text"
|
||
assert "thinking_blocks" not in result
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_anthropic_provider_thinking_blocks_round_trip():
|
||
"""End-to-end: thinking_blocks in response stay intact for Anthropic provider."""
|
||
response_message = {
|
||
"role": "assistant",
|
||
"content": "Final answer",
|
||
"thinking_blocks": [
|
||
{
|
||
"type": "thinking",
|
||
"thinking": "Let me reason...",
|
||
"signature": "c2lnX2E=",
|
||
},
|
||
],
|
||
}
|
||
|
||
reasoning_value = _extract_reasoning_value(response_message)
|
||
assert isinstance(reasoning_value, list)
|
||
|
||
parts = _convert_reasoning_value_to_parts(reasoning_value)
|
||
assert len(parts) == 1
|
||
assert parts[0].thought_signature == b"sig_a"
|
||
|
||
all_parts = parts + [
|
||
types.Part(text="Final answer"),
|
||
types.Part.from_function_call(name="add", args={"a": 1, "b": 2}),
|
||
]
|
||
content = types.Content(role="model", parts=all_parts)
|
||
|
||
msg = await _content_to_message_param(content, provider="anthropic")
|
||
assert isinstance(msg["content"], list)
|
||
assert msg["content"][0] == {
|
||
"type": "thinking",
|
||
"thinking": "Let me reason...",
|
||
"signature": "c2lnX2E=",
|
||
}
|
||
assert msg["content"][1] == {"type": "text", "text": "Final answer"}
|
||
assert msg["tool_calls"] is not None
|
||
assert len(msg["tool_calls"]) == 1
|
||
assert msg["tool_calls"][0]["function"]["name"] == "add"
|
||
assert msg.get("reasoning_content") is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_no_signature_falls_back():
|
||
"""Anthropic model with thought parts but no signatures uses reasoning_content."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking without sig", thought=True),
|
||
types.Part(text="Response"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content, model="anthropic/claude-4-sonnet"
|
||
)
|
||
assert result.get("reasoning_content") == "thinking without sig"
|
||
assert "thinking_blocks" not in result
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"provider,model,expected",
|
||
[
|
||
("anthropic", "anthropic/claude-3-5-sonnet", True),
|
||
("anthropic", "", True), # anthropic always routes to Claude
|
||
("bedrock", "bedrock/anthropic.claude-3-5-sonnet", True),
|
||
("bedrock", "bedrock/meta.llama3-70b-instruct-v1:0", False),
|
||
("vertex_ai", "vertex_ai/claude-3-5-sonnet@20241022", True),
|
||
("vertex_ai", "vertex_ai/gemini-2.5-flash", False),
|
||
("openai", "openai/gpt-4o", False),
|
||
("", "", False),
|
||
],
|
||
)
|
||
def test_is_anthropic_route(provider, model, expected):
|
||
assert _is_anthropic_route(provider, model) is expected
|
||
|
||
|
||
def test_convert_reasoning_value_to_parts_empty_thinking_does_not_fall_through():
|
||
"""An empty thinking block is skipped, not parsed via the text fallback."""
|
||
thinking_blocks = [
|
||
{
|
||
"type": "thinking",
|
||
"thinking": "",
|
||
"text": "leaked",
|
||
"signature": "",
|
||
},
|
||
]
|
||
parts = _convert_reasoning_value_to_parts(thinking_blocks)
|
||
assert parts == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_bedrock_non_claude_no_thinking_blocks():
|
||
"""bedrock + non-Claude model must not get Anthropic thinking-block formatting."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking text", thought=True),
|
||
types.Part(text="Answer"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content,
|
||
provider="bedrock",
|
||
model="bedrock/meta.llama3-70b-instruct-v1:0",
|
||
)
|
||
assert result.get("reasoning_content") == "thinking text"
|
||
assert "thinking_blocks" not in result
|
||
body = result.get("content")
|
||
assert not (
|
||
isinstance(body, list)
|
||
and any(isinstance(b, dict) and b.get("type") == "thinking" for b in body)
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_bedrock_claude_embeds_thinking_blocks():
|
||
"""bedrock + Claude model embeds thinking blocks in the content list."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking text", thought=True),
|
||
types.Part(text="Answer"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content,
|
||
provider="bedrock",
|
||
model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||
)
|
||
assert isinstance(result["content"], list)
|
||
assert result["content"][0] == {
|
||
"type": "thinking",
|
||
"thinking": "thinking text",
|
||
}
|
||
assert result.get("reasoning_content") is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_vertex_gemini_no_thinking_blocks():
|
||
"""vertex_ai + Gemini model must not get Anthropic thinking-block formatting."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking text", thought=True),
|
||
types.Part(text="Answer"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content,
|
||
provider="vertex_ai",
|
||
model="vertex_ai/gemini-2.5-flash",
|
||
)
|
||
assert result.get("reasoning_content") == "thinking text"
|
||
assert "thinking_blocks" not in result
|
||
body = result.get("content")
|
||
assert not (
|
||
isinstance(body, list)
|
||
and any(isinstance(b, dict) and b.get("type") == "thinking" for b in body)
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_provider_embeds_thinking_blocks():
|
||
"""provider 'anthropic' always embeds thinking blocks in the content list."""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
types.Part(text="thinking text", thought=True),
|
||
types.Part(text="Answer"),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content,
|
||
provider="anthropic",
|
||
model="anthropic/claude-3-5-sonnet",
|
||
)
|
||
assert isinstance(result["content"], list)
|
||
assert result["content"][0] == {
|
||
"type": "thinking",
|
||
"thinking": "thinking text",
|
||
}
|
||
assert result.get("reasoning_content") is None
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_content_to_message_param_anthropic_aggregates_streaming_split_thinking():
|
||
"""Streaming splits one Anthropic thinking block across many parts:
|
||
text-only chunks followed by a signature-only chunk at block_stop.
|
||
_content_to_message_param must re-join them into one thinking_block.
|
||
"""
|
||
content = types.Content(
|
||
role="model",
|
||
parts=[
|
||
# Text-only chunks from streaming deltas (no signature)
|
||
types.Part(text="The user wants ", thought=True),
|
||
types.Part(text="GST research ", thought=True),
|
||
types.Part(text="on secondment.", thought=True),
|
||
# Final signature-only chunk (empty text, signature carries the whole block)
|
||
types.Part(
|
||
text="", thought=True, thought_signature=b"ErEDClsIDBACGAIfull"
|
||
),
|
||
# Non-thought response content
|
||
types.Part.from_function_call(name="create_plan", args={"q": "test"}),
|
||
],
|
||
)
|
||
result = await _content_to_message_param(
|
||
content, model="anthropic/claude-4-sonnet"
|
||
)
|
||
# One aggregated thinking block with combined text and the block's signature
|
||
blocks = result["thinking_blocks"]
|
||
assert len(blocks) == 1
|
||
assert blocks[0]["type"] == "thinking"
|
||
assert blocks[0]["thinking"] == "The user wants GST research on secondment."
|
||
assert blocks[0]["signature"] == "RXJFRENsc0lEQkFDR0FJZnVsbA=="
|
||
# Legacy reasoning_content is not set when the Anthropic branch takes
|
||
assert result.get("reasoning_content") is None
|
||
|
||
|
||
def test_model_response_to_chunk_preserves_signature_only_delta():
|
||
"""Anthropic streams a final thinking delta where content and
|
||
reasoning_content are empty but thinking_blocks carries the signature.
|
||
_has_meaningful_signal must recognize thinking_blocks as signal so the
|
||
signature survives into a ReasoningChunk.
|
||
"""
|
||
stream = ModelResponseStream(
|
||
id="x",
|
||
created=0,
|
||
model="claude",
|
||
choices=[
|
||
StreamingChoices(
|
||
index=0,
|
||
delta=Delta(
|
||
role=None,
|
||
content="",
|
||
reasoning_content="",
|
||
thinking_blocks=[{
|
||
"type": "thinking",
|
||
"thinking": "",
|
||
"signature": "SignatureOnlyChunk",
|
||
}],
|
||
),
|
||
)
|
||
],
|
||
)
|
||
chunks = list(_model_response_to_chunk(stream))
|
||
reasoning_chunks = [c for c, _ in chunks if isinstance(c, ReasoningChunk)]
|
||
assert len(reasoning_chunks) == 1
|
||
parts = reasoning_chunks[0].parts
|
||
assert len(parts) == 1
|
||
assert parts[0].text == ""
|
||
assert parts[0].thought is True
|
||
assert parts[0].thought_signature == b"SignatureOnlyChunk"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.parametrize(
|
||
"log_level,should_call",
|
||
[
|
||
(logging.WARNING, False),
|
||
(logging.INFO, False),
|
||
(logging.DEBUG, True),
|
||
],
|
||
)
|
||
async def test_generate_content_async_skips_request_log_build_above_debug(
|
||
mock_acompletion, lite_llm_instance, log_level, should_call
|
||
):
|
||
del mock_acompletion # unused; lite_llm_instance is wired to it
|
||
litellm_logger = logging.getLogger("google_adk.google.adk.models.lite_llm")
|
||
original_level = litellm_logger.level
|
||
litellm_logger.setLevel(log_level)
|
||
try:
|
||
with patch(
|
||
"google.adk.models.lite_llm._build_request_log",
|
||
return_value="log",
|
||
) as mock_build:
|
||
async for _ in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION
|
||
):
|
||
pass
|
||
|
||
assert mock_build.called is should_call
|
||
finally:
|
||
litellm_logger.setLevel(original_level)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"file_uri, expected",
|
||
[
|
||
("file-abc123", True),
|
||
("assistant-abc123", True),
|
||
("https://example.com/file.pdf", False),
|
||
("not-a-file-id", False),
|
||
("", False),
|
||
("FILE-abc123", False),
|
||
],
|
||
)
|
||
def test_looks_like_openai_file_id(file_uri, expected):
|
||
"""Both `file-` and `assistant-` (Azure assistants) prefixes count as OpenAI file IDs."""
|
||
assert _looks_like_openai_file_id(file_uri) is expected
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"file_uri, expected",
|
||
[
|
||
("file-abc123", "file-<redacted>"),
|
||
("assistant-abc123", "assistant-<redacted>"),
|
||
],
|
||
)
|
||
def test_redact_file_uri_for_log_openai_prefixes(file_uri, expected):
|
||
"""OpenAI-style IDs are redacted while preserving the prefix kind."""
|
||
assert _redact_file_uri_for_log(file_uri) == expected
|
||
|
||
|
||
def test_redact_file_uri_for_log_uses_display_name_when_provided():
|
||
assert (
|
||
_redact_file_uri_for_log("file-abc123", display_name="my.pdf") == "my.pdf"
|
||
)
|
||
|
||
|
||
def test_redact_file_uri_for_log_http_url_keeps_scheme_and_tail():
|
||
assert (
|
||
_redact_file_uri_for_log("https://example.com/path/file.pdf")
|
||
== "https://<redacted>/file.pdf"
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_passes_http_options_headers_as_extra_headers(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
"""Test that http_options.headers from LlmRequest are forwarded to litellm."""
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(
|
||
headers={"X-User-Id": "user-123", "X-Trace-Id": "trace-abc"}
|
||
)
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert "extra_headers" in kwargs
|
||
assert kwargs["extra_headers"]["X-User-Id"] == "user-123"
|
||
assert kwargs["extra_headers"]["X-Trace-Id"] == "trace-abc"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_merges_http_options_with_existing_extra_headers(
|
||
mock_response,
|
||
):
|
||
"""Test that http_options.headers merge with pre-existing extra_headers."""
|
||
mock_acompletion = AsyncMock(return_value=mock_response)
|
||
mock_client = MockLLMClient(mock_acompletion, Mock())
|
||
# Create instance with pre-existing extra_headers via kwargs
|
||
lite_llm_with_extra = LiteLlm(
|
||
model="test_model",
|
||
llm_client=mock_client,
|
||
extra_headers={"X-Api-Key": "secret-key"},
|
||
)
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(headers={"X-User-Id": "user-456"})
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_with_extra.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert "extra_headers" in kwargs
|
||
# Both existing and new headers should be present
|
||
assert kwargs["extra_headers"]["X-Api-Key"] == "secret-key"
|
||
assert kwargs["extra_headers"]["X-User-Id"] == "user-456"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_http_options_headers_override_existing(
|
||
mock_response,
|
||
):
|
||
"""Test that http_options.headers override same-key extra_headers from init."""
|
||
mock_acompletion = AsyncMock(return_value=mock_response)
|
||
mock_client = MockLLMClient(mock_acompletion, Mock())
|
||
lite_llm_with_extra = LiteLlm(
|
||
model="test_model",
|
||
llm_client=mock_client,
|
||
extra_headers={"X-Override-Me": "old-value"},
|
||
)
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(headers={"X-Override-Me": "new-value"})
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_with_extra.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
# Request-level headers should override init-level headers
|
||
assert kwargs["extra_headers"]["X-Override-Me"] == "new-value"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_passes_http_options_timeout(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
"""Test that http_options.timeout is forwarded to litellm."""
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(timeout=30000)
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert "timeout" in kwargs
|
||
assert kwargs["timeout"] == 30000
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_passes_http_options_retry_options(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
"""Test that http_options.retry_options is forwarded to litellm."""
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(
|
||
retry_options=types.HttpRetryOptions(
|
||
attempts=3,
|
||
)
|
||
)
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert "num_retries" in kwargs
|
||
assert kwargs["num_retries"] == 3
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_generate_content_async_passes_http_options_extra_body(
|
||
mock_acompletion, lite_llm_instance
|
||
):
|
||
"""Test that http_options.extra_body is forwarded to litellm."""
|
||
|
||
llm_request = LlmRequest(
|
||
contents=[
|
||
types.Content(
|
||
role="user", parts=[types.Part.from_text(text="Test prompt")]
|
||
)
|
||
],
|
||
config=types.GenerateContentConfig(
|
||
http_options=types.HttpOptions(
|
||
extra_body={"custom_field": "custom_value", "priority": "high"}
|
||
)
|
||
),
|
||
)
|
||
|
||
async for _ in lite_llm_instance.generate_content_async(llm_request):
|
||
pass
|
||
|
||
mock_acompletion.assert_called_once()
|
||
_, kwargs = mock_acompletion.call_args
|
||
assert "extra_body" in kwargs
|
||
assert kwargs["extra_body"]["custom_field"] == "custom_value"
|
||
assert kwargs["extra_body"]["priority"] == "high"
|
||
|
||
|
||
def _split_into_chunks(text, sizes):
|
||
pieces = []
|
||
pos = 0
|
||
for size in sizes:
|
||
pieces.append(text[pos : pos + size])
|
||
pos += size
|
||
if pos < len(text):
|
||
pieces.append(text[pos:])
|
||
return pieces
|
||
|
||
|
||
def test_brace_depth_tracker_simple_object():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed('{"a": 1}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_empty_object():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed("{}") is True
|
||
|
||
|
||
def test_brace_depth_tracker_only_opens():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed('{"a": ') is False
|
||
|
||
|
||
def test_brace_depth_tracker_completes_after_more_fragments():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed('{"a": ') is False
|
||
assert tracker.feed('"b"') is False
|
||
assert tracker.feed("}") is True
|
||
|
||
|
||
def test_brace_depth_tracker_nested_objects():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed('{"a": {"b": {"c": 1}}}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_nested_split_across_fragments():
|
||
tracker = _BraceDepthTracker()
|
||
fragments = _split_into_chunks(
|
||
'{"a": {"b": {"c": 1}, "d": [1, 2, 3]}, "e": "f"}', [3, 5, 4, 7, 9, 2, 1]
|
||
)
|
||
closes = [tracker.feed(f) for f in fragments]
|
||
assert sum(closes) == 1
|
||
assert closes[-1] is True
|
||
|
||
|
||
def test_brace_depth_tracker_string_with_braces_ignored():
|
||
tracker = _BraceDepthTracker()
|
||
# Braces inside strings must not affect depth.
|
||
assert tracker.feed('{"x": "{}{{}}"}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_string_with_braces_split_across_fragments():
|
||
tracker = _BraceDepthTracker()
|
||
fragments = ['{"x": "', "abc{def", "}ghi", '"}']
|
||
closes = [tracker.feed(f) for f in fragments]
|
||
assert closes == [False, False, False, True]
|
||
|
||
|
||
def test_brace_depth_tracker_escaped_quote_in_string():
|
||
tracker = _BraceDepthTracker()
|
||
# Escaped quote should not end the string; the trailing } closes the obj.
|
||
assert tracker.feed(r'{"x": "a\"b}c"}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_escaped_backslash_then_quote_ends_string():
|
||
tracker = _BraceDepthTracker()
|
||
# \\ is an escaped backslash; the next " ends the string. Then } closes.
|
||
assert tracker.feed(r'{"x": "a\\"}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_escape_split_across_fragments():
|
||
tracker = _BraceDepthTracker()
|
||
# Backslash arrives in one fragment, the escaped quote in the next.
|
||
fragments = ['{"x": "a', "\\", '"', 'b"}']
|
||
closes = [tracker.feed(f) for f in fragments]
|
||
assert closes == [False, False, False, True]
|
||
|
||
|
||
def test_brace_depth_tracker_two_consecutive_objects():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed('{"a": 1}{"b": 2}') is True
|
||
|
||
|
||
def test_brace_depth_tracker_one_char_at_a_time():
|
||
tracker = _BraceDepthTracker()
|
||
text = '{"key": {"nested": "v{}al"}, "n": 42}'
|
||
closes = [tracker.feed(ch) for ch in text]
|
||
assert sum(closes) == 1
|
||
assert closes[-1] is True
|
||
|
||
|
||
def test_brace_depth_tracker_leading_whitespace_ignored():
|
||
tracker = _BraceDepthTracker()
|
||
assert tracker.feed(' \n {"a": 1}') is True
|
||
|
||
|
||
def _function_chunks_for_args(arg_fragments):
|
||
return [
|
||
FunctionChunk(
|
||
id="call_xyz" if i == 0 else None,
|
||
name="my_func" if i == 0 else None,
|
||
args=fragment,
|
||
index=0,
|
||
)
|
||
for i, fragment in enumerate(arg_fragments)
|
||
]
|
||
|
||
|
||
def _stream_chunks_from_function_chunks(function_chunks):
|
||
stream = []
|
||
for chunk in function_chunks:
|
||
delta_kwargs = {"role": "assistant"}
|
||
if chunk.args is not None:
|
||
delta_kwargs["tool_calls"] = [
|
||
ChatCompletionDeltaToolCall(
|
||
type="function",
|
||
id=chunk.id,
|
||
function=Function(name=chunk.name, arguments=chunk.args),
|
||
index=chunk.index,
|
||
)
|
||
]
|
||
stream.append(
|
||
ModelResponseStream(
|
||
choices=[
|
||
StreamingChoices(
|
||
finish_reason=None, delta=Delta(**delta_kwargs)
|
||
)
|
||
]
|
||
)
|
||
)
|
||
stream.append(
|
||
ModelResponseStream(
|
||
choices=[StreamingChoices(finish_reason="tool_calls", delta=Delta())]
|
||
)
|
||
)
|
||
return stream
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_call_args_assembled_from_many_fragments(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
full_args = (
|
||
'{"city": "San Francisco", "details": {"radius": 5, "tags": ["a{}",'
|
||
' "b\\"c"]}}'
|
||
)
|
||
fragments = _split_into_chunks(full_args, [4, 6, 1, 8, 3, 11, 2, 9, 7, 5, 1])
|
||
mock_completion.return_value = iter(
|
||
_stream_chunks_from_function_chunks(_function_chunks_for_args(fragments))
|
||
)
|
||
|
||
responses = [
|
||
r
|
||
async for r in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
|
||
assert len(responses) == 1
|
||
function_call = responses[0].content.parts[0].function_call
|
||
assert function_call.name == "my_func"
|
||
assert function_call.id == "call_xyz"
|
||
assert function_call.args == json.loads(full_args)
|
||
|
||
|
||
async def _count_full_buffer_loads(
|
||
lite_llm_instance, mock_completion, full_args, fragments
|
||
):
|
||
mock_completion.return_value = iter(
|
||
_stream_chunks_from_function_chunks(_function_chunks_for_args(fragments))
|
||
)
|
||
real_loads = json.loads
|
||
with patch(
|
||
"google.adk.models.lite_llm.json.loads", side_effect=real_loads
|
||
) as patched_loads:
|
||
responses = [
|
||
r
|
||
async for r in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
full_buffer_calls = [
|
||
c
|
||
for c in patched_loads.call_args_list
|
||
if c.args and c.args[0] == full_args
|
||
]
|
||
return responses, len(full_buffer_calls)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_call_json_loads_count_independent_of_fragment_count(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
# The previous implementation called json.loads(buffer) after every
|
||
# fragment, so the count grew with the fragment count (O(N) calls and
|
||
# O(N^2) total parse cost). The fix makes the count constant.
|
||
full_args = '{"a": 1, "b": {"c": 2}}'
|
||
one_chunk = [full_args]
|
||
one_char_at_a_time = _split_into_chunks(full_args, [1] * len(full_args))
|
||
|
||
_, count_one_chunk = await _count_full_buffer_loads(
|
||
lite_llm_instance, mock_completion, full_args, one_chunk
|
||
)
|
||
_, count_many_chunks = await _count_full_buffer_loads(
|
||
lite_llm_instance, mock_completion, full_args, one_char_at_a_time
|
||
)
|
||
assert count_one_chunk == count_many_chunks
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete(
|
||
mock_completion, lite_llm_instance
|
||
):
|
||
# The closing brace inside the string must not advance fallback_index.
|
||
# If it did, the second tool call would be merged into a single bucket
|
||
# and the assembled args would be invalid.
|
||
full_args_a = '{"text": "a{b}c"}'
|
||
full_args_b = '{"x": 1}'
|
||
fragments_a = _split_into_chunks(full_args_a, [1] * len(full_args_a))
|
||
fragments_b = _split_into_chunks(full_args_b, [1] * len(full_args_b))
|
||
|
||
function_chunks = _function_chunks_for_args(fragments_a)
|
||
# Second tool call: provider emits no index, relies on fallback_index advance.
|
||
for i, fragment in enumerate(fragments_b):
|
||
function_chunks.append(
|
||
FunctionChunk(
|
||
id="call_2" if i == 0 else None,
|
||
name="other_func" if i == 0 else None,
|
||
args=fragment,
|
||
index=0,
|
||
)
|
||
)
|
||
|
||
mock_completion.return_value = iter(
|
||
_stream_chunks_from_function_chunks(function_chunks)
|
||
)
|
||
|
||
responses = [
|
||
r
|
||
async for r in lite_llm_instance.generate_content_async(
|
||
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||
)
|
||
]
|
||
|
||
assert len(responses) == 1
|
||
parts = responses[0].content.parts
|
||
assert len(parts) == 2
|
||
args_by_name = {p.function_call.name: p.function_call.args for p in parts}
|
||
assert args_by_name["my_func"] == json.loads(full_args_a)
|
||
assert args_by_name["other_func"] == json.loads(full_args_b)
|