From 9c4c44536904f5cf3301a5abb910a5666344a8c5 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 20 Feb 2026 13:26:59 -0800 Subject: [PATCH] feat: Add /chat/completions integration to ApigeeLlm PiperOrigin-RevId: 873049983 --- src/google/adk/models/apigee_llm.py | 634 +++++++++++++++++- tests/unittests/models/test_apigee_llm.py | 174 ++++- .../models/test_completions_http_client.py | 440 ++++++++++++ 3 files changed, 1235 insertions(+), 13 deletions(-) create mode 100644 tests/unittests/models/test_completions_http_client.py diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index 92f94c75..90a91f32 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -12,21 +12,31 @@ # See the License for the specific language governing permissions and # limitations under the License. - from __future__ import annotations +import asyncio +import atexit +import base64 +import collections.abc +import enum from functools import cached_property +import json import logging import os +from typing import Any +from typing import AsyncGenerator from typing import Optional from typing import TYPE_CHECKING from google.adk import version as adk_version from google.genai import types +import httpx +import tenacity from typing_extensions import override from ..utils.env_utils import is_env_enabled from .google_llm import Gemini +from .llm_response import LlmResponse if TYPE_CHECKING: from google.genai import Client @@ -49,6 +59,20 @@ class ApigeeLlm(Gemini): model: The name of the Gemini model. """ + class ApiType(str, enum.Enum): + """The supported API types for Apigee LLM.""" + + UNKNOWN = 'unknown' + CHAT_COMPLETIONS = 'chat_completions' + GENAI = 'genai' + + @classmethod + def _missing_(cls, value): + # Empty string or None should return UNKNOWN. + if not value: + return cls.UNKNOWN + return super()._missing_(value) + def __init__( self, *, @@ -56,6 +80,7 @@ class ApigeeLlm(Gemini): proxy_url: str | None = None, custom_headers: dict[str, str] | None = None, retry_options: Optional[types.HttpRetryOptions] = None, + api_type: ApiType | str = ApiType.UNKNOWN, ): """Initializes the Apigee LLM backend. @@ -80,19 +105,31 @@ class ApigeeLlm(Gemini): - `apigee/vertex_ai/gemini-2.5-flash` - `apigee/gemini/v1/gemini-2.5-flash` - `apigee/vertex_ai/v1beta/gemini-2.5-flash` - proxy_url: The URL of the Apigee proxy. custom_headers: A dictionary of headers to be sent with the request. + If needed, you can add authorization headers here, for example: + {'Authorization': f'Bearer {API_KEY}'}. ApigeeLlm already handles + authorization headers in Vertex AI and Gemini API calls. retry_options: Allow google-genai to retry failed responses. - """ + api_type: The type of API to use. One of `ApiType` or string. + """ # fmt: skip super().__init__(model=model, retry_options=retry_options) # Validate the model string. Create a helper method to validate the model # string. if not _validate_model_string(model): raise ValueError(f'Invalid model string: {model}') - - self._isvertexai = _identify_vertexai(model) + if isinstance(api_type, str): + api_type = ApigeeLlm.ApiType(api_type) + if api_type and api_type != ApigeeLlm.ApiType.UNKNOWN: + self._api_type = api_type + elif model.startswith(('apigee/gemini/', 'apigee/vertex_ai/')): + self._api_type = ApigeeLlm.ApiType.GENAI + elif model.startswith('apigee/openai/'): + self._api_type = ApigeeLlm.ApiType.CHAT_COMPLETIONS + else: + self._api_type = ApigeeLlm.ApiType.GENAI + self._isvertexai = _identify_vertexai(model, self._api_type) # Set the project and location for Vertex AI. if self._isvertexai: @@ -131,6 +168,42 @@ class ApigeeLlm(Gemini): r'apigee\/.*', ] + @cached_property + def _completions_http_client(self) -> CompletionsHTTPClient: + """Provides the completions HTTP client.""" + return CompletionsHTTPClient( + base_url=self._proxy_url, + headers=self._merge_tracking_headers(self._custom_headers), + retry_options=self.retry_options, + ) + + @override + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + if self._api_type == ApigeeLlm.ApiType.CHAT_COMPLETIONS: + await self._preprocess_other_requests(llm_request) + async for ( + response + ) in self._completions_http_client.generate_content_async( + llm_request, stream + ): + yield response + else: + async for response in super().generate_content_async(llm_request, stream): + yield response + + async def _preprocess_other_requests(self, llm_request: LlmRequest) -> None: + """Preprocesses the request for non-Gemini/Vertex AI models.""" + llm_request.model = _get_model_id(llm_request.model) + if llm_request.config and llm_request.config.tools: + # Check if computer use is configured + for tool in llm_request.config.tools: + if isinstance(tool, types.Tool) and tool.computer_use: + llm_request.config.system_instruction = None + await self._adapt_computer_use_tool(llm_request) + self._maybe_append_user_content(llm_request) + @cached_property def api_client(self) -> Client: """Provides the api client. @@ -167,11 +240,25 @@ class ApigeeLlm(Gemini): await super()._preprocess_request(llm_request) -def _identify_vertexai(model: str) -> bool: - """Returns True if the model spec starts with apigee/vertex_ai.""" - return not model.startswith('apigee/gemini/') and ( - model.startswith('apigee/vertex_ai/') - or is_env_enabled(_GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME) +def _identify_vertexai(model: str, api_type: ApigeeLlm.ApiType) -> bool: + """Returns if a model is Vertex AI. + + 1. The api_type is GENAI or UNKNOWN. + 2. The model is provider is Vertex AI model or the + GOOGLE_GENAI_USE_VERTEXAI environment variable is set to TRUE or 1. + + Args: + model: The model string. + api_type: The type of API to use. + """ + if api_type not in (ApigeeLlm.ApiType.GENAI, ApigeeLlm.ApiType.UNKNOWN): + return False + if model.startswith('apigee/gemini/'): + return False + if model.startswith('apigee/openai/'): + return False + return model.startswith('apigee/vertex_ai/') or is_env_enabled( + _GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME ) @@ -240,7 +327,7 @@ def _validate_model_string(model: str) -> bool: # and model_id are present. This is a valid format. if len(components) == 3: # Format: // - if components[0] not in ('vertex_ai', 'gemini'): + if components[0] not in ('vertex_ai', 'gemini', 'openai'): return False if not components[1].startswith('v'): return False @@ -249,10 +336,533 @@ def _validate_model_string(model: str) -> bool: # If the model string has 2 components, it means either the provider or the # version (but not both), and model_id are present. if len(components) == 2: - if components[0] in ['vertex_ai', 'gemini']: + if components[0] in ['vertex_ai', 'gemini', 'openai']: return True if components[0].startswith('v'): return True return False return False + + +class CompletionsHTTPClient: + """A generic HTTP client for completions, compatible with OpenAI API.""" + + def __init__( + self, + base_url: str, + headers: dict[str, str] | None = None, + retry_options: Optional[types.HttpRetryOptions] = None, + ): + self._base_url = base_url + self._headers = headers or {} + self.retry_options = retry_options + + def __del__(self) -> None: + self.close() + + @cached_property + def _client(self) -> httpx.AsyncClient: + """Provides the httpx client.""" + client = httpx.AsyncClient( + base_url=self._base_url, + headers=self._headers, + timeout=None, + follow_redirects=True, + ) + atexit.register(self._cleanup_client, client) + return client + + @staticmethod + def _cleanup_client(client: httpx.AsyncClient) -> None: + """Cleans up the httpx client.""" + if client.is_closed: + return + try: + loop = asyncio.get_running_loop() + loop.create_task(client.aclose()) + except RuntimeError: + try: + # This fails if aynscio.run is already called in main and is being closed. + asyncio.run(client.aclose()) + except RuntimeError: + pass + + def close(self) -> None: + if '_client' not in self.__dict__: + return + self._cleanup_client(self._client) + + async def aclose(self) -> None: + if '_client' not in self.__dict__: + return + if self._client.is_closed: + return + await self._client.aclose() + + def _get_retry_kwargs(self) -> dict[str, Any]: + """Returns the retry kwargs for tenacity.""" + if not self.retry_options: + return {'stop': tenacity.stop_after_attempt(1), 'reraise': True} + + default_attempts = 5 + default_initial_delay = 1.0 + default_max_delay = 60.0 + default_exp_base = 2 + default_jitter = 1 + default_status_codes = (408, 429, 500, 502, 503, 504) + + opts = self.retry_options + stop = tenacity.stop_after_attempt( + opts.attempts if opts.attempts is not None else default_attempts + ) + + retriable_codes = ( + opts.http_status_codes + if opts.http_status_codes is not None + else default_status_codes + ) + + retry_network = tenacity.retry_if_exception_type(httpx.NetworkError) + + def is_retriable(e: Exception) -> bool: + if isinstance(e, httpx.HTTPStatusError): + return e.response.status_code in retriable_codes + return False + + retry_status = tenacity.retry_if_exception(is_retriable) + + wait = tenacity.wait_exponential_jitter( + initial=( + opts.initial_delay + if opts.initial_delay is not None + else default_initial_delay + ), + max=( + opts.max_delay if opts.max_delay is not None else default_max_delay + ), + exp_base=( + opts.exp_base if opts.exp_base is not None else default_exp_base + ), + jitter=opts.jitter if opts.jitter is not None else default_jitter, + ) + + return { + 'stop': stop, + 'retry': tenacity.retry_any(retry_network, retry_status), + 'reraise': True, + 'wait': wait, + } + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool + ) -> AsyncGenerator[LlmResponse, None]: + """Generates content using the OpenAI-compatible HTTP API.""" + payload = self._construct_payload(llm_request, stream) + headers = self._headers.copy() + headers['Content-Type'] = 'application/json' + + url = self._base_url + if not url: + raise ValueError('Base URL is not set.') + + if not url.endswith('/chat/completions'): + url = f"{url.rstrip('/')}/chat/completions" + + if stream: + raise NotImplementedError('Streaming is not supported yet.') + else: + response = await self._httpx_post_with_retry(url, payload, headers) + data = response.json() + yield self._parse_response(data) + + async def _httpx_post_with_retry( + self, url: str, payload: dict[str, Any], headers: dict[str, str] + ) -> httpx.Response: + """Sends a POST request and handles retries.""" + retry_kwargs = self._get_retry_kwargs() + async for attempt in tenacity.AsyncRetrying(**retry_kwargs): + with attempt: + response = await self._client.post(url, json=payload, headers=headers) + response.raise_for_status() + return response + + async def _handle_streaming_response( + self, response: httpx.Response + ) -> AsyncGenerator[LlmResponse, None]: + """Handles streaming response from OpenAI-compatible API.""" + raise NotImplementedError('Streaming is not supported yet.') + + def _construct_payload( + self, llm_request: LlmRequest, stream: bool + ) -> dict[str, Any]: + """Constructs the payload from the LlmRequest.""" + messages = [] + if llm_request.config and llm_request.config.system_instruction: + content = self._serialize_system_instruction( + llm_request.config.system_instruction + ) + if content: + messages.append({ + 'role': 'system', + 'content': content, + }) + + for content in llm_request.contents: + messages += self._content_to_messages(content) + + payload = { + 'model': _get_model_id(llm_request.model), + 'messages': messages, + 'stream': stream, + } + + if llm_request.config: + self._map_config_parameters(llm_request.config, payload) + self._map_tools(llm_request.config, payload) + + return payload + + def _map_config_parameters( + self, config: types.GenerateContentConfig, payload: dict[str, Any] + ) -> None: + """Maps configuration parameters to the payload.""" + if config.temperature is not None: + payload['temperature'] = config.temperature + if config.top_p is not None: + payload['top_p'] = config.top_p + if config.max_output_tokens is not None: + payload['max_tokens'] = config.max_output_tokens + if config.stop_sequences: + payload['stop'] = config.stop_sequences + if config.frequency_penalty is not None: + payload['frequency_penalty'] = config.frequency_penalty + if config.presence_penalty is not None: + payload['presence_penalty'] = config.presence_penalty + if config.seed is not None: + payload['seed'] = config.seed + if config.candidate_count is not None: + payload['n'] = config.candidate_count + if config.response_logprobs: + payload['logprobs'] = True + if config.logprobs is not None: + payload['top_logprobs'] = config.logprobs + + if config.response_json_schema: + payload['response_format'] = { + 'type': 'json_schema', + 'json_schema': config.response_json_schema, + } + elif config.response_mime_type == 'application/json': + payload['response_format'] = {'type': 'json_object'} + + def _map_tools( + self, config: types.GenerateContentConfig, payload: dict[str, Any] + ) -> None: + """Maps tools and tool configuration to the payload.""" + if config.tools: + tools = [] + for tool in config.tools: + if tool.function_declarations: + for func in tool.function_declarations: + tools.append(self._function_declaration_to_tool(func)) + if tools: + payload['tools'] = tools + if config.tool_config and config.tool_config.function_calling_config: + mode = config.tool_config.function_calling_config.mode + if mode == types.FunctionCallingConfigMode.ANY: + payload['tool_choice'] = 'required' + elif mode == types.FunctionCallingConfigMode.NONE: + payload['tool_choice'] = 'none' + elif mode == types.FunctionCallingConfigMode.AUTO: + payload['tool_choice'] = 'auto' + + def _content_to_messages( + self, content: types.Content + ) -> list[dict[str, Any]]: + """Converts a Content object to /chat/completions messages.""" + role = content.role + if role == 'model': + role = 'assistant' + + tool_calls = [] + content_parts = [] + + function_responses = [] + + for part in content.parts or []: + self._process_content_part(content, part, tool_calls, content_parts) + if part.function_response: + function_responses.append({ + 'role': 'tool', + 'tool_call_id': part.function_response.id, + 'content': json.dumps(part.function_response.response), + }) + if function_responses: + return function_responses + + message = {'role': role} + if tool_calls: + message['tool_calls'] = tool_calls + if not content_parts: + message['content'] = None + + if content_parts: + if len(content_parts) == 1 and content_parts[0]['type'] == 'text': + message['content'] = content_parts[0]['text'] + else: + message['content'] = content_parts + return [message] + + def _process_content_part( + self, + content: types.Content, + part: types.Part, + tool_calls: list[dict[str, Any]], + content_parts: list[dict[str, Any]], + ) -> None: + """Processes a single Part and updates tool_calls or content_parts.""" + if content.role != 'user' and ( + part.inline_data + or ( + part.file_data + and part.file_data.mime_type + and part.file_data.mime_type.startswith('image') + ) + ): + logger.warning('Image data is not supported for assistant turns.') + return + + if part.function_call: + tool_call = { + 'id': part.function_call.id or 'call_' + part.function_call.name, + 'type': 'function', + 'function': { + 'name': part.function_call.name, + 'arguments': ( + json.dumps(part.function_call.args) + if part.function_call.args + else '{}' + ), + }, + } + if part.thought_signature: + sig = part.thought_signature + if isinstance(sig, bytes): + sig = base64.b64encode(sig).decode('utf-8') + tool_call['extra_content'] = { + 'google': { + 'thought_signature': sig, + }, + } + tool_calls.append(tool_call) + elif part.function_response: + # Handled in the loop to return immediately + pass + elif part.text: + content_parts.append({'type': 'text', 'text': part.text}) + elif part.inline_data: + mime_type = part.inline_data.mime_type + data = base64.b64encode(part.inline_data.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: + if part.file_data.file_uri: + content_parts.append({ + 'type': 'image_url', + 'image_url': {'url': part.file_data.file_uri}, + }) + elif part.executable_code: + logger.warning( + 'Executable code is not supported in the standard Chat Completions' + ' API.' + ) + elif part.code_execution_result: + logger.warning( + 'Code execution result is not supported in the standard Chat' + ' Completions API.' + ) + + def _function_declaration_to_tool( + self, func: types.FunctionDeclaration + ) -> dict[str, Any]: + """Converts a FunctionDeclaration to an OpenAI tool dictionary.""" + parameters = {} + if func.parameters_json_schema: + parameters = func.parameters_json_schema + elif func.parameters: + parameters = func.parameters.model_dump(exclude_none=True) + + return { + 'type': 'function', + 'function': { + 'name': func.name, + 'description': func.description, + 'parameters': parameters, + }, + } + + def _serialize_system_instruction( + self, system_instruction: Optional[types.ContentUnion] + ) -> str | None: + """Serializes system instruction to a string from ContentUnion type.""" + if not system_instruction: + return None + if isinstance(system_instruction, str): + return system_instruction + if isinstance(system_instruction, types.Part): + return system_instruction.text + if isinstance(system_instruction, types.Content): + return ''.join( + part.text for part in system_instruction.parts if part.text + ) + if isinstance(system_instruction, dict): + part = types.Part(**system_instruction) + return part.text + if isinstance(system_instruction, collections.abc.Iterable): + parts = [] + for item in system_instruction: + if isinstance(item, str): + parts.append(types.Part(text=item)) + elif isinstance(item, types.Part): + parts.append(item) + elif isinstance(item, dict): + parts.append(types.Part(**item)) + return ''.join(part.text for part in parts if part.text) + return None + + def _parse_logprobs( + self, logprobs_data: dict[str, Any] | None + ) -> types.LogprobsResult | None: + """Parses OpenAI logprobs data into LogprobsResult.""" + if not logprobs_data or 'content' not in logprobs_data: + return None + + chosen_candidates = [] + top_candidates = [] + + for item in logprobs_data['content']: + chosen_candidates.append( + types.LogprobsResultCandidate( + token=item.get('token'), + log_probability=item.get('logprob'), + # OpenAI text format usually doesn't expose ID easily here + token_id=None, + ) + ) + + if 'top_logprobs' in item: + current_top_candidates = [] + for top_item in item['top_logprobs']: + current_top_candidates.append( + types.LogprobsResultCandidate( + token=top_item.get('token'), + log_probability=top_item.get('logprob'), + token_id=None, + ) + ) + top_candidates.append( + types.LogprobsResultTopCandidates(candidates=current_top_candidates) + ) + + return types.LogprobsResult( + chosen_candidates=chosen_candidates, top_candidates=top_candidates + ) + + def _parse_response(self, response: dict[str, Any]) -> LlmResponse: + """Parses an OpenAI response dictionary into an LlmResponse.""" + choices = response.get('choices', []) + if not choices: + return LlmResponse() + + choice = choices[0] + message = choice.get('message', {}) + role = message.get('role', 'model') + if role == 'assistant': + role = 'model' + + parts = [] + content_str = message.get('content') + if content_str: + parts.append(types.Part.from_text(text=content_str)) + + tool_calls = message.get('tool_calls') + if tool_calls: + for tool_call in tool_calls: + call_type = tool_call.get('type', 'unknown') + # TODO: Add support for 'custom' type. + if call_type != 'function': + raise ValueError( + f'Unsupported tool_call type: {call_type} in call {tool_call}' + ) + func = tool_call.get('function', {}) + part = self._parse_function_call(func) + parts.append(part) + + function_call = message.get('function_call') + if function_call: + part = self._parse_function_call(function_call) + parts.append(part) + + usage = response.get('usage', {}) + usage_metadata = types.GenerateContentResponseUsageMetadata( + prompt_token_count=usage.get('prompt_tokens', 0), + candidates_token_count=usage.get('completion_tokens', 0), + total_token_count=usage.get('total_tokens', 0), + ) + + logprobs_result = self._parse_logprobs(choice.get('logprobs')) + + custom_metadata = { + 'id': response.get('id'), + 'created': response.get('created'), + 'model': response.get('model'), + 'system_fingerprint': response.get('system_fingerprint'), + 'service_tier': response.get('service_tier'), + } + custom_metadata = { + k: v for k, v in custom_metadata.items() if v is not None + } + + return LlmResponse( + content=types.Content(role=role, parts=parts), + usage_metadata=usage_metadata, + finish_reason=self._map_finish_reason(choice.get('finish_reason')), + logprobs_result=logprobs_result, + model_version=response.get('model'), + custom_metadata=custom_metadata, + ) + + def _map_finish_reason(self, reason: str | None) -> types.FinishReason: + if reason == 'stop': + return types.FinishReason.STOP + if reason == 'length': + return types.FinishReason.MAX_TOKENS + if reason == 'tool_calls': + return types.FinishReason.STOP + if reason == 'content_filter': + return types.FinishReason.SAFETY + return types.FinishReason.FINISH_REASON_UNSPECIFIED + + def _parse_function_call(self, func: dict[str, Any]) -> types.Part: + """Parses a function call dictionary into a Part.""" + name = func.get('name') + args_str = func.get('arguments', '{}') + try: + args = json.loads(args_str) + except json.JSONDecodeError: + args = {} + tool_part = types.Part.from_function_call(name=name, args=args) + if tool_part.function_call: + tool_part.function_call.id = func.get('id', None) + # Add support for gemini's thought_signature. + thought_signature = ( + func.get('extra_content', {}) + .get('google', {}) + .get('thought_signature', '') + ) + if thought_signature: + if isinstance(thought_signature, str): + thought_signature = base64.b64decode(thought_signature) + tool_part.thought_signature = thought_signature + return tool_part diff --git a/tests/unittests/models/test_apigee_llm.py b/tests/unittests/models/test_apigee_llm.py index 67894ea8..c57bc9fc 100644 --- a/tests/unittests/models/test_apigee_llm.py +++ b/tests/unittests/models/test_apigee_llm.py @@ -19,6 +19,7 @@ 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.genai import types from google.genai.types import Content @@ -441,7 +442,6 @@ async def test_model_string_parsing_and_client_initialization( @pytest.mark.parametrize( 'invalid_model_string', [ - 'apigee/openai/v1/gpt', 'apigee/', # Missing model_id 'apigee', # Invalid format 'gemini-pro', # Invalid format @@ -455,3 +455,175 @@ async def test_invalid_model_strings_raise_value_error(invalid_model_string): ValueError, match=f'Invalid model string: {invalid_model_string}' ): ApigeeLlm(model=invalid_model_string, proxy_url=PROXY_URL) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'model', + [ + 'apigee/openai/gpt-4o', + 'apigee/openai/v1/gpt-4o', + 'apigee/openai/v1/gpt-3.5-turbo', + ], +) +async def test_validate_model_for_chat_completion_providers(model): + """Tests that new providers like OpenAI are accepted.""" + # Should not raise ValueError + ApigeeLlm(model=model, proxy_url=PROXY_URL) + + +@pytest.mark.parametrize( + ('model', 'api_type', 'expected_api_type'), + [ + # Default case (input defaults to UNKNOWN) + ( + 'apigee/openai/gpt-4o', + ApigeeLlm.ApiType.UNKNOWN, + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/openai/v1/gpt-3.5-turbo', + ApigeeLlm.ApiType.UNKNOWN, + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/gemini/v1/gemini-pro', + ApigeeLlm.ApiType.UNKNOWN, + ApigeeLlm.ApiType.GENAI, + ), + ( + 'apigee/vertex_ai/gemini-pro', + ApigeeLlm.ApiType.UNKNOWN, + ApigeeLlm.ApiType.GENAI, + ), + ( + 'apigee/vertex_ai/v1beta/gemini-1.5-pro', + ApigeeLlm.ApiType.UNKNOWN, + ApigeeLlm.ApiType.GENAI, + ), + # Override by setting the ApiType + ( + 'apigee/gemini/pro', + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/gemini/pro', + ApigeeLlm.ApiType.GENAI, + ApigeeLlm.ApiType.GENAI, + ), + ( + 'apigee/openai/gpt-4o', + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/openai/gpt-4o', + ApigeeLlm.ApiType.GENAI, + ApigeeLlm.ApiType.GENAI, + ), + # Override by setting the ApiType as a string + ( + 'apigee/gemini/pro', + 'chat_completions', + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/gemini/pro', + 'genai', + ApigeeLlm.ApiType.GENAI, + ), + ( + 'apigee/openai/gpt-4o', + 'chat_completions', + ApigeeLlm.ApiType.CHAT_COMPLETIONS, + ), + ( + 'apigee/openai/gpt-4o', + 'genai', + ApigeeLlm.ApiType.GENAI, + ), + ], +) +def test_api_type_resolution(model, api_type, expected_api_type): + """Tests that api_type is resolved correctly.""" + llm = ApigeeLlm( + model=model, + proxy_url=PROXY_URL, + api_type=api_type, + ) + assert llm._api_type == expected_api_type + + +@pytest.mark.parametrize( + ('input_value', 'expected_type'), + [ + ('chat_completions', ApigeeLlm.ApiType.CHAT_COMPLETIONS), + ('genai', ApigeeLlm.ApiType.GENAI), + ('unknown', ApigeeLlm.ApiType.UNKNOWN), + ('', ApigeeLlm.ApiType.UNKNOWN), + (None, ApigeeLlm.ApiType.UNKNOWN), + ], +) +def test_apitype_creation(input_value, expected_type): + """Tests the creation of ApiType enum members.""" + assert ApigeeLlm.ApiType(input_value) == expected_type + + +def test_apitype_creation_invalid(): + """Tests that invalid ApiType raises ValueError.""" + with pytest.raises(ValueError): + ApigeeLlm.ApiType('invalid') + + +def test_invalid_api_type_raises_error(): + """Tests that invalid string for api_type raises ValueError.""" + with pytest.raises(ValueError): + ApigeeLlm( + model='apigee/gemini-pro', + proxy_url=PROXY_URL, + api_type='invalid_type', + ) + + +@pytest.mark.asyncio +async def test_generate_content_async_dispatch_to_completions_client( + llm_request, +): + """Tests that generate_content_async uses CompletionsHTTPClient for OpenAI models.""" + llm_request.model = 'apigee/openai/gpt-4o' + with ( + mock.patch.object( + CompletionsHTTPClient, + 'generate_content_async', + ) as mock_completions_generate_content, + mock.patch('google.genai.Client') as mock_genai_client, + ): + apigee_llm = ApigeeLlm(model='apigee/openai/gpt-4o', proxy_url=PROXY_URL) + _ = [ + r + async for r in apigee_llm.generate_content_async( + llm_request, stream=False + ) + ] + mock_completions_generate_content.assert_called_once() + mock_genai_client.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'model', + [ + 'apigee/openai/gpt-4o', + 'apigee/openai/v1/gpt-3.5-turbo', + ], +) +async def test_api_key_injection_openai(model): + """Tests that api_key is injected for OpenAI models.""" + apigee_llm = ApigeeLlm( + model=model, + proxy_url=PROXY_URL, + custom_headers={'Authorization': 'Bearer sk-test-key'}, + ) + client = apigee_llm._completions_http_client + assert client._headers['Authorization'] == 'Bearer sk-test-key' diff --git a/tests/unittests/models/test_completions_http_client.py b/tests/unittests/models/test_completions_http_client.py new file mode 100644 index 00000000..f16376d7 --- /dev/null +++ b/tests/unittests/models/test_completions_http_client.py @@ -0,0 +1,440 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest import mock +from unittest.mock import AsyncMock + +from google.adk.models.apigee_llm import CompletionsHTTPClient +from google.adk.models.llm_request import LlmRequest +from google.genai import types +import httpx +import pytest + + +@pytest.fixture +def client(): + return CompletionsHTTPClient(base_url='https://example.com') + + +@pytest.fixture(name='llm_request') +def fixture_llm_request(): + return LlmRequest( + model='apigee/open_llama', + contents=[ + types.Content(role='user', parts=[types.Part.from_text(text='Hello')]) + ], + ) + + +@pytest.mark.asyncio +async def test_construct_payload_basic_payload(client, llm_request): + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + call_args = mock_post.call_args + url = call_args[0][0] + kwargs = call_args[1] + + assert url == 'https://example.com/chat/completions' + payload = kwargs['json'] + assert payload['model'] == 'open_llama' + assert payload['stream'] is False + assert len(payload['messages']) == 1 + assert payload['messages'][0]['role'] == 'user' + assert payload['messages'][0]['content'] == 'Hello' + + +@pytest.mark.asyncio +async def test_construct_payload_with_config(client, llm_request): + llm_request.config = types.GenerateContentConfig( + temperature=0.7, + top_p=0.9, + max_output_tokens=100, + stop_sequences=['STOP'], + frequency_penalty=0.5, + presence_penalty=0.5, + seed=42, + candidate_count=2, + response_mime_type='application/json', + ) + + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + payload = mock_post.call_args[1]['json'] + + assert payload['temperature'] == 0.7 + assert payload['top_p'] == 0.9 + assert payload['max_tokens'] == 100 + assert payload['stop'] == ['STOP'] + assert payload['frequency_penalty'] == 0.5 + assert payload['presence_penalty'] == 0.5 + assert payload['seed'] == 42 + assert payload['n'] == 2 + assert payload['response_format'] == {'type': 'json_object'} + + +@pytest.mark.asyncio +async def test_construct_payload_with_tools(client, llm_request): + tool = types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name='get_weather', + description='Get weather', + parameters=types.Schema( + type=types.Type.OBJECT, + properties={'location': types.Schema(type=types.Type.STRING)}, + ), + ) + ] + ) + llm_request.config = types.GenerateContentConfig(tools=[tool]) + + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + payload = mock_post.call_args[1]['json'] + assert 'tools' in payload + assert payload['tools'][0]['function']['name'] == 'get_weather' + + +@pytest.mark.asyncio +async def test_construct_payload_system_instruction(client, llm_request): + llm_request.config = types.GenerateContentConfig( + system_instruction='You are a helpful assistant.' + ) + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{'message': {'role': 'assistant', 'content': 'Hi'}}] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + payload = mock_post.call_args[1]['json'] + assert payload['messages'][0]['role'] == 'system' + assert payload['messages'][0]['content'] == 'You are a helpful assistant.' + # Ensure user message follows system + assert payload['messages'][1]['role'] == 'user' + + +@pytest.mark.asyncio +async def test_construct_payload_multimodal_content(client): + # Mock inline_data for image + image_data = b'fake_image_bytes' + llm_request = LlmRequest( + model='apigee/open_llama', + contents=[ + types.Content( + role='user', + parts=[ + types.Part.from_text(text='What is this?'), + types.Part.from_bytes( + data=image_data, mime_type='image/jpeg' + ), + ], + ) + ], + ) + + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [ + {'message': {'role': 'assistant', 'content': 'It is an image'}} + ] + } + + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + payload = mock_post.call_args[1]['json'] + assert len(payload['messages']) == 1 + message = payload['messages'][0] + assert message['role'] == 'user' + assert isinstance(message['content'], list) + assert len(message['content']) == 2 + assert message['content'][0] == {'type': 'text', 'text': 'What is this?'} + assert message['content'][1]['type'] == 'image_url' + # Base64 encoding of b'fake_image_bytes' is 'ZmFrZV9pbWFnZV9ieXRlcw==' + assert message['content'][1]['image_url']['url'] == ( + 'data:image/jpeg;base64,ZmFrZV9pbWFnZV9ieXRlcw==' + ) + + +@pytest.mark.asyncio +async def test_construct_payload_image_file_uri(client): + llm_request = LlmRequest( + model='apigee/open_llama', + contents=[ + types.Content( + role='user', + parts=[ + types.Part.from_uri( + file_uri='https://example.com/image.jpg', + mime_type='image/jpeg', + ) + ], + ) + ], + ) + + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [ + {'message': {'role': 'assistant', 'content': 'It is an image'}} + ] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + payload = mock_post.call_args[1]['json'] + assert len(payload['messages']) == 1 + message = payload['messages'][0] + assert message['role'] == 'user' + assert isinstance(message['content'], list) + assert message['content'][0] == { + 'type': 'image_url', + 'image_url': {'url': 'https://example.com/image.jpg'}, + } + + +@pytest.mark.asyncio +async def test_generate_content_async_function_call_response( + client, llm_request +): + # Mock response with tool call + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'role': 'assistant', + 'content': None, + 'tool_calls': [{ + 'id': 'call_123', + 'type': 'function', + 'function': { + 'name': 'get_weather', + 'arguments': '{"location": "London"}', + }, + }], + } + }] + } + mock_response.status_code = 200 + + with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response): + responses = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + assert len(responses) == 1 + part = responses[0].content.parts[0] + assert part.function_call + assert part.function_call.name == 'get_weather' + assert part.function_call.args == {'location': 'London'} + assert part.function_call.id == 'call_123' + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('response_json_schema', 'response_mime_type', 'expected_response_format'), + [ + # Case 1: Only response_json_schema is provided + ( + {'type': 'object', 'properties': {'name': {'type': 'string'}}}, + None, + { + 'type': 'json_schema', + 'json_schema': { + 'type': 'object', + 'properties': {'name': {'type': 'string'}}, + }, + }, + ), + # Case 2: Both provided, schema takes precedence + ( + {'type': 'object', 'properties': {'name': {'type': 'string'}}}, + 'application/json', + { + 'type': 'json_schema', + 'json_schema': { + 'type': 'object', + 'properties': {'name': {'type': 'string'}}, + }, + }, + ), + # Case 3: Only response_mime_type is provided + ( + None, + 'application/json', + {'type': 'json_object'}, + ), + ], +) +async def test_construct_payload_response_format( + client, + llm_request, + response_json_schema, + response_mime_type, + expected_response_format, +): + llm_request.config = types.GenerateContentConfig( + response_json_schema=response_json_schema, + response_mime_type=response_mime_type, + ) + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{'message': {'role': 'assistant', 'content': '{}'}}] + } + mock_response.status_code = 200 + + with mock.patch.object( + httpx.AsyncClient, 'post', return_value=mock_response + ) as mock_post: + _ = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + mock_post.assert_called_once() + payload = mock_post.call_args[1]['json'] + assert payload['response_format'] == expected_response_format + + +@pytest.mark.asyncio +async def test_generate_content_async_invalid_tool_call_type_raises_error( + client, llm_request +): + # Mock response with invalid tool call type + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'role': 'assistant', + 'content': None, + 'tool_calls': [{ + 'id': 'call_123', + # Invalid type + 'type': 'custom', + 'custom': { + 'name': 'read_string', + 'input': 'Hi! The this is a custom tool call!', + }, + }], + } + }] + } + mock_response.status_code = 200 + + with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response): + with pytest.raises(ValueError, match='Unsupported tool_call type: custom'): + _ = [ + r + async for r in client.generate_content_async( + llm_request, stream=False + ) + ] + + +@pytest.mark.asyncio +async def test_generate_content_async_function_call_response( + client, llm_request +): + # Mock response with deprecated function call + mock_response = AsyncMock(spec=httpx.Response) + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'role': 'assistant', + 'content': None, + 'function_call': { + 'name': 'get_weather', + 'arguments': '{"location": "London"}', + }, + } + }] + } + mock_response.status_code = 200 + + with mock.patch.object(httpx.AsyncClient, 'post', return_value=mock_response): + responses = [ + r + async for r in client.generate_content_async(llm_request, stream=False) + ] + + assert len(responses) == 1 + part = responses[0].content.parts[0] + assert part.function_call + assert part.function_call.name == 'get_weather' + assert part.function_call.args == {'location': 'London'} + assert part.function_call.id is None