github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 7 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

refactor(types): make google.adk.models pass strict mypy

Browse Source 
Not annotations-only. This is one component's slice of a repo-wide typing
cleanup, and the wider change was found to contain behavior changes that have
not all been individually triaged, so please review it as a functional change.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 958623646
This commit is contained in:
George Weale
2026-08-03 15:57:30 -07:00
committed by Copybara-Service
parent da46f8a4b0
commit 7f82142adb
15 changed files with 843 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/models/anthropic_llm.py
+106 -59
View File
@@ -26,10 +26,13 @@ import os
import re
from typing import Any
from typing import AsyncGenerator
from typing import cast
from typing import get_args
from typing import Iterable
from typing import Literal
from typing import Optional
from typing import TYPE_CHECKING
from typing import TypeAlias
from typing import Union
import warnings
@@ -57,6 +60,24 @@ __all__ = ["AnthropicLlm", "Claude", "AnthropicGenerateContentConfig"]
logger = logging.getLogger("google_adk." + __name__)
_ImageMediaType: TypeAlias = Literal[
"image/jpeg",
"image/png",
"image/gif",
"image/webp",
]
_ANTHROPIC_IMAGE_MEDIA_TYPES = frozenset[str](get_args(_ImageMediaType))
_MessageBlockParam: TypeAlias = Union[
anthropic_types.TextBlockParam,
anthropic_types.ThinkingBlockParam,
anthropic_types.RedactedThinkingBlockParam,
anthropic_types.ImageBlockParam,
anthropic_types.DocumentBlockParam,
anthropic_types.ToolUseBlockParam,
anthropic_types.ToolResultBlockParam,
]
_RATE_LIMIT_POSSIBLE_FIX_MESSAGE = (
"On how to mitigate this issue, please refer to:\n\n"
@@ -277,21 +298,30 @@ def to_google_genai_finish_reason(
def _is_image_part(part: types.Part) -> bool:
return (
part.inline_data
and part.inline_data.mime_type
and part.inline_data.mime_type.startswith("image")
inline_data = part.inline_data
return bool(
inline_data is not None
and inline_data.mime_type is not None
and inline_data.mime_type.startswith("image/")
)
def _is_pdf_part(part: types.Part) -> bool:
return (
part.inline_data
and part.inline_data.mime_type
and part.inline_data.mime_type.split(";")[0].strip() == "application/pdf"
inline_data = part.inline_data
return bool(
inline_data is not None
and inline_data.mime_type is not None
and inline_data.mime_type.split(";", 1)[0].strip() == "application/pdf"
)
def _normalize_image_media_type(mime_type: str) -> _ImageMediaType:
normalized = mime_type.split(";", 1)[0].strip().lower()
if normalized not in _ANTHROPIC_IMAGE_MEDIA_TYPES:
raise ValueError(f"Unsupported Anthropic image MIME type: {mime_type}")
return cast(_ImageMediaType, normalized)
class _ToolUseIdSanitizer:
"""Maps invalid tool_use IDs to deterministic fallbacks.
@@ -316,14 +346,7 @@ class _ToolUseIdSanitizer:
def _part_to_message_block(
part: types.Part,
sanitizer: _ToolUseIdSanitizer,
) -> Union[
anthropic_types.TextBlockParam,
anthropic_types.ThinkingBlockParam,
anthropic_types.ImageBlockParam,
anthropic_types.DocumentBlockParam,
anthropic_types.ToolUseBlockParam,
anthropic_types.ToolResultBlockParam,
]:
) -> _MessageBlockParam:
if part.thought and part.text:
signature = ""
if part.thought_signature:
@@ -343,17 +366,20 @@ def _part_to_message_block(
if part.text:
return anthropic_types.TextBlockParam(text=part.text, type="text")
elif part.function_call:
assert part.function_call.name
function_call = part.function_call
assert function_call.name
tool_input: dict[str, object] = dict(function_call.args or {})
return anthropic_types.ToolUseBlockParam(
id=sanitizer.sanitize(part.function_call.id),
name=part.function_call.name,
input=part.function_call.args,
id=sanitizer.sanitize(function_call.id),
name=function_call.name,
input=tool_input,
type="tool_use",
)
elif part.function_response:
function_response = part.function_response
content = ""
response_data = part.function_response.response
response_data = function_response.response or {}
if (
"content" in response_data
@@ -393,36 +419,52 @@ def _part_to_message_block(
content = json.dumps(response_data)
return anthropic_types.ToolResultBlockParam(
tool_use_id=sanitizer.sanitize(part.function_response.id),
tool_use_id=sanitizer.sanitize(function_response.id),
type="tool_result",
content=content,
is_error=False,
)
elif _is_image_part(part):
data = base64.b64encode(part.inline_data.data).decode()
inline_data = part.inline_data
if (
inline_data is None
or inline_data.data is None
or inline_data.mime_type is None
):
raise ValueError("Anthropic image parts require MIME type and data")
data = base64.b64encode(inline_data.data).decode()
image_source = anthropic_types.Base64ImageSourceParam(
type="base64",
media_type=_normalize_image_media_type(inline_data.mime_type),
data=data,
)
return anthropic_types.ImageBlockParam(
type="image",
source=dict(
type="base64", media_type=part.inline_data.mime_type, data=data
),
source=image_source,
)
elif _is_pdf_part(part):
data = base64.b64encode(part.inline_data.data).decode()
inline_data = part.inline_data
if inline_data is None or inline_data.data is None:
raise ValueError("Anthropic PDF parts require data")
data = base64.b64encode(inline_data.data).decode()
pdf_source = anthropic_types.Base64PDFSourceParam(
type="base64",
media_type="application/pdf",
data=data,
)
return anthropic_types.DocumentBlockParam(
type="document",
source=dict(
type="base64", media_type=part.inline_data.mime_type, data=data
),
source=pdf_source,
)
elif part.executable_code:
return anthropic_types.TextBlockParam(
type="text",
text="Code:```python\n" + part.executable_code.code + "\n```",
text="Code:```python\n" + (part.executable_code.code or "") + "\n```",
)
elif part.code_execution_result:
return anthropic_types.TextBlockParam(
text="Execution Result:```code_output\n"
+ part.code_execution_result.output
+ (part.code_execution_result.output or "")
+ "\n```",
type="text",
)
@@ -458,13 +500,7 @@ def _content_to_message_param(
def part_to_message_block(
part: types.Part,
) -> Union[
anthropic_types.TextBlockParam,
anthropic_types.ImageBlockParam,
anthropic_types.DocumentBlockParam,
anthropic_types.ToolUseBlockParam,
anthropic_types.ToolResultBlockParam,
]:
) -> _MessageBlockParam:
return _part_to_message_block(part, _ToolUseIdSanitizer())
@@ -497,7 +533,10 @@ def content_block_to_part(
part = types.Part.from_function_call(
name=content_block.name, args=content_block.input
)
part.function_call.id = content_block.id
function_call = part.function_call
if function_call is None:
raise ValueError("Function-call part factory returned no function call")
function_call.id = content_block.id
return part
raise NotImplementedError(
f"Unsupported content block type: {type(content_block)}"
@@ -538,7 +577,7 @@ def message_to_generate_content_response(
)
def _update_type_string(value: Any) -> None:
def _update_type_string(value: object) -> None:
"""Lowercases nested JSON schema type strings for Anthropic compatibility."""
if isinstance(value, list):
for item in value:
@@ -678,14 +717,14 @@ class AnthropicLlm(BaseLlm):
NotGiven,
],
) -> dict[str, Any]:
system = NOT_GIVEN
system: str | NotGiven = NOT_GIVEN
if llm_request.config:
system_str = extract_system_instruction(llm_request.config)
if system_str:
system = system_str
model_to_use = self._resolve_model_name(llm_request.model)
kwargs = {
kwargs: dict[str, Any] = {
"model": model_to_use,
"system": system,
"messages": messages,
@@ -750,15 +789,18 @@ class AnthropicLlm(BaseLlm):
_content_to_message_param(content, sanitizer)
for content in llm_request.contents or []
]
tools = NOT_GIVEN
if (
llm_request.config
and llm_request.config.tools
and llm_request.config.tools[0].function_declarations
):
tools: Iterable[anthropic_types.ToolUnionParam] | NotGiven = NOT_GIVEN
function_declarations: list[types.FunctionDeclaration] = []
if llm_request.config and llm_request.config.tools:
for configured_tool in llm_request.config.tools:
if isinstance(configured_tool, types.Tool):
function_declarations.extend(
configured_tool.function_declarations or []
)
if function_declarations:
tools = [
function_declaration_to_tool_param(tool)
for tool in llm_request.config.tools[0].function_declarations
for tool in function_declarations
]
tool_choice = (
anthropic_types.ToolChoiceAutoParam(type="auto")
@@ -912,10 +954,10 @@ class AnthropicLlm(BaseLlm):
)
for idx in all_indices:
if idx in thinking_blocks:
acc = thinking_blocks[idx]
part = types.Part(text=acc.thinking, thought=True)
if acc.signature:
part.thought_signature = acc.signature.encode("utf-8")
thinking_acc = thinking_blocks[idx]
part = types.Part(text=thinking_acc.thinking, thought=True)
if thinking_acc.signature:
part.thought_signature = thinking_acc.signature.encode("utf-8")
all_parts.append(part)
if idx in redacted_thinking_blocks:
all_parts.append(
@@ -927,10 +969,15 @@ class AnthropicLlm(BaseLlm):
if idx in text_blocks:
all_parts.append(types.Part.from_text(text=text_blocks[idx]))
if idx in tool_use_blocks:
acc = tool_use_blocks[idx]
args = json.loads(acc.args_json) if acc.args_json else {}
part = types.Part.from_function_call(name=acc.name, args=args)
part.function_call.id = acc.id
tool_acc = tool_use_blocks[idx]
args = json.loads(tool_acc.args_json) if tool_acc.args_json else {}
part = types.Part.from_function_call(name=tool_acc.name, args=args)
function_call = part.function_call
if function_call is None:
raise ValueError(
"Function-call part factory returned no function call"
)
function_call.id = tool_acc.id
all_parts.append(part)
yield LlmResponse(
@@ -946,7 +993,7 @@ class AnthropicLlm(BaseLlm):
)
@cached_property
def _anthropic_client(self) -> AsyncAnthropic:
def _anthropic_client(self) -> AsyncAnthropic | AsyncAnthropicVertex:
return AsyncAnthropic()
src/google/adk/models/apigee_llm.py
+63 -32
View File
@@ -93,7 +93,7 @@ class ApigeeLlm(Gemini):
retry_options: Optional[types.HttpRetryOptions] = None,
api_type: ApiType | str = ApiType.UNKNOWN,
credentials: Credentials | None = None,
):
) -> None:
"""Initializes the Apigee LLM backend.
Args:
@@ -147,23 +147,27 @@ class ApigeeLlm(Gemini):
else:
self._api_type = ApigeeLlm.ApiType.GENAI
self._isvertexai = _identify_vertexai(model, self._api_type)
self._project: str | None = None
self._location: str | None = None
# Set the project and location for Vertex AI.
if self._isvertexai:
self._project = os.environ.get(_PROJECT_ENV_VARIABLE_NAME)
self._location = os.environ.get(_LOCATION_ENV_VARIABLE_NAME)
project = os.environ.get(_PROJECT_ENV_VARIABLE_NAME)
location = os.environ.get(_LOCATION_ENV_VARIABLE_NAME)
if not self._project:
if not project:
raise ValueError(
f'The {_PROJECT_ENV_VARIABLE_NAME} environment variable must be'
' set.'
)
if not self._location:
if not location:
raise ValueError(
f'The {_LOCATION_ENV_VARIABLE_NAME} environment variable must be'
' set.'
)
self._project = project
self._location = location
self._api_version = _identify_api_version(model)
self._proxy_url = proxy_url or os.environ.get(
@@ -190,11 +194,19 @@ class ApigeeLlm(Gemini):
def _completions_http_client(self) -> CompletionsHTTPClient:
"""Provides the completions HTTP client."""
return CompletionsHTTPClient(
base_url=self._proxy_url,
base_url=self._require_proxy_url(),
headers=self._merge_tracking_headers(self._custom_headers),
retry_options=self.retry_options,
)
def _require_proxy_url(self) -> str:
if not self._proxy_url:
raise ValueError(
'Apigee proxy URL is not set. Pass proxy_url or set '
f'{_APIGEE_PROXY_URL_ENV_VARIABLE_NAME}.'
)
return self._proxy_url
@override
async def generate_content_async(
self, llm_request: LlmRequest, stream: bool = False
@@ -231,17 +243,16 @@ class ApigeeLlm(Gemini):
"""
from google.genai import Client
kwargs_for_http_options = {}
if self._api_version:
kwargs_for_http_options['api_version'] = self._api_version
http_options = types.HttpOptions(
base_url=self._proxy_url,
api_version=self._api_version or None,
base_url=self._require_proxy_url(),
headers=self._merge_tracking_headers(self._custom_headers),
retry_options=self.retry_options,
**kwargs_for_http_options,
)
kwargs_for_client = {}
# Built conditionally: passing project/location/credentials as explicit
# Nones is not equivalent to omitting them.
kwargs_for_client: dict[str, Any] = {}
kwargs_for_client['enterprise'] = self._isvertexai
if self._isvertexai:
kwargs_for_client['project'] = self._project
@@ -299,8 +310,10 @@ def _identify_api_version(model: str) -> str:
return ''
def _get_model_id(model: str) -> str:
def _get_model_id(model: str | None) -> str:
"""Returns the model ID for the model spec."""
if not model:
raise ValueError('Model is not set.')
model = model.removeprefix('apigee/')
components = model.split('/')
@@ -482,7 +495,7 @@ class CompletionsHTTPClient:
retry_network = tenacity.retry_if_exception_type(httpx.NetworkError)
def is_retriable(e: Exception) -> bool:
def is_retriable(e: BaseException) -> bool:
if isinstance(e, httpx.HTTPStatusError):
return e.response.status_code in retriable_codes
return False
@@ -564,6 +577,7 @@ class CompletionsHTTPClient:
)
response.raise_for_status()
return response
raise RuntimeError('HTTP retry loop completed without making an attempt')
async def _handle_streaming(
self,
@@ -603,15 +617,15 @@ class CompletionsHTTPClient:
self, llm_request: LlmRequest, stream: bool
) -> dict[str, Any]:
"""Constructs the payload from the LlmRequest."""
messages = []
messages: list[dict[str, Any]] = []
if llm_request.config and llm_request.config.system_instruction:
content = self._serialize_system_instruction(
system_content = self._serialize_system_instruction(
llm_request.config.system_instruction
)
if content:
if system_content:
messages.append({
'role': 'system',
'content': content,
'content': system_content,
})
for content in llm_request.contents:
@@ -667,8 +681,13 @@ class CompletionsHTTPClient:
) -> None:
"""Maps tools and tool configuration to the payload."""
if config.tools:
tools = []
tools: list[dict[str, Any]] = []
for tool in config.tools:
if not isinstance(tool, types.Tool):
raise TypeError(
'OpenAI-compatible Apigee requests require '
'google.genai.types.Tool values.'
)
if tool.function_declarations:
for func in tool.function_declarations:
tools.append(self._function_declaration_to_tool(func))
@@ -746,11 +765,14 @@ class CompletionsHTTPClient:
return
if part.function_call:
function_name = part.function_call.name
if not function_name:
raise ValueError('Function calls must include a name.')
tool_call = {
'id': part.function_call.id or 'call_' + part.function_call.name,
'id': part.function_call.id or f'call_{function_name}',
'type': 'function',
'function': {
'name': part.function_call.name,
'name': function_name,
'arguments': (
json.dumps(part.function_call.args)
if part.function_call.args
@@ -759,7 +781,7 @@ class CompletionsHTTPClient:
},
}
if part.thought_signature:
sig = part.thought_signature
sig: str | bytes = part.thought_signature
if isinstance(sig, bytes):
sig = base64.b64encode(sig).decode('utf-8')
tool_call['extra_content'] = {
@@ -782,7 +804,12 @@ class CompletionsHTTPClient:
content_parts.append({'type': 'text', 'text': before})
elif part.inline_data:
mime_type = part.inline_data.mime_type
data = base64.b64encode(part.inline_data.data).decode('utf-8')
if not mime_type:
raise ValueError('Inline data must include a MIME type.')
inline_data = part.inline_data.data
if inline_data is None:
raise ValueError('Inline data must include data.')
data = base64.b64encode(inline_data).decode('utf-8')
url = f'data:{mime_type};base64,{data}'
content_parts.append({'type': 'image_url', 'image_url': {'url': url}})
elif part.file_data:
@@ -833,7 +860,7 @@ class CompletionsHTTPClient:
return system_instruction.text
if isinstance(system_instruction, types.Content):
return ''.join(
part.text for part in system_instruction.parts if part.text
part.text for part in system_instruction.parts or [] if part.text
)
if isinstance(system_instruction, dict):
part = types.Part(**system_instruction)
@@ -1174,6 +1201,10 @@ class ChatCompletionsResponseHandler:
)
part = self.tool_call_parts[index]
chunk_part = types.Part(function_call=types.FunctionCall())
function_call = part.function_call
chunk_function_call = chunk_part.function_call
if function_call is None or chunk_function_call is None:
raise RuntimeError('Tool-call parts must contain a function call.')
call_type = tool_call.get('type')
# TODO: Add support for 'custom' type.
if call_type is not None and call_type != 'function':
@@ -1185,22 +1216,22 @@ class ChatCompletionsResponseHandler:
if args_delta:
try:
args = json.loads(args_delta)
chunk_part.function_call.args = args
if not part.function_call.args:
part.function_call.args = dict(args)
chunk_function_call.args = args
if not function_call.args:
function_call.args = dict(args)
else:
part.function_call.args.update(args)
function_call.args.update(args)
except json.JSONDecodeError as e:
raise ValueError(f'Failed to parse arguments: {args_delta}') from e
func_name = func.get('name')
if func_name:
part.function_call.name = func_name
chunk_part.function_call.name = func_name
function_call.name = func_name
chunk_function_call.name = func_name
tool_call_id = tool_call.get('id')
if tool_call_id:
part.function_call.id = tool_call_id
chunk_part.function_call.id = tool_call_id
function_call.id = tool_call_id
chunk_function_call.id = tool_call_id
# Add support for gemini's thought_signature.
thought_signature = (
src/google/adk/models/base_llm.py
+4 -1
View File
@@ -15,6 +15,7 @@
from __future__ import annotations
from abc import abstractmethod
from contextlib import AbstractAsyncContextManager
from typing import AsyncGenerator
from typing import TYPE_CHECKING
import warnings
@@ -272,7 +273,9 @@ class BaseLlm(BaseModel):
)
)
def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
def connect(
self, llm_request: LlmRequest
) -> AbstractAsyncContextManager[BaseLlmConnection]:
"""Creates a live connection to the LLM.
Args:
src/google/adk/models/base_llm_connection.py
+2 -2
View File
@@ -87,8 +87,8 @@ class BaseLlmConnection:
Yields:
LlmResponse: The model response.
"""
# We need to yield here to help type checkers infer the correct type.
yield
# A value-bearing yield keeps this abstract method an async generator.
yield LlmResponse()
@abstractmethod
async def close(self) -> None:
src/google/adk/models/cache_metadata.py
+5 -5
View File
@@ -15,7 +15,6 @@
from __future__ import annotations
import time
from typing import Optional
from pydantic import BaseModel
from pydantic import ConfigDict
@@ -58,14 +57,14 @@ class CacheMetadata(BaseModel):
frozen=True, # Cache metadata should be immutable
)
cache_name: Optional[str] = Field(
cache_name: str | None = Field(
default=None,
description=(
"Full resource name of the cached content (None if no active cache)"
),
)
expire_time: Optional[float] = Field(
expire_time: float | None = Field(
default=None,
description="Unix timestamp when cache expires (None if no active cache)",
)
@@ -74,7 +73,7 @@ class CacheMetadata(BaseModel):
description="Hash of cacheable contents used to detect changes"
)
invocations_used: Optional[int] = Field(
invocations_used: int | None = Field(
default=None,
ge=0,
description=(
@@ -91,7 +90,7 @@ class CacheMetadata(BaseModel):
),
)
created_at: Optional[float] = Field(
created_at: float | None = Field(
default=None,
description=(
"Unix timestamp when cache was created (None if no active cache)"
@@ -123,6 +122,7 @@ class CacheMetadata(BaseModel):
f"Fingerprint-only: {self.contents_count} contents, "
f"fingerprint={self.fingerprint[:8]}..."
)
assert self.expire_time is not None and self.invocations_used is not None
cache_id = self.cache_name.split("/")[-1]
time_until_expiry_minutes = (self.expire_time - time.time()) / 60
return (
src/google/adk/models/gemini_context_cache_manager.py
+77 -35
View File
@@ -22,11 +22,13 @@ import json
import logging
import time
from typing import Any
from typing import cast
from typing import Optional
from typing import TYPE_CHECKING
from google.genai import types
from ..agents.context_cache_config import ContextCacheConfig
from ..utils.feature_decorator import experimental
from .cache_metadata import CacheMetadata
from .llm_request import LlmRequest
@@ -53,6 +55,31 @@ def _minimum_cache_tokens(model: Optional[str]) -> Optional[int]:
return None
def _require_cache_config(llm_request: LlmRequest) -> ContextCacheConfig:
cache_config = llm_request.cache_config
if cache_config is None:
raise ValueError("Context caching requires a cache configuration.")
return cache_config
def _require_model(llm_request: LlmRequest) -> str:
model = llm_request.model
if model is None:
raise ValueError("Context caching requires a model name.")
return model
def _content_union_character_count(value: types.ContentUnion) -> int:
"""Returns a stable rough size for a system-instruction value."""
if isinstance(value, str):
return len(value)
if isinstance(value, list):
return sum(
len(item) if isinstance(item, str) else len(str(item)) for item in value
)
return len(str(value))
@experimental
class GeminiContextCacheManager:
"""Manages context cache lifecycle for Gemini models.
@@ -86,6 +113,9 @@ class GeminiContextCacheManager:
Returns:
Cache metadata to be included in response, or None if caching failed
"""
_require_model(llm_request)
_require_cache_config(llm_request)
# Check if we have existing cache metadata and if it's valid
if llm_request.cache_metadata:
logger.debug(
@@ -99,6 +129,8 @@ class GeminiContextCacheManager:
llm_request.cache_metadata.cache_name,
)
cache_name = llm_request.cache_metadata.cache_name
if cache_name is None:
raise RuntimeError("A valid cache must have active metadata.")
cache_contents_count = llm_request.cache_metadata.contents_count
self._apply_cache_to_request(
llm_request, cache_name, cache_contents_count
@@ -141,8 +173,11 @@ class GeminiContextCacheManager:
llm_request, cache_contents_count
)
if cache_metadata:
cache_name = cache_metadata.cache_name
if cache_name is None:
raise RuntimeError("A newly created cache must be active.")
self._apply_cache_to_request(
llm_request, cache_metadata.cache_name, cache_contents_count
llm_request, cache_name, cache_contents_count
)
return cache_metadata
@@ -239,25 +274,26 @@ class GeminiContextCacheManager:
if not cache_metadata:
return False
# Fingerprint-only metadata is not a valid active cache
if cache_metadata.cache_name is None:
# Fingerprint-only metadata is not a valid active cache.
cache_name = cache_metadata.cache_name
expire_time = cache_metadata.expire_time
invocations_used = cache_metadata.invocations_used
if cache_name is None or expire_time is None or invocations_used is None:
return False
cache_config = _require_cache_config(llm_request)
# Check if cache has expired
if time.time() >= cache_metadata.expire_time:
logger.info("Cache expired: %s", cache_metadata.cache_name)
if time.time() >= expire_time:
logger.info("Cache expired: %s", cache_name)
return False
# Check if cache has been used for too many invocations
if (
cache_metadata.invocations_used
> llm_request.cache_config.cache_intervals
):
if invocations_used > cache_config.cache_intervals:
logger.info(
"Cache exceeded cache intervals: %s (%d > %d intervals)",
cache_metadata.cache_name,
cache_metadata.invocations_used,
llm_request.cache_config.cache_intervals,
cache_name,
invocations_used,
cache_config.cache_intervals,
)
return False
@@ -359,6 +395,8 @@ class GeminiContextCacheManager:
Returns:
Cache metadata if successful, None otherwise
"""
cache_config = _require_cache_config(llm_request)
# Check if we have token count from previous response for cache size validation
if llm_request.cacheable_contents_token_count is None:
logger.info(
@@ -367,14 +405,11 @@ class GeminiContextCacheManager:
)
return None
if (
llm_request.cacheable_contents_token_count
< llm_request.cache_config.min_tokens
):
if llm_request.cacheable_contents_token_count < cache_config.min_tokens:
logger.info(
"Previous request too small for caching (%d < %d tokens)",
llm_request.cacheable_contents_token_count,
llm_request.cache_config.min_tokens,
cache_config.min_tokens,
)
return None
@@ -447,7 +482,9 @@ class GeminiContextCacheManager:
# System instruction
if llm_request.config and llm_request.config.system_instruction:
total_chars += len(llm_request.config.system_instruction)
total_chars += _content_union_character_count(
llm_request.config.system_instruction
)
# Tools
if llm_request.config and llm_request.config.tools:
@@ -461,7 +498,7 @@ class GeminiContextCacheManager:
if cache_contents_count is not None:
contents = contents[:cache_contents_count]
for content in contents:
for part in content.parts:
for part in content.parts or []:
if part.text:
total_chars += len(part.text)
@@ -519,12 +556,15 @@ class GeminiContextCacheManager:
from ..telemetry.tracing import tracer
with tracer.start_as_current_span("create_cache") as span:
cache_request_config = _require_cache_config(llm_request)
model = _require_model(llm_request)
# Prepare cache contents (first N contents + system instruction + tools)
cache_contents = llm_request.contents[:cache_contents_count] or None
cache_config = types.CreateCachedContentConfig(
contents=cache_contents,
ttl=llm_request.cache_config.ttl_string,
ttl=cache_request_config.ttl_string,
display_name=(
f"adk-cache-{int(time.time())}-{cache_contents_count}contents"
),
@@ -535,35 +575,34 @@ class GeminiContextCacheManager:
cache_config.system_instruction = llm_request.config.system_instruction
logger.debug(
"Added system instruction to cache config (length=%d)",
len(llm_request.config.system_instruction),
_content_union_character_count(
llm_request.config.system_instruction
),
)
# Add tools if present
if llm_request.config and llm_request.config.tools:
cache_config.tools = llm_request.config.tools
cache_config.tools = cast(list[types.Tool], llm_request.config.tools)
# Add tool config if present
if llm_request.config and llm_request.config.tool_config:
cache_config.tool_config = llm_request.config.tool_config
# Pass through HTTP options (e.g. timeout) from cache config
if (
llm_request.cache_config
and llm_request.cache_config.create_http_options
):
cache_config.http_options = llm_request.cache_config.create_http_options
if cache_request_config.create_http_options:
cache_config.http_options = cache_request_config.create_http_options
span.set_attribute("cache_contents_count", cache_contents_count)
span.set_attribute("model", llm_request.model)
span.set_attribute("ttl_seconds", llm_request.cache_config.ttl_seconds)
span.set_attribute("model", model)
span.set_attribute("ttl_seconds", cache_request_config.ttl_seconds)
logger.debug(
"Creating cache with model %s and config: %s",
llm_request.model,
model,
cache_config,
)
cached_content = await self.genai_client.aio.caches.create(
model=llm_request.model,
model=model,
config=cache_config,
)
# Set precise creation timestamp right after cache creation
@@ -572,15 +611,18 @@ class GeminiContextCacheManager:
expire_time = (
server_expire_time.timestamp()
if isinstance(server_expire_time, datetime)
else created_at + llm_request.cache_config.ttl_seconds
else created_at + cache_request_config.ttl_seconds
)
logger.info("Cache created successfully: %s", cached_content.name)
cache_name = cached_content.name
if not cache_name:
raise RuntimeError("The cache service returned no cache name.")
logger.info("Cache created successfully: %s", cache_name)
span.set_attribute("cache_name", cached_content.name)
span.set_attribute("cache_name", cache_name)
# Return complete cache metadata with precise timing
return CacheMetadata(
cache_name=cached_content.name,
cache_name=cache_name,
expire_time=expire_time,
fingerprint=self._generate_cache_fingerprint(
llm_request, cache_contents_count
src/google/adk/models/gemini_llm_connection.py
+21 -5
View File
@@ -16,6 +16,7 @@ from __future__ import annotations
import logging
from typing import AsyncGenerator
from typing import cast
from typing import Union
from google.genai import types
@@ -92,8 +93,9 @@ class GeminiLlmConnection(BaseLlmConnection):
if contents:
logger.debug('Sending history to live connection: %s', contents)
turns: list[types.Content | types.ContentDict] = [*contents]
await self._gemini_session.send_client_content(
turns=contents,
turns=turns,
turn_complete=contents[-1].role == 'user',
)
else:
@@ -124,7 +126,16 @@ class GeminiLlmConnection(BaseLlmConnection):
assert content.parts
if content.parts[0].function_response:
# All parts have to be function responses.
function_responses = [part.function_response for part in content.parts]
function_responses = [
function_response
for part in content.parts
if (function_response := part.function_response) is not None
]
if len(function_responses) != len(content.parts):
raise ValueError(
'Function-response content cannot mix function and non-function'
' parts.'
)
logger.debug('Sending LLM function response: %s', function_responses)
await self._gemini_session.send_tool_response(
function_responses=function_responses
@@ -307,7 +318,12 @@ class GeminiLlmConnection(BaseLlmConnection):
tool_call_parts: list[types.Part] = []
last_grounding_metadata = None
tool_call_metadata = None
async with Aclosing(self._gemini_session.receive()) as agen:
async with Aclosing(
cast(
AsyncGenerator[types.LiveServerMessage, None],
self._gemini_session.receive(),
)
) as agen:
# Pending cleanup: reuse StreamingResponseAggregator to accumulate
# partial content and emit responses as needed, once that aggregator
# handles the live-connection message shapes.
@@ -510,7 +526,7 @@ class GeminiLlmConnection(BaseLlmConnection):
text,
is_thought,
last_grounding_metadata,
message.server_content.interrupted,
bool(message.server_content.interrupted),
)
text = ''
is_thought = False
@@ -582,7 +598,7 @@ class GeminiLlmConnection(BaseLlmConnection):
last_grounding_metadata = None
tool_call_parts.extend([
types.Part(function_call=function_call)
for function_call in message.tool_call.function_calls
for function_call in message.tool_call.function_calls or []
])
if not self._is_gemini_3_x_live:
if tool_call_metadata is None:
src/google/adk/models/gemma_llm.py
+18 -9
View File
@@ -20,6 +20,8 @@ import logging
import re
from typing import Any
from typing import AsyncGenerator
from typing import cast
from typing import TYPE_CHECKING
from google.adk.models.google_llm import Gemini
from google.adk.models.llm_request import LlmRequest
@@ -220,7 +222,8 @@ class Gemma(GemmaFunctionCallingMixin, Gemini):
if system_instruction := llm_request.config.system_instruction:
contents = llm_request.contents
instruction_content = Content(
role='user', parts=[Part.from_text(text=system_instruction)]
role='user',
parts=[Part.from_text(text=cast(str, system_instruction))],
)
# NOTE: if history is preserved, we must include the system instructions ONLY once at the beginning
@@ -248,8 +251,9 @@ class Gemma(GemmaFunctionCallingMixin, Gemini):
LlmResponse: The model response.
"""
# print(f'{llm_request=}')
assert llm_request.model.startswith('gemma-'), (
f'Requesting a non-Gemma model ({llm_request.model}) with the Gemma LLM'
model = llm_request.model
assert model is not None and model.startswith('gemma-'), (
f'Requesting a non-Gemma model ({model}) with the Gemma LLM'
' is not supported.'
)
@@ -276,7 +280,7 @@ def _convert_content_parts_for_gemma(
has_function_response_part = False
has_function_call_part = False
for part in content_item.parts:
for part in content_item.parts or []:
if func_response := part.function_response:
has_function_response_part = True
response_text = (
@@ -355,11 +359,16 @@ def _get_last_valid_json_substring(text: str) -> tuple[bool, str | None]:
return False, None
try:
from google.adk.models.lite_llm import LiteLlm # noqa: F401
except ImportError as e:
logger.debug('LiteLlm not available; Gemma3Ollama will not be defined: %s', e)
LiteLlm = None
if TYPE_CHECKING:
from google.adk.models.lite_llm import LiteLlm
else:
try:
from google.adk.models.lite_llm import LiteLlm # noqa: F401
except ImportError as e:
logger.debug(
'LiteLlm not available; Gemma3Ollama will not be defined: %s', e
)
LiteLlm = None
if LiteLlm is not None:
src/google/adk/models/google_llm.py
+32 -20
View File
@@ -23,6 +23,7 @@ import logging
import re
from typing import Any
from typing import AsyncGenerator
from typing import AsyncIterator
from typing import cast
from typing import Optional
from typing import TYPE_CHECKING
@@ -198,6 +199,9 @@ class Gemini(BaseLlm):
"""
await self._preprocess_request(llm_request)
self._maybe_append_user_content(llm_request)
model = llm_request.model
if model is None:
raise ValueError('Gemini requests require a model name.')
# Handle context caching if configured
cache_metadata = None
@@ -230,7 +234,7 @@ class Gemini(BaseLlm):
if not llm_request.config.http_options:
llm_request.config.http_options = types.HttpOptions()
llm_request.config.http_options.headers = self._merge_tracking_headers(
llm_request.config.http_options.headers
llm_request.config.http_options.headers or {}
)
_, api_version = self._base_url_and_api_version
if api_version:
@@ -250,8 +254,8 @@ class Gemini(BaseLlm):
if stream:
responses = await self.api_client.aio.models.generate_content_stream(
model=llm_request.model,
contents=llm_request.contents,
model=model,
contents=cast(list[types.ContentUnion], llm_request.contents),
config=llm_request.config,
)
@@ -274,7 +278,7 @@ class Gemini(BaseLlm):
if (close_result := aggregator.close()) is not None:
# Populate cache metadata in the final aggregated response for
# streaming
if cache_metadata:
if cache_metadata and cache_manager is not None:
cache_manager.populate_cache_metadata_in_response(
close_result, cache_metadata
)
@@ -282,8 +286,8 @@ class Gemini(BaseLlm):
else:
response = await self.api_client.aio.models.generate_content(
model=llm_request.model,
contents=llm_request.contents,
model=model,
contents=cast(list[types.ContentUnion], llm_request.contents),
config=llm_request.config,
)
logger.info('Response received from the model.')
@@ -291,7 +295,7 @@ class Gemini(BaseLlm):
logger.debug(_build_response_log(response))
llm_response = LlmResponse.create(response)
if cache_metadata:
if cache_metadata and cache_manager is not None:
cache_manager.populate_cache_metadata_in_response(
llm_response, cache_metadata
)
@@ -425,7 +429,9 @@ class Gemini(BaseLlm):
return Client(**kwargs)
@contextlib.asynccontextmanager
async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
async def connect(
self, llm_request: LlmRequest
) -> AsyncIterator[BaseLlmConnection]:
"""Connects to the Gemini model and returns an llm connection.
Args:
@@ -455,12 +461,14 @@ class Gemini(BaseLlm):
if self.speech_config is not None:
llm_request.live_connect_config.speech_config = self.speech_config
llm_request.live_connect_config.system_instruction = types.Content(
role='system',
parts=[
types.Part.from_text(text=llm_request.config.system_instruction)
],
)
system_instruction = llm_request.config.system_instruction
if system_instruction is not None:
if not isinstance(system_instruction, str):
raise TypeError('Live Gemini system instructions must be text.')
llm_request.live_connect_config.system_instruction = types.Content(
role='system',
parts=[types.Part.from_text(text=system_instruction)],
)
logger.info(
'Trying to connect to live model: %s with api backend: %s',
@@ -489,13 +497,16 @@ class Gemini(BaseLlm):
)
logger.debug('Connecting to live with llm_request:%s', llm_request)
logger.debug('Live connect config: %s', llm_request.live_connect_config)
model = llm_request.model
if model is None:
raise ValueError('Live Gemini requests require a model name.')
async with self._live_api_client.aio.live.connect(
model=llm_request.model, config=llm_request.live_connect_config
model=model, config=llm_request.live_connect_config
) as live_session:
yield GeminiLlmConnection(
live_session,
api_backend=self._api_backend,
model_version=llm_request.model,
model_version=model,
)
async def _adapt_computer_use_tool(self, llm_request: LlmRequest) -> None:
@@ -595,10 +606,10 @@ def _build_request_log(req: LlmRequest) -> str:
if req.config.tools:
for idx, tool in enumerate(req.config.tools):
if not isinstance(tool, types.Tool):
continue
if tool.function_declarations:
function_decls = cast(
list[types.FunctionDeclaration], tool.function_declarations
)
function_decls = tool.function_declarations
function_decl_tool_index = idx
break
@@ -615,7 +626,8 @@ def _build_request_log(req: LlmRequest) -> str:
exclude_none=True,
exclude={
'parts': {
i: _EXCLUDED_PART_FIELD for i in range(len(content.parts))
i: _EXCLUDED_PART_FIELD
for i in range(len(content.parts or []))
}
},
)
src/google/adk/models/lite_llm.py
+290 -140
View File
@@ -27,6 +27,7 @@ import re
import sys
from typing import Any
from typing import AsyncGenerator
from typing import cast
from typing import Dict
from typing import Generator
from typing import Iterable
@@ -35,6 +36,7 @@ from typing import Literal
from typing import Optional
from typing import Tuple
from typing import TYPE_CHECKING
from typing import TypeAlias
from typing import TypedDict
from typing import Union
from urllib.parse import urlparse
@@ -50,11 +52,15 @@ if not TYPE_CHECKING and importlib.util.find_spec("litellm") is None:
from pydantic import BaseModel
from pydantic import Field
from pydantic import PrivateAttr
from typing_extensions import NotRequired
from typing_extensions import override
from typing_extensions import Required
from ..utils._google_client_headers import merge_tracking_headers
from ._capabilities import LlmCapabilities
from .base_llm import BaseLlm
from .interactions_utils import extract_system_instruction
from .llm_request import LlmRequest
from .llm_response import LlmResponse
@@ -62,14 +68,13 @@ if TYPE_CHECKING:
import litellm
from litellm import acompletion
from litellm import ChatCompletionAssistantMessage
from litellm import ChatCompletionAssistantToolCall
from litellm import ChatCompletionMessageToolCall
from litellm import ChatCompletionSystemMessage
from litellm import ChatCompletionToolCallFunctionChunk
from litellm import ChatCompletionToolMessage
from litellm import ChatCompletionUserMessage
from litellm import completion
from litellm import CustomStreamWrapper
from litellm import Function
from litellm import Message
from litellm import ModelResponse
from litellm import ModelResponseStream
@@ -79,14 +84,13 @@ else:
litellm = None
acompletion = None
ChatCompletionAssistantMessage = None
ChatCompletionAssistantToolCall = None
ChatCompletionMessageToolCall = None
ChatCompletionSystemMessage = None
ChatCompletionToolMessage = None
ChatCompletionUserMessage = None
completion = None
CustomStreamWrapper = None
Function = None
ChatCompletionToolCallFunctionChunk = None
Message = None
ModelResponse = None
Delta = None
@@ -104,7 +108,9 @@ _UNQUOTED_KEY_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
# Mapping of major MIME type prefixes to LiteLLM content types for URL blocks.
# Audio is handled separately as `input_audio` content blocks because LiteLLM
# (and OpenAI) do not accept an `audio_url` content type.
_MEDIA_URL_CONTENT_TYPE_BY_MAJOR_MIME_TYPE = {
_MEDIA_URL_CONTENT_TYPE_BY_MAJOR_MIME_TYPE: dict[
str, Literal["image_url", "video_url"]
] = {
"image": "image_url",
"video": "video_url",
}
@@ -249,13 +255,12 @@ _THOUGHT_SIGNATURE_SEPARATOR = "__thought__"
_LITELLM_IMPORTED = False
_LITELLM_GLOBAL_SYMBOLS = (
"ChatCompletionAssistantMessage",
"ChatCompletionAssistantToolCall",
"ChatCompletionMessageToolCall",
"ChatCompletionSystemMessage",
"ChatCompletionToolMessage",
"ChatCompletionUserMessage",
"CustomStreamWrapper",
"Function",
"ChatCompletionToolCallFunctionChunk",
"Message",
"ModelResponse",
"ModelResponseStream",
@@ -458,7 +463,9 @@ def _normalize_mime_type(mime_type: str) -> str:
return mime_type.split(";", 1)[0].strip().lower()
def _media_url_content_type(mime_type: str) -> str | None:
def _media_url_content_type(
mime_type: str,
) -> Literal["image_url", "video_url"] | None:
"""Returns the LiteLLM URL content type for known media MIME types."""
major_mime_type = _normalize_mime_type(mime_type).split("/", 1)[0]
return _MEDIA_URL_CONTENT_TYPE_BY_MAJOR_MIME_TYPE.get(major_mime_type)
@@ -652,6 +659,130 @@ class ChatCompletionFileUrlObject(TypedDict, total=False):
format: str
class _TextContentObject(TypedDict):
type: Literal["text"]
text: str
class _AudioData(TypedDict):
data: str
format: str
class _AudioContentObject(TypedDict):
type: Literal["input_audio"]
input_audio: _AudioData
class _UrlData(TypedDict):
url: str
class _ImageContentObject(TypedDict):
type: Literal["image_url"]
image_url: _UrlData
class _VideoContentObject(TypedDict):
type: Literal["video_url"]
video_url: _UrlData
class _FileContentObject(TypedDict):
type: Literal["file"]
file: ChatCompletionFileUrlObject
_ContentObject: TypeAlias = Union[
_TextContentObject,
_AudioContentObject,
_ImageContentObject,
_VideoContentObject,
_FileContentObject,
]
_MessageContent: TypeAlias = Union[str, list[_ContentObject]]
class _ThinkingBlock(TypedDict):
type: Required[Literal["thinking"]]
thinking: Required[str]
signature: NotRequired[str]
_AssistantContentObject: TypeAlias = Union[_ContentObject, _ThinkingBlock]
_AssistantContent: TypeAlias = Union[
str, Iterable[_AssistantContentObject], None
]
class _OutboundToolCallFunction(TypedDict):
name: str
arguments: str
class _OutboundToolCall(TypedDict):
type: Required[Literal["function"]]
id: Required[str]
function: Required[_OutboundToolCallFunction]
provider_specific_fields: NotRequired[dict[str, str]]
extra_content: NotRequired[dict[str, dict[str, str]]]
class _AssistantMessagePayload(TypedDict):
role: Required[Literal["assistant"]]
content: Required[_AssistantContent]
tool_calls: NotRequired[list[_OutboundToolCall] | None]
reasoning_content: NotRequired[str | None]
thinking_blocks: NotRequired[list[_ThinkingBlock] | None]
class _GemmaToolMessagePayload(TypedDict):
role: Literal["tool_responses"]
tool_call_id: str
content: str
def _assistant_message(
*,
content: _AssistantContent,
tool_calls: list[_OutboundToolCall] | None = None,
reasoning_content: str | None = None,
thinking_blocks: list[_ThinkingBlock] | None = None,
) -> Message:
"""Build an assistant payload including LiteLLM provider extensions."""
payload = _AssistantMessagePayload(
role="assistant",
content=content,
tool_calls=tool_calls,
reasoning_content=reasoning_content,
)
if thinking_blocks is not None:
payload["thinking_blocks"] = thinking_blocks
# LiteLLM's Message union omits fields accepted by provider adapters.
return cast(Message, payload)
def _tool_message(
*,
role: Literal["tool", "tool_responses"],
tool_call_id: str,
content: str,
) -> Message:
"""Build a standard tool result or Gemma's provider-specific variant."""
if role == "tool":
return ChatCompletionToolMessage(
role="tool",
tool_call_id=tool_call_id,
content=content,
)
payload = _GemmaToolMessagePayload(
role="tool_responses",
tool_call_id=tool_call_id,
content=content,
)
return cast(Message, payload)
class FunctionChunk(BaseModel):
id: Optional[str]
name: Optional[str]
@@ -759,7 +890,7 @@ def _part_has_payload(part: types.Part) -> bool:
return True
if part.inline_data and part.inline_data.data:
return True
if part.file_data and (part.file_data.file_uri or part.file_data.data):
if part.file_data and part.file_data.file_uri:
return True
if part.function_response:
return True
@@ -779,13 +910,12 @@ def _append_fallback_user_content_if_missing(
parts = content.parts or []
if any(_part_has_payload(part) for part in parts):
return
if not parts:
content.parts = []
content.parts.append(
parts.append(
types.Part.from_text(
text="Handle the requests as specified in the System Instruction."
)
)
content.parts = parts
return
llm_request.contents.append(
types.Content(
@@ -1014,9 +1144,11 @@ async def _content_to_message_param(
tool_messages: list[Message] = []
non_tool_parts: list[types.Part] = []
for part in content.parts:
content_parts_or_empty = content.parts or []
for part in content_parts_or_empty:
if part.function_response:
response = part.function_response.response
function_response = part.function_response
response = function_response.response
response_content = (
response
if isinstance(response, str)
@@ -1026,11 +1158,13 @@ async def _content_to_message_param(
# from the tool call, instead of OpenAI-compatible 'tool' role used by other models.
# Earlier Gemma versions before version 4 do not support tool use,
# so this check is intentionally scoped to only look for "gemma4" in the model name.
tool_role = "tool_responses" if _is_gemma4_model(model) else "tool"
tool_role: Literal["tool", "tool_responses"] = (
"tool_responses" if _is_gemma4_model(model) else "tool"
)
tool_messages.append(
ChatCompletionToolMessage(
_tool_message(
role=tool_role,
tool_call_id=part.function_response.id,
tool_call_id=function_response.id or "",
content=response_content,
)
)
@@ -1055,35 +1189,39 @@ async def _content_to_message_param(
role = _to_litellm_role(content.role)
if role == "user":
user_parts = [part for part in content.parts if not part.thought]
user_parts = [part for part in content_parts_or_empty if not part.thought]
message_content = (
await _get_content(user_parts, provider=provider, model=model) or None
)
return ChatCompletionUserMessage(role="user", content=message_content)
return ChatCompletionUserMessage(
role="user",
content=cast(OpenAIMessageContent, message_content),
)
else: # assistant/model
tool_calls = []
tool_calls: list[_OutboundToolCall] = []
content_parts: list[types.Part] = []
reasoning_parts: list[types.Part] = []
for part in content.parts:
for part in content_parts_or_empty:
if part.function_call:
tool_call_id = part.function_call.id or ""
tool_call_dict: ChatCompletionAssistantToolCall = {
"type": "function",
"id": tool_call_id,
"function": {
"name": part.function_call.name,
"arguments": _safe_json_serialize(part.function_call.args),
function_call = part.function_call
if not function_call.name:
raise ValueError("LiteLLM function calls require a name")
tool_call_id = function_call.id or ""
tool_call_dict = _OutboundToolCall(
type="function",
id=tool_call_id,
function={
"name": function_call.name,
"arguments": _safe_json_serialize(function_call.args),
},
}
)
# Preserve thought_signature for Gemini thinking models.
# LiteLLM's Gemini prompt conversion reads provider_specific_fields,
# while the OpenAI-compatible Gemini endpoint path expects the
# extra_content.google.thought_signature payload to survive.
# See https://ai.google.dev/gemini-api/docs/thought-signatures.
if part.thought_signature:
sig = part.thought_signature
if isinstance(sig, bytes):
sig = base64.b64encode(sig).decode("utf-8")
sig = base64.b64encode(part.thought_signature).decode("utf-8")
tool_call_dict["provider_specific_fields"] = {
"thought_signature": sig
}
@@ -1104,11 +1242,9 @@ async def _content_to_message_param(
if final_content and isinstance(final_content, list):
# when the content is a single text object, we can use it directly.
# this is needed for ollama_chat provider which fails if content is a list
final_content = (
final_content[0].get("text", "")
if final_content[0].get("type", None) == "text"
else final_content
)
first_content = final_content[0]
if first_content["type"] == "text":
final_content = first_content["text"]
# For Anthropic models, rebuild thinking_blocks with signatures so that
# thinking is preserved across tool call boundaries. Without this,
@@ -1119,25 +1255,23 @@ async def _content_to_message_param(
# Aggregate them back into one thinking block for outbound.
if model and _is_anthropic_model(model) and reasoning_parts:
aggregated_parts = _aggregate_streaming_thought_parts(reasoning_parts)
thinking_blocks = []
thinking_blocks: list[_ThinkingBlock] = []
for part in aggregated_parts:
if part.text and part.thought_signature:
sig = part.thought_signature
if isinstance(sig, bytes):
sig = base64.b64encode(sig).decode("utf-8")
thinking_blocks.append({
"type": "thinking",
"thinking": part.text,
"signature": sig,
})
signature = base64.b64encode(part.thought_signature).decode("utf-8")
thinking_blocks.append(
_ThinkingBlock(
type="thinking",
thinking=part.text,
signature=signature,
)
)
if thinking_blocks:
msg = ChatCompletionAssistantMessage(
role=role,
return _assistant_message(
content=final_content,
tool_calls=tool_calls or None,
thinking_blocks=thinking_blocks,
)
msg["thinking_blocks"] = thinking_blocks # type: ignore[typeddict-unknown-key]
return msg
# Anthropic routes require thinking blocks to be embedded directly in the
# message content list. LiteLLM's prompt template for Anthropic drops the
@@ -1147,29 +1281,26 @@ async def _content_to_message_param(
# multi-turn conversations. On multi-model platforms (bedrock, vertex_ai)
# this must only apply to actual Claude models, not Gemini/Llama/etc.
if reasoning_parts and _is_anthropic_route(provider, model):
content_list = []
content_list: list[_AssistantContentObject] = []
for part in reasoning_parts:
if part.text:
block = {"type": "thinking", "thinking": part.text}
block = _ThinkingBlock(type="thinking", thinking=part.text)
if part.thought_signature:
sig = part.thought_signature
if isinstance(sig, bytes):
sig = base64.b64encode(sig).decode("utf-8")
block["signature"] = sig
block["signature"] = base64.b64encode(
part.thought_signature
).decode("utf-8")
content_list.append(block)
if isinstance(final_content, list):
content_list.extend(final_content)
elif final_content:
content_list.append({"type": "text", "text": final_content})
return ChatCompletionAssistantMessage(
role=role,
content_list.append(_TextContentObject(type="text", text=final_content))
return _assistant_message(
content=content_list or None,
tool_calls=tool_calls or None,
)
reasoning_content = _merge_reasoning_texts(reasoning_parts)
return ChatCompletionAssistantMessage(
role=role,
return _assistant_message(
content=final_content,
tool_calls=tool_calls or None,
reasoning_content=reasoning_content or None,
@@ -1194,7 +1325,9 @@ def _ensure_tool_results(messages: List[Message], model: str) -> List[Message]:
healed_messages: List[Message] = []
pending_tool_call_ids: List[str] = []
expected_tool_role = "tool_responses" if _is_gemma4_model(model) else "tool"
expected_tool_role: Literal["tool", "tool_responses"] = (
"tool_responses" if _is_gemma4_model(model) else "tool"
)
for message in messages:
role = message.get("role")
@@ -1205,7 +1338,7 @@ def _ensure_tool_results(messages: List[Message], model: str) -> List[Message]:
pending_tool_call_ids,
)
healed_messages.extend(
ChatCompletionToolMessage(
_tool_message(
role=expected_tool_role,
tool_call_id=tool_call_id,
content=_MISSING_TOOL_RESULT_MESSAGE,
@@ -1233,7 +1366,7 @@ def _ensure_tool_results(messages: List[Message], model: str) -> List[Message]:
pending_tool_call_ids,
)
healed_messages.extend(
ChatCompletionToolMessage(
_tool_message(
role=expected_tool_role,
tool_call_id=tool_call_id,
content=_MISSING_TOOL_RESULT_MESSAGE,
@@ -1249,7 +1382,7 @@ async def _get_content(
*,
provider: str = "",
model: str = "",
) -> OpenAIMessageContent:
) -> _MessageContent:
"""Converts a list of parts to litellm content.
Callers may need to filter out thought parts before calling this helper if
@@ -1279,13 +1412,10 @@ async def _get_content(
):
return _decode_inline_text_data(part.inline_data.data)
content_objects = []
content_objects: list[_ContentObject] = []
for part in parts_list:
if part.text:
content_objects.append({
"type": "text",
"text": part.text,
})
content_objects.append(_TextContentObject(type="text", text=part.text))
elif (
part.inline_data
and part.inline_data.data
@@ -1294,31 +1424,35 @@ async def _get_content(
mime_type = _normalize_mime_type(part.inline_data.mime_type)
if mime_type.startswith("text/"):
decoded_text = _decode_inline_text_data(part.inline_data.data)
content_objects.append({
"type": "text",
"text": decoded_text,
})
content_objects.append(
_TextContentObject(type="text", text=decoded_text)
)
continue
base64_string = base64.b64encode(part.inline_data.data).decode("utf-8")
if mime_type.startswith("audio/"):
content_objects.append({
"type": "input_audio",
"input_audio": {
"data": base64_string,
"format": _audio_format_from_mime_type(mime_type),
},
})
content_objects.append(
_AudioContentObject(
type="input_audio",
input_audio={
"data": base64_string,
"format": _audio_format_from_mime_type(mime_type),
},
)
)
continue
data_uri = f"data:{mime_type};base64,{base64_string}"
# LiteLLM providers extract the MIME type from the data URI; avoid
# passing a separate `format` field that some backends reject.
url_content_type = _media_url_content_type(mime_type)
if url_content_type:
content_objects.append({
"type": url_content_type,
url_content_type: {"url": data_uri},
})
if url_content_type == "image_url":
content_objects.append(
_ImageContentObject(type="image_url", image_url={"url": data_uri})
)
elif url_content_type == "video_url":
content_objects.append(
_VideoContentObject(type="video_url", video_url={"url": data_uri})
)
elif mime_type in _SUPPORTED_FILE_CONTENT_MIME_TYPES:
# OpenAI/Azure require file_id from uploaded file, not inline data
if provider in _FILE_ID_REQUIRED_PROVIDERS:
@@ -1327,15 +1461,16 @@ async def _get_content(
purpose="assistants",
custom_llm_provider=provider,
)
content_objects.append({
"type": "file",
"file": {"file_id": file_response.id, "format": mime_type},
})
content_objects.append(
_FileContentObject(
type="file",
file={"file_id": file_response.id, "format": mime_type},
)
)
else:
content_objects.append({
"type": "file",
"file": {"file_data": data_uri},
})
content_objects.append(
_FileContentObject(type="file", file={"file_data": data_uri})
)
else:
raise ValueError(
"LiteLlm(BaseLlm) does not support content part with MIME type "
@@ -1346,10 +1481,11 @@ async def _get_content(
provider in _FILE_ID_REQUIRED_PROVIDERS
and _looks_like_openai_file_id(part.file_data.file_uri)
):
content_objects.append({
"type": "file",
"file": {"file_id": part.file_data.file_uri},
})
content_objects.append(
_FileContentObject(
type="file", file={"file_id": part.file_data.file_uri}
)
)
continue
# Resolve MIME type early: needed before the media-URL shortcut below,
@@ -1357,27 +1493,37 @@ async def _get_content(
# deferred until after all early-continue paths so that providers which
# always fall back to text (anthropic, non-Gemini Vertex AI) are never
# asked for a MIME type they cannot supply.
mime_type = part.file_data.mime_type
if not mime_type:
mime_type = _infer_mime_type_from_uri(part.file_data.file_uri)
if not mime_type and part.file_data.display_name:
file_mime_type = part.file_data.mime_type
if not file_mime_type:
file_mime_type = _infer_mime_type_from_uri(part.file_data.file_uri)
if not file_mime_type and part.file_data.display_name:
guessed_mime_type, _ = mimetypes.guess_type(part.file_data.display_name)
mime_type = guessed_mime_type
if mime_type:
mime_type = _normalize_mime_type(mime_type)
file_mime_type = guessed_mime_type
if file_mime_type:
file_mime_type = _normalize_mime_type(file_mime_type)
# For OpenAI/Azure: HTTP media URLs (image, video, audio) are sent as
# typed URL blocks and must be handled before the generic text fallback.
if provider in _FILE_ID_REQUIRED_PROVIDERS and _is_http_url(
part.file_data.file_uri
):
if mime_type:
url_content_type = _media_url_content_type(mime_type)
if url_content_type:
content_objects.append({
"type": url_content_type,
url_content_type: {"url": part.file_data.file_uri},
})
if file_mime_type:
url_content_type = _media_url_content_type(file_mime_type)
if url_content_type == "image_url":
content_objects.append(
_ImageContentObject(
type="image_url",
image_url={"url": part.file_data.file_uri},
)
)
continue
if url_content_type == "video_url":
content_objects.append(
_VideoContentObject(
type="video_url",
video_url={"url": part.file_data.file_uri},
)
)
continue
if not _is_file_uri_supported(provider, model, part.file_data.file_uri):
@@ -1395,8 +1541,8 @@ async def _get_content(
# 'application/octet-stream' cause a downstream ValueError from LiteLLM
# regardless of whether the value was set explicitly by the caller or
# arrived via a default fallback; raise early with an actionable message.
if not mime_type or mime_type == "application/octet-stream":
type_label = mime_type or "(unknown)"
if not file_mime_type or file_mime_type == "application/octet-stream":
type_label = file_mime_type or "(unknown)"
raise ValueError(
f"Cannot process file_uri {part.file_data.file_uri!r}: MIME type"
f" {type_label!r} is not supported. Please set a specific MIME"
@@ -1406,11 +1552,8 @@ async def _get_content(
file_object: ChatCompletionFileUrlObject = {
"file_id": part.file_data.file_uri,
}
file_object["format"] = mime_type
content_objects.append({
"type": "file",
"file": file_object,
})
file_object["format"] = file_mime_type
content_objects.append(_FileContentObject(type="file", file=file_object))
return content_objects
@@ -1469,7 +1612,7 @@ def _flatten_ollama_content(
for block in blocks:
if isinstance(block, dict) and block.get("type") == "text":
text_value = block.get("text")
if text_value:
if isinstance(text_value, str) and text_value:
text_parts.append(text_value)
if text_parts:
@@ -1554,23 +1697,17 @@ def _build_tool_call_from_json_dict(
if isinstance(call_index, int):
index = call_index
function = Function(
function = ChatCompletionToolCallFunctionChunk(
name=name,
arguments=arguments_payload,
)
# Some LiteLLM types carry an `index` field only in streaming contexts,
# so guard the assignment to stay compatible with older versions.
if hasattr(function, "index"):
function.index = index # type: ignore[attr-defined]
tool_call = ChatCompletionMessageToolCall(
type="function",
id=str(call_id),
function=function,
index=index,
)
# Same reasoning as above: not every ChatCompletionMessageToolCall exposes it.
if hasattr(tool_call, "index"):
tool_call.index = index # type: ignore[attr-defined]
return tool_call
@@ -1880,7 +2017,7 @@ def _function_declaration_to_tool_param(
assert function_declaration.name
parameters = {
parameters: dict[str, Any] = {
"type": "object",
"properties": {},
}
@@ -1899,7 +2036,7 @@ def _function_declaration_to_tool_param(
elif function_declaration.parameters_json_schema:
parameters = function_declaration.parameters_json_schema
tool_params = {
tool_params: dict[str, Any] = {
"type": "function",
"function": {
"name": function_declaration.name,
@@ -2149,7 +2286,7 @@ def _message_to_generate_content_response(
message: Message,
*,
is_partial: bool = False,
model_version: str = None,
model_version: Optional[str] = None,
thought_parts: Optional[List[types.Part]] = None,
) -> LlmResponse:
"""Converts a litellm message to LlmResponse.
@@ -2183,7 +2320,12 @@ def _message_to_generate_content_response(
name=tool_call.function.name,
args=_parse_tool_call_arguments(tool_call.function.arguments),
)
part.function_call.id = tool_call.id
function_call = part.function_call
if function_call is None:
raise ValueError(
"Function-call part factory returned no function call"
)
function_call.id = tool_call.id
if thought_signature:
part.thought_signature = thought_signature
parts.append(part)
@@ -2365,12 +2507,13 @@ async def _get_completion_inputs(
elif message_param_or_list: # Ensure it's not None before appending
messages.append(message_param_or_list)
if llm_request.config.system_instruction:
system_instruction = extract_system_instruction(llm_request.config)
if system_instruction:
messages.insert(
0,
ChatCompletionSystemMessage(
role="system",
content=llm_request.config.system_instruction,
content=system_instruction,
),
)
messages = _ensure_tool_results(messages, model)
@@ -2506,7 +2649,8 @@ def _build_request_log(req: LlmRequest) -> str:
exclude_none=True,
exclude={
"parts": {
i: _EXCLUDED_PART_FIELD for i in range(len(content.parts))
i: _EXCLUDED_PART_FIELD
for i in range(len(content.parts or []))
}
},
)
@@ -2715,7 +2859,7 @@ class LiteLlm(BaseLlm):
llm_client: LiteLLMClient = Field(default_factory=LiteLLMClient, exclude=True)
"""The LLM client to use for the model."""
_additional_args: Dict[str, Any] = None
_additional_args: Dict[str, Any] = PrivateAttr(default_factory=dict)
def __init__(self, model: str, **kwargs: Any) -> None:
"""Initializes the LiteLlm class.
@@ -2865,11 +3009,11 @@ class LiteLlm(BaseLlm):
ChatCompletionMessageToolCall(
type="function",
id=func_data["id"],
function=Function(
function=ChatCompletionToolCallFunctionChunk(
name=func_data["name"],
arguments=args,
index=index,
),
index=index,
)
)
@@ -2896,7 +3040,10 @@ class LiteLlm(BaseLlm):
)
mapped_finish_reason = _map_finish_reason(finish_reason)
llm_response.finish_reason = mapped_finish_reason
if mapped_finish_reason != types.FinishReason.STOP:
if (
mapped_finish_reason is not None
and mapped_finish_reason != types.FinishReason.STOP
):
llm_response.error_code = mapped_finish_reason
llm_response.error_message = _finish_reason_to_error_message(
mapped_finish_reason
@@ -2917,7 +3064,10 @@ class LiteLlm(BaseLlm):
)
mapped_finish_reason = _map_finish_reason(finish_reason)
llm_response.finish_reason = mapped_finish_reason
if mapped_finish_reason != types.FinishReason.STOP:
if (
mapped_finish_reason is not None
and mapped_finish_reason != types.FinishReason.STOP
):
llm_response.error_code = mapped_finish_reason
llm_response.error_message = _finish_reason_to_error_message(
mapped_finish_reason
src/google/adk/models/llm_request.py
+1 -1
View File
@@ -148,7 +148,7 @@ class LlmRequest(BaseModel):
# Process all parts, creating references for non-text parts
non_text_count = 0
for part in instructions.parts:
for part in instructions.parts or []:
if part.text:
# Text part - add to system instruction
text_parts.append(part.text)
tests/unittests/agents/test_gemini_context_cache_manager.py
+15
View File
@@ -28,6 +28,7 @@ from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import Client
from google.genai import types
import pytest
class TestGeminiContextCacheManager:
@@ -866,6 +867,20 @@ class TestGeminiContextCacheManager:
)
assert isinstance(fingerprint, str)
async def test_handle_context_caching_requires_configuration(self):
llm_request = self.create_llm_request()
llm_request.cache_config = None
with pytest.raises(ValueError, match="cache configuration"):
await self.manager.handle_context_caching(llm_request)
async def test_handle_context_caching_requires_model(self):
llm_request = self.create_llm_request()
llm_request.model = None
with pytest.raises(ValueError, match="model name"):
await self.manager.handle_context_caching(llm_request)
def test_parameter_types_enforcement(self):
"""Test that method calls with correct parameter types work properly."""
# Create proper objects
tests/unittests/models/test_anthropic_llm.py
+58 -1
View File
@@ -653,6 +653,48 @@ async def test_anthropic_llm_generate_content_async(
assert responses[0].content.parts[0].text == "Hello, how can I help you?"
@pytest.mark.asyncio
async def test_generate_content_async_collects_declarations_from_all_tools(
generate_content_response,
):
llm = AnthropicLlm(model="claude-sonnet-4-20250514")
llm_request = LlmRequest(
contents=[Content(role="user", parts=[Part.from_text(text="Run both")])],
config=types.GenerateContentConfig(
tools=[
types.Tool(
function_declarations=[
types.FunctionDeclaration(name="first_tool")
]
),
types.Tool(
function_declarations=[
types.FunctionDeclaration(name="second_tool")
]
),
]
),
)
mock_client = MagicMock()
mock_client.messages.create = AsyncMock(
return_value=generate_content_response
)
with mock.patch.object(llm, "_anthropic_client", mock_client):
_ = [
response
async for response in llm.generate_content_async(
llm_request, stream=False
)
]
_, kwargs = mock_client.messages.create.call_args
assert [tool["name"] for tool in kwargs["tools"]] == [
"first_tool",
"second_tool",
]
def test_claude_vertex_client_uses_tracking_headers():
"""Tests that Claude vertex client is called with tracking headers."""
with mock.patch.object(
@@ -814,10 +856,25 @@ def test_part_to_message_block_with_pdf_mime_type_parameters():
assert isinstance(result, dict)
assert result["type"] == "document"
assert result["source"]["type"] == "base64"
assert result["source"]["media_type"] == "application/pdf; name=doc.pdf"
assert result["source"]["media_type"] == "application/pdf"
assert result["source"]["data"] == base64.b64encode(pdf_data).decode()
@pytest.mark.parametrize("mime_type", ["image/png", "application/pdf"])
def test_part_to_message_block_rejects_media_without_data(mime_type):
part = Part(inline_data=types.Blob(mime_type=mime_type))
with pytest.raises(ValueError, match="require.*data"):
part_to_message_block(part)
def test_part_to_message_block_rejects_unsupported_image_mime_type():
part = Part(inline_data=types.Blob(mime_type="image/bmp", data=b"bitmap"))
with pytest.raises(ValueError, match="Unsupported Anthropic image MIME"):
part_to_message_block(part)
content_to_message_param_test_cases = [
(
"user_role_with_text_and_image",
tests/unittests/models/test_apigee_llm.py
+140 -47
View File
@@ -15,12 +15,16 @@
from __future__ import annotations
import os
from typing import AsyncGenerator
from typing import cast
from unittest import mock
from unittest.mock import AsyncMock
from google.adk.models.apigee_llm import ApigeeLlm
from google.adk.models.apigee_llm import CompletionsHTTPClient
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.auth.credentials import Credentials
from google.genai import types
from google.genai.types import Content
from google.genai.types import Part
@@ -33,8 +37,17 @@ VERTEX_BASE_MODEL_ID = 'gemini-pro'
PROXY_URL = 'https://test.apigee.net'
def _response_parts(response: LlmResponse) -> list[types.Part]:
assert response.content is not None
raw_parts = response.content.parts
assert isinstance(raw_parts, list)
parts = [part for part in raw_parts if isinstance(part, types.Part)]
assert len(parts) == len(raw_parts)
return parts
@pytest.fixture
def llm_request():
def llm_request() -> LlmRequest:
"""Provides a sample LlmRequest for testing."""
return LlmRequest(
model=APIGEE_GEMINI_MODEL_ID,
@@ -49,8 +62,8 @@ def llm_request():
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_generate_content_async_non_streaming(
mock_client_constructor, llm_request
):
mock_client_constructor: mock.MagicMock, llm_request: LlmRequest
) -> None:
"""Tests the generate_content_async method for non-streaming responses."""
apigee_llm_instance = ApigeeLlm(
model=APIGEE_GEMINI_MODEL_ID,
@@ -77,7 +90,8 @@ async def test_generate_content_async_non_streaming(
assert len(responses) == 1
llm_response = responses[0]
assert llm_response.content.parts[0].text == 'Test response'
assert _response_parts(llm_response)[0].text == 'Test response'
assert llm_response.content is not None
assert llm_response.content.role == 'model'
mock_client_constructor.assert_called_once()
@@ -99,8 +113,8 @@ async def test_generate_content_async_non_streaming(
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_generate_content_async_streaming(
mock_client_constructor, llm_request
):
mock_client_constructor: mock.MagicMock, llm_request: LlmRequest
) -> None:
"""Tests the generate_content_async method for streaming responses."""
apigee_llm_instance = ApigeeLlm(
model=APIGEE_GEMINI_MODEL_ID,
@@ -137,7 +151,9 @@ async def test_generate_content_async_streaming(
),
]
async def mock_stream_generator():
async def mock_stream_generator() -> (
AsyncGenerator[types.GenerateContentResponse, None]
):
for r in mock_responses:
yield r
@@ -154,7 +170,7 @@ async def test_generate_content_async_streaming(
assert responses
full_text_parts = []
for r in responses:
for p in r.content.parts:
for p in _response_parts(r):
if p.text:
full_text_parts.append(p.text)
full_text = ''.join(full_text_parts)
@@ -170,8 +186,8 @@ async def test_generate_content_async_streaming(
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_generate_content_async_with_custom_headers(
mock_client_constructor, llm_request
):
mock_client_constructor: mock.MagicMock, llm_request: LlmRequest
) -> None:
"""Tests that custom headers are passed in the request."""
custom_headers = {
'X-Custom-Header': 'custom-value',
@@ -209,7 +225,9 @@ async def test_generate_content_async_with_custom_headers(
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_vertex_model_path_parsing(mock_client_constructor):
async def test_vertex_model_path_parsing(
mock_client_constructor: mock.MagicMock,
) -> None:
"""Tests that Vertex AI model paths are parsed correctly."""
apigee_llm = ApigeeLlm(model=APIGEE_VERTEX_MODEL_ID, proxy_url=PROXY_URL)
llm_request = LlmRequest(
@@ -251,7 +269,9 @@ async def test_vertex_model_path_parsing(mock_client_constructor):
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_proxy_url_from_env_variable(mock_client_constructor):
async def test_proxy_url_from_env_variable(
mock_client_constructor: mock.MagicMock,
) -> None:
"""Tests that proxy_url is read from environment variable."""
with mock.patch.dict(
os.environ, {'APIGEE_PROXY_URL': 'https://env.proxy.url'}
@@ -287,6 +307,20 @@ async def test_proxy_url_from_env_variable(mock_client_constructor):
assert kwargs['http_options'].base_url == 'https://env.proxy.url'
def test_clients_require_an_apigee_proxy_url(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv('APIGEE_PROXY_URL', raising=False)
genai_llm = ApigeeLlm(model=APIGEE_GEMINI_MODEL_ID)
with pytest.raises(ValueError, match='Apigee proxy URL is not set'):
_ = genai_llm.api_client
completions_llm = ApigeeLlm(model='apigee/openai/gpt-4o')
with pytest.raises(ValueError, match='Apigee proxy URL is not set'):
_ = completions_llm._completions_http_client
@pytest.mark.parametrize(
('model_string', 'env_vars'),
[
@@ -315,8 +349,8 @@ async def test_proxy_url_from_env_variable(mock_client_constructor):
],
)
def test_vertex_model_missing_project_or_location_raises_error(
model_string, env_vars
):
model_string: str, env_vars: dict[str, str]
) -> None:
"""Tests that ValueError is raised for Vertex models if project or location is missing."""
with mock.patch.dict(os.environ, env_vars, clear=True):
with pytest.raises(ValueError, match='environment variable must be set'):
@@ -384,15 +418,15 @@ def test_vertex_model_missing_project_or_location_raises_error(
)
@mock.patch('google.genai.Client')
async def test_model_string_parsing_and_client_initialization(
mock_client_constructor,
model_string,
use_vertexai_env,
expected_is_vertexai,
expected_api_version,
expected_model_id,
):
mock_client_constructor: mock.MagicMock,
model_string: str,
use_vertexai_env: str | None,
expected_is_vertexai: bool,
expected_api_version: str | None,
expected_model_id: str,
) -> None:
"""Tests model string parsing and genai.Client initialization."""
env_vars = {}
env_vars: dict[str, str] = {}
if use_vertexai_env is not None:
env_vars['GOOGLE_GENAI_USE_ENTERPRISE'] = use_vertexai_env
@@ -449,7 +483,9 @@ async def test_model_string_parsing_and_client_initialization(
'apigee/unknown/model',
],
)
async def test_invalid_model_strings_raise_value_error(invalid_model_string):
async def test_invalid_model_strings_raise_value_error(
invalid_model_string: str,
) -> None:
"""Tests that invalid model strings raise a ValueError."""
with pytest.raises(
ValueError, match=f'Invalid model string: {invalid_model_string}'
@@ -466,7 +502,9 @@ async def test_invalid_model_strings_raise_value_error(invalid_model_string):
'apigee/openai/v1/gpt-3.5-turbo',
],
)
async def test_validate_model_for_chat_completion_providers(model):
async def test_validate_model_for_chat_completion_providers(
model: str,
) -> None:
"""Tests that new providers like OpenAI are accepted."""
# Should not raise ValueError
ApigeeLlm(model=model, proxy_url=PROXY_URL)
@@ -545,7 +583,11 @@ async def test_validate_model_for_chat_completion_providers(model):
),
],
)
def test_api_type_resolution(model, api_type, expected_api_type):
def test_api_type_resolution(
model: str,
api_type: ApigeeLlm.ApiType | str,
expected_api_type: ApigeeLlm.ApiType,
) -> None:
"""Tests that api_type is resolved correctly."""
llm = ApigeeLlm(
model=model,
@@ -565,18 +607,20 @@ def test_api_type_resolution(model, api_type, expected_api_type):
(None, ApigeeLlm.ApiType.UNKNOWN),
],
)
def test_apitype_creation(input_value, expected_type):
def test_apitype_creation(
input_value: str | None, expected_type: ApigeeLlm.ApiType
) -> None:
"""Tests the creation of ApiType enum members."""
assert ApigeeLlm.ApiType(input_value) == expected_type
def test_apitype_creation_invalid():
def test_apitype_creation_invalid() -> None:
"""Tests that invalid ApiType raises ValueError."""
with pytest.raises(ValueError):
ApigeeLlm.ApiType('invalid')
def test_invalid_api_type_raises_error():
def test_invalid_api_type_raises_error() -> None:
"""Tests that invalid string for api_type raises ValueError."""
with pytest.raises(ValueError):
ApigeeLlm(
@@ -588,8 +632,8 @@ def test_invalid_api_type_raises_error():
@pytest.mark.asyncio
async def test_generate_content_async_dispatch_to_completions_client(
llm_request,
):
llm_request: LlmRequest,
) -> None:
"""Tests that generate_content_async uses CompletionsHTTPClient for OpenAI models."""
llm_request.model = 'apigee/openai/gpt-4o'
with (
@@ -682,7 +726,7 @@ async def test_streaming_chat_completions_honors_request_timeout():
'apigee/openai/v1/gpt-3.5-turbo',
],
)
async def test_api_key_injection_openai(model):
async def test_api_key_injection_openai(model: str) -> None:
"""Tests that api_key is injected for OpenAI models."""
apigee_llm = ApigeeLlm(
model=model,
@@ -693,7 +737,7 @@ async def test_api_key_injection_openai(model):
assert client._headers['Authorization'] == 'Bearer sk-test-key'
def test_parse_response_usage_metadata():
def test_parse_response_usage_metadata() -> None:
"""Tests that CompletionsHTTPClient parses usage metadata correctly including reasoning tokens."""
client = CompletionsHTTPClient(base_url='http://test')
response_dict = {
@@ -709,19 +753,21 @@ def test_parse_response_usage_metadata():
},
}
llm_response = client._parse_response(response_dict)
assert llm_response.usage_metadata.prompt_token_count == 10
assert llm_response.usage_metadata.candidates_token_count == 5
assert llm_response.usage_metadata.total_token_count == 15
assert llm_response.usage_metadata.thoughts_token_count == 4
usage_metadata = llm_response.usage_metadata
assert usage_metadata is not None
assert usage_metadata.prompt_token_count == 10
assert usage_metadata.candidates_token_count == 5
assert usage_metadata.total_token_count == 15
assert usage_metadata.thoughts_token_count == 4
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_api_client_passes_credentials_when_provided(
mock_client_constructor, llm_request
):
mock_client_constructor: mock.MagicMock, llm_request: LlmRequest
) -> None:
"""Tests that credentials passed to __init__ are forwarded to genai.Client."""
mock_credentials = mock.Mock()
mock_credentials = cast(Credentials, mock.Mock())
mock_client_instance = mock.Mock()
mock_client_instance.aio.models.generate_content = AsyncMock(
@@ -752,8 +798,8 @@ async def test_api_client_passes_credentials_when_provided(
@pytest.mark.asyncio
@mock.patch('google.genai.Client')
async def test_api_client_omits_credentials_when_not_provided(
mock_client_constructor, llm_request
):
mock_client_constructor: mock.MagicMock, llm_request: LlmRequest
) -> None:
"""Tests that credentials kwarg is not forwarded when not supplied."""
mock_client_instance = mock.Mock()
mock_client_instance.aio.models.generate_content = AsyncMock(
@@ -780,7 +826,7 @@ async def test_api_client_omits_credentials_when_not_provided(
assert 'credentials' not in kwargs
def test_parse_response_with_refusal():
def test_parse_response_with_refusal() -> None:
"""Tests that CompletionsHTTPClient parses refusal correctly."""
client = CompletionsHTTPClient(base_url='http://test')
@@ -794,8 +840,9 @@ def test_parse_response_with_refusal():
}],
}
llm_response = client._parse_response(response_dict)
assert len(llm_response.content.parts) == 1
assert llm_response.content.parts[0].text == '[[REFUSAL]]: I refuse to answer'
response_parts = _response_parts(llm_response)
assert len(response_parts) == 1
assert response_parts[0].text == '[[REFUSAL]]: I refuse to answer'
response_dict_mixed = {
'choices': [{
@@ -808,9 +855,10 @@ def test_parse_response_with_refusal():
}],
}
llm_response_mixed = client._parse_response(response_dict_mixed)
assert len(llm_response_mixed.content.parts) == 1
mixed_parts = _response_parts(llm_response_mixed)
assert len(mixed_parts) == 1
assert (
llm_response_mixed.content.parts[0].text
mixed_parts[0].text
== 'Here is some content\n[[REFUSAL]]: But I refuse to answer the rest'
)
@@ -846,7 +894,9 @@ def test_parse_response_with_refusal():
),
],
)
def test_construct_payload_with_refusal(parts, expected_message):
def test_construct_payload_with_refusal(
parts: list[types.Part], expected_message: dict[str, object]
) -> None:
"""Tests that CompletionsHTTPClient constructs payload with refusal correctly."""
client = CompletionsHTTPClient(base_url='http://test')
req = LlmRequest(
@@ -861,3 +911,46 @@ def test_construct_payload_with_refusal(parts, expected_message):
payload = client._construct_payload(req, stream=False)
messages = payload['messages']
assert messages == [expected_message]
def test_construct_payload_rejects_non_genai_tools() -> None:
def unsupported_tool() -> None:
pass
request = LlmRequest(
model='apigee/openai/gpt-4o',
contents=[],
config=types.GenerateContentConfig(tools=[unsupported_tool]),
)
client = CompletionsHTTPClient(base_url='http://test')
with pytest.raises(TypeError, match='require google.genai.types.Tool'):
client._construct_payload(request, stream=False)
def test_content_conversion_rejects_unnamed_function_call() -> None:
content = types.Content(
role='model',
parts=[types.Part(function_call=types.FunctionCall())],
)
client = CompletionsHTTPClient(base_url='http://test')
with pytest.raises(ValueError, match='must include a name'):
client._content_to_messages(content)
@pytest.mark.parametrize(
'blob',
[
types.Blob(mime_type='image/png'),
types.Blob(data=b'image'),
],
)
def test_content_conversion_rejects_incomplete_inline_data(
blob: types.Blob,
) -> None:
content = types.Content(role='user', parts=[types.Part(inline_data=blob)])
client = CompletionsHTTPClient(base_url='http://test')
with pytest.raises(ValueError, match='Inline data must include'):
client._content_to_messages(content)
tests/unittests/models/test_llm_request.py
+11
View File
@@ -298,6 +298,17 @@ def test_append_instructions_empty_string_list():
assert len(request.contents) == 0
def test_append_instructions_content_without_parts_is_noop():
"""An SDK Content with omitted parts is an empty instruction."""
request = LlmRequest()
user_contents = request.append_instructions(types.Content(role='user'))
assert user_contents == []
assert request.config.system_instruction is None
assert request.contents == []
def test_append_instructions_invalid_input():
"""Test append_instructions with invalid input types."""
request = LlmRequest()