From 7f82142adbc770ebb7d47d2795cdb911fb593fa1 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 3 Aug 2026 15:57:30 -0700 Subject: [PATCH] refactor(types): make google.adk.models pass strict mypy 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 PiperOrigin-RevId: 958623646 --- src/google/adk/models/anthropic_llm.py | 165 ++++--- src/google/adk/models/apigee_llm.py | 95 ++-- src/google/adk/models/base_llm.py | 5 +- src/google/adk/models/base_llm_connection.py | 4 +- src/google/adk/models/cache_metadata.py | 10 +- .../models/gemini_context_cache_manager.py | 112 +++-- .../adk/models/gemini_llm_connection.py | 26 +- src/google/adk/models/gemma_llm.py | 27 +- src/google/adk/models/google_llm.py | 52 ++- src/google/adk/models/lite_llm.py | 430 ++++++++++++------ src/google/adk/models/llm_request.py | 2 +- .../test_gemini_context_cache_manager.py | 15 + tests/unittests/models/test_anthropic_llm.py | 59 ++- tests/unittests/models/test_apigee_llm.py | 187 ++++++-- tests/unittests/models/test_llm_request.py | 11 + 15 files changed, 843 insertions(+), 357 deletions(-) diff --git a/src/google/adk/models/anthropic_llm.py b/src/google/adk/models/anthropic_llm.py index 49113c6e..bd999bf2 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -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() diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index 84d41f6a..4fae3c30 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -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 = ( diff --git a/src/google/adk/models/base_llm.py b/src/google/adk/models/base_llm.py index 6ff701ca..63f2dadd 100644 --- a/src/google/adk/models/base_llm.py +++ b/src/google/adk/models/base_llm.py @@ -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: diff --git a/src/google/adk/models/base_llm_connection.py b/src/google/adk/models/base_llm_connection.py index 8b8e01d1..46bab7f6 100644 --- a/src/google/adk/models/base_llm_connection.py +++ b/src/google/adk/models/base_llm_connection.py @@ -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: diff --git a/src/google/adk/models/cache_metadata.py b/src/google/adk/models/cache_metadata.py index d899ab47..1e76b094 100644 --- a/src/google/adk/models/cache_metadata.py +++ b/src/google/adk/models/cache_metadata.py @@ -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 ( diff --git a/src/google/adk/models/gemini_context_cache_manager.py b/src/google/adk/models/gemini_context_cache_manager.py index bf179ac6..bbe0d067 100644 --- a/src/google/adk/models/gemini_context_cache_manager.py +++ b/src/google/adk/models/gemini_context_cache_manager.py @@ -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 diff --git a/src/google/adk/models/gemini_llm_connection.py b/src/google/adk/models/gemini_llm_connection.py index 1a146223..cc380f22 100644 --- a/src/google/adk/models/gemini_llm_connection.py +++ b/src/google/adk/models/gemini_llm_connection.py @@ -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: diff --git a/src/google/adk/models/gemma_llm.py b/src/google/adk/models/gemma_llm.py index 8fea7152..0599f2f1 100644 --- a/src/google/adk/models/gemma_llm.py +++ b/src/google/adk/models/gemma_llm.py @@ -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: diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 02b1f261..839d8d41 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -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 [])) } }, ) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 4cf6e5c2..4656c1a9 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -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 diff --git a/src/google/adk/models/llm_request.py b/src/google/adk/models/llm_request.py index 48fc51df..96a9c540 100644 --- a/src/google/adk/models/llm_request.py +++ b/src/google/adk/models/llm_request.py @@ -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) diff --git a/tests/unittests/agents/test_gemini_context_cache_manager.py b/tests/unittests/agents/test_gemini_context_cache_manager.py index 4e3ae7c3..03501668 100644 --- a/tests/unittests/agents/test_gemini_context_cache_manager.py +++ b/tests/unittests/agents/test_gemini_context_cache_manager.py @@ -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 diff --git a/tests/unittests/models/test_anthropic_llm.py b/tests/unittests/models/test_anthropic_llm.py index 0dafa121..b08ae692 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -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", diff --git a/tests/unittests/models/test_apigee_llm.py b/tests/unittests/models/test_apigee_llm.py index 38c10d61..0bd7996a 100644 --- a/tests/unittests/models/test_apigee_llm.py +++ b/tests/unittests/models/test_apigee_llm.py @@ -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) diff --git a/tests/unittests/models/test_llm_request.py b/tests/unittests/models/test_llm_request.py index ca4ef5f4..6cc61ddb 100644 --- a/tests/unittests/models/test_llm_request.py +++ b/tests/unittests/models/test_llm_request.py @@ -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()