encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add /chat/completions integration to ApigeeLlm

Browse Source 
PiperOrigin-RevId: 873049983
This commit is contained in:
Google Team Member
2026-02-20 13:27:22 -08:00
committed by Copybara-Service
parent 09ee3c3695
commit 9c4c445369
3 changed files with 1235 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/models/apigee_llm.py
+622 -12
View File
File diff suppressed because it is too large Load Diff
tests/unittests/models/test_apigee_llm.py
+173 -1
View File
@@ -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'
tests/unittests/models/test_completions_http_client.py
+440
View File
@@ -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