mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
George WealeandCopybara-Service
47221cd5c1
For providers that typically require file IDs (like OpenAI and Azure), if a file URI is an HTTP/HTTPS URL and the MIME type is image, video, or audio, convert it to the corresponding URL-based content type (e.g., "image_url") instead of using the generic "file" type Close #4112 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 863311703
3892 lines
120 KiB
Python
3892 lines
120 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 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 Mock
|
|
import warnings
|
|
|
|
from google.adk.models.lite_llm import _content_to_message_param
|
|
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 _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_tool_calls_from_text
|
|
from google.adk.models.lite_llm import _redirect_litellm_loggers_to_stdout
|
|
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 _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 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 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 = [
|
|
ModelResponse(
|
|
model="test_model",
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
content="zero, ",
|
|
),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
model="test_model",
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
content="one, ",
|
|
),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
model="test_model",
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
content="two:",
|
|
),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
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.0-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.0-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.0-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.0-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.0-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.0-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.0-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
|
|
|
|
|
|
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.0-flash",
|
|
config=types.GenerateContentConfig(response_schema=_StructuredOutput),
|
|
)
|
|
|
|
_, _, response_format, _ = await _get_completion_inputs(
|
|
llm_request, model="gemini/gemini-2.0-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.0-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",
|
|
]
|
|
|
|
|
|
MULTIPLE_FUNCTION_CALLS_STREAM = [
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ChatCompletionDeltaToolCall(
|
|
type="function",
|
|
id=None,
|
|
function=Function(
|
|
name=None,
|
|
arguments='ue1"}',
|
|
),
|
|
index=0,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ChatCompletionDeltaToolCall(
|
|
type="function",
|
|
id=None,
|
|
function=Function(
|
|
name=None,
|
|
arguments='ue2"}',
|
|
),
|
|
index=1,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason="tool_calls",
|
|
)
|
|
]
|
|
),
|
|
]
|
|
|
|
|
|
STREAM_WITH_EMPTY_CHUNK = [
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
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.
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ChatCompletionDeltaToolCall(
|
|
type="function",
|
|
id=None,
|
|
function=Function(
|
|
name=None,
|
|
arguments="",
|
|
),
|
|
index=0,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
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 = [
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ChatCompletionDeltaToolCall(
|
|
type="function",
|
|
id=None,
|
|
function=Function(
|
|
name=None,
|
|
arguments='ue1"}',
|
|
),
|
|
index=0,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
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,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(
|
|
role="assistant",
|
|
tool_calls=[
|
|
ChatCompletionDeltaToolCall(
|
|
type="function",
|
|
id=None,
|
|
function=Function(
|
|
name=None,
|
|
arguments='ue2"}',
|
|
),
|
|
index=0,
|
|
)
|
|
],
|
|
),
|
|
)
|
|
]
|
|
),
|
|
ModelResponse(
|
|
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."
|
|
)
|
|
|
|
|
|
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,
|
|
},
|
|
)
|
|
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
|
|
|
|
mock_acompletion.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_async_ollama_chat_flattens_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"]
|
|
assert isinstance(message_content, str)
|
|
assert "Describe this image." in message_content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_async_custom_provider_flattens_content(
|
|
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"]
|
|
assert isinstance(message_content, str)
|
|
assert "Describe this image." in message_content
|
|
|
|
|
|
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",
|
|
),
|
|
(
|
|
[
|
|
{"type": "text", "text": "Describe this image."},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "http://example.com"},
|
|
},
|
|
],
|
|
"Describe this image.",
|
|
),
|
|
],
|
|
)
|
|
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_serializes_non_text_blocks_to_json():
|
|
from google.adk.models.lite_llm import _flatten_ollama_content
|
|
|
|
blocks = [
|
|
{"type": "image_url", "image_url": {"url": "http://example.com"}},
|
|
]
|
|
flattened = _flatten_ollama_content(blocks)
|
|
assert isinstance(flattened, str)
|
|
assert json.loads(flattened) == 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 handling of file_data without mime_type (GcsArtifactService scenario).
|
|
|
|
When using GcsArtifactService, artifacts may have file_uri (gs://...) but
|
|
without mime_type set. LiteLLM's Vertex AI backend requires the format
|
|
field to be present, so we infer MIME type from the URI extension or use
|
|
a default fallback to ensure compatibility.
|
|
|
|
See: https://github.com/google/adk-python/issues/3787
|
|
"""
|
|
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,
|
|
],
|
|
)
|
|
|
|
message = await _content_to_message_param(content)
|
|
assert message == {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "Analyze this file."},
|
|
{
|
|
"type": "file",
|
|
"file": {
|
|
"file_id": (
|
|
"gs://agent-artifact-bucket/app/user/session/artifact/0"
|
|
),
|
|
"format": "application/octet-stream",
|
|
},
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
@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.
|
|
|
|
See: https://github.com/google/adk-python/issues/3787
|
|
"""
|
|
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_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_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_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_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'
|
|
|
|
|
|
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_falls_back_to_text(
|
|
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",
|
|
)
|
|
)
|
|
]
|
|
content = await _get_content(parts, provider=provider, model=model)
|
|
assert content == [
|
|
{"type": "text", "text": '[File reference: "document.pdf"]'}
|
|
]
|
|
|
|
|
|
@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",
|
|
),
|
|
pytest.param(
|
|
"https://example.com/audio.mp3",
|
|
"audio/mpeg",
|
|
"audio_url",
|
|
id="audio",
|
|
),
|
|
],
|
|
)
|
|
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
|
|
@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_falls_back_to_text(
|
|
provider, model
|
|
):
|
|
file_uri = "https://example.com/document.pdf"
|
|
parts = [
|
|
types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=file_uri,
|
|
mime_type="application/pdf",
|
|
display_name="document.pdf",
|
|
)
|
|
)
|
|
]
|
|
content = await _get_content(parts, provider=provider, model=model)
|
|
assert content == [
|
|
{"type": "text", "text": '[File reference: "document.pdf"]'}
|
|
]
|
|
|
|
|
|
@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_falls_back_to_text():
|
|
parts = [
|
|
types.Part(
|
|
file_data=types.FileData(
|
|
file_uri="gs://bucket/path/to/document.pdf",
|
|
mime_type="application/pdf",
|
|
display_name="document.pdf",
|
|
)
|
|
)
|
|
]
|
|
content = await _get_content(
|
|
parts, provider="anthropic", model="anthropic/claude-3-5"
|
|
)
|
|
assert content == [
|
|
{"type": "text", "text": '[File reference: "document.pdf"]'}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_content_file_uri_anthropic_openai_file_id_falls_back_to_text():
|
|
parts = [types.Part(file_data=types.FileData(file_uri="file-abc123"))]
|
|
content = await _get_content(
|
|
parts, provider="anthropic", model="anthropic/claude-3-5"
|
|
)
|
|
assert content == [
|
|
{"type": "text", "text": '[File reference: "file-abc123"]'}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_content_file_uri_vertex_ai_non_gemini_falls_back_to_text():
|
|
parts = [
|
|
types.Part(
|
|
file_data=types.FileData(
|
|
file_uri="gs://bucket/path/to/document.pdf",
|
|
mime_type="application/pdf",
|
|
display_name="document.pdf",
|
|
)
|
|
)
|
|
]
|
|
content = await _get_content(
|
|
parts, provider="vertex_ai", model="vertex_ai/claude-3-5"
|
|
)
|
|
assert content == [
|
|
{"type": "text", "text": '[File reference: "document.pdf"]'}
|
|
]
|
|
|
|
|
|
@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.
|
|
|
|
See: https://github.com/google/adk-python/issues/3787
|
|
"""
|
|
# 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 extension uses default MIME type.
|
|
|
|
When file_data has a file_uri without a recognizable extension and no explicit
|
|
mime_type, a default MIME type should be used to ensure compatibility with
|
|
LiteLLM backends.
|
|
|
|
See: https://github.com/google/adk-python/issues/3787
|
|
"""
|
|
# Use Part constructor directly to create file_data without mime_type
|
|
# (types.Part.from_uri requires a valid mime_type when it can't infer)
|
|
parts = [
|
|
types.Part(file_data=types.FileData(file_uri="gs://bucket/artifact/0"))
|
|
]
|
|
content = await _get_content(parts)
|
|
assert content[0] == {
|
|
"type": "file",
|
|
"file": {
|
|
"file_id": "gs://bucket/artifact/0",
|
|
"format": "application/octet-stream",
|
|
},
|
|
}
|
|
|
|
|
|
@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
|
|
async def test_get_content_audio():
|
|
parts = [
|
|
types.Part.from_bytes(data=b"test_audio_data", mime_type="audio/mpeg")
|
|
]
|
|
content = await _get_content(parts)
|
|
assert content[0]["type"] == "audio_url"
|
|
assert (
|
|
content[0]["audio_url"]["url"]
|
|
== "data:audio/mpeg;base64,dGVzdF9hdWRpb19kYXRh"
|
|
)
|
|
assert "format" not in content[0]["audio_url"]
|
|
|
|
|
|
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",
|
|
}
|
|
}
|
|
]
|
|
),
|
|
[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",
|
|
),
|
|
(
|
|
ModelResponse(
|
|
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')],
|
|
UsageMetadataChunk(
|
|
prompt_tokens=0, completion_tokens=0, total_tokens=0
|
|
),
|
|
None,
|
|
),
|
|
(
|
|
ModelResponse(choices=[{"finish_reason": "tool_calls"}]),
|
|
[None],
|
|
UsageMetadataChunk(
|
|
prompt_tokens=0, completion_tokens=0, total_tokens=0
|
|
),
|
|
"tool_calls",
|
|
),
|
|
(
|
|
ModelResponse(choices=[{}]),
|
|
[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",
|
|
),
|
|
],
|
|
)
|
|
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
|
|
assert finished == expected_finished
|
|
|
|
if expected_usage_chunk is None:
|
|
assert usage_chunk is None
|
|
else:
|
|
assert usage_chunk is not None
|
|
assert usage_chunk == expected_usage_chunk
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acompletion_additional_args(mock_acompletion, 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=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"] == "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"
|
|
|
|
|
|
@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(
|
|
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([
|
|
ModelResponse(
|
|
model="test_model",
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(role="assistant", content="Hello "),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
model="test_model",
|
|
choices=[
|
|
StreamingChoices(
|
|
finish_reason=None,
|
|
delta=Delta(role="assistant", content="world"),
|
|
)
|
|
],
|
|
),
|
|
ModelResponse(
|
|
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,
|
|
ModelResponse(
|
|
usage={
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
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
|
|
|
|
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,
|
|
ModelResponse(
|
|
usage={
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 15,
|
|
"cached_tokens": 8,
|
|
},
|
|
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
|
|
|
|
|
|
@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_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-1.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-1.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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finish_reason_unknown_maps_to_other(
|
|
mock_acompletion, lite_llm_instance
|
|
):
|
|
"""Test that unknown finish_reason values map to FinishReason.OTHER."""
|
|
mock_response = ModelResponse(
|
|
choices=[
|
|
Choices(
|
|
message=ChatCompletionAssistantMessage(
|
|
role="assistant",
|
|
content="Test response",
|
|
),
|
|
finish_reason="unknown_reason_type",
|
|
)
|
|
]
|
|
)
|
|
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)
|