wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
832 lines
24 KiB
Python
832 lines
24 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import json
|
|
from unittest import mock
|
|
from unittest.mock import AsyncMock
|
|
|
|
from google.adk.models.apigee_llm import ChatCompletionsResponseHandler
|
|
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://localhost')
|
|
|
|
|
|
@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://localhost/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://localhost/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://localhost/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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_async_streaming_function_call():
|
|
local_client = CompletionsHTTPClient(base_url='https://localhost')
|
|
llm_request = LlmRequest(
|
|
model='apigee/test',
|
|
contents=[
|
|
types.Content(role='user', parts=[types.Part.from_text(text='hi')])
|
|
],
|
|
)
|
|
|
|
# Mock chunks simulating split arguments
|
|
chunk_data_0 = {
|
|
'id': 'chatcmpl-123',
|
|
'object': 'chat.completion.chunk',
|
|
'created': 1234567890,
|
|
'model': 'gpt-3.5-turbo',
|
|
'service_tier': 'default',
|
|
'choices': [{
|
|
'index': 0,
|
|
'delta': {
|
|
'tool_calls': [{
|
|
'index': 0,
|
|
'id': 'call_123',
|
|
'type': 'function',
|
|
'function': {'name': 'get_weather', 'arguments': ''},
|
|
}]
|
|
},
|
|
'finish_reason': None,
|
|
}],
|
|
}
|
|
chunk_data_1 = {
|
|
'id': 'chatcmpl-123',
|
|
'object': 'chat.completion.chunk',
|
|
'created': 1234567890,
|
|
'model': 'gpt-3.5-turbo',
|
|
'service_tier': 'default',
|
|
'choices': [{
|
|
'index': 0,
|
|
'delta': {
|
|
'tool_calls': [{
|
|
'index': 0,
|
|
'function': {'arguments': '{"location": "London"}'},
|
|
}]
|
|
},
|
|
'finish_reason': None,
|
|
}],
|
|
}
|
|
chunk_data_2 = {
|
|
'id': 'chatcmpl-123',
|
|
'object': 'chat.completion.chunk',
|
|
'created': 1234567890,
|
|
'model': 'gpt-3.5-turbo',
|
|
'service_tier': 'default',
|
|
'choices': [{
|
|
'index': 0,
|
|
'delta': {
|
|
'tool_calls': [{
|
|
'index': 0,
|
|
'function': {'arguments': '{"country": "UK"}'},
|
|
}]
|
|
},
|
|
'finish_reason': None,
|
|
}],
|
|
}
|
|
chunk_data_3 = {
|
|
'id': 'chatcmpl-123',
|
|
'object': 'chat.completion.chunk',
|
|
'created': 1234567890,
|
|
'model': 'gpt-3.5-turbo',
|
|
'service_tier': 'default',
|
|
'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'tool_calls'}],
|
|
'usage': {
|
|
'prompt_tokens': 10,
|
|
'completion_tokens': 20,
|
|
'total_tokens': 30,
|
|
},
|
|
}
|
|
|
|
chunks = [
|
|
f'{json.dumps(chunk_data_0)}\n',
|
|
f'{json.dumps(chunk_data_1)}\n',
|
|
f'{json.dumps(chunk_data_2)}\n',
|
|
f'{json.dumps(chunk_data_3)}\n',
|
|
]
|
|
|
|
async def mock_aiter_lines():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.aiter_lines.return_value = mock_aiter_lines()
|
|
mock_response.status_code = 200
|
|
|
|
mock_stream_ctx = mock.AsyncMock()
|
|
mock_stream_ctx.__aenter__.return_value = mock_response
|
|
|
|
with mock.patch.object(
|
|
httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
|
|
):
|
|
responses = [
|
|
r
|
|
async for r in local_client.generate_content_async(
|
|
llm_request, stream=True
|
|
)
|
|
]
|
|
# Check that we get 5 responses (one per chunk + extra final accumulated)
|
|
assert len(responses) == 5
|
|
|
|
# Check 1st response: partial tool call, empty args
|
|
assert responses[0].partial is True
|
|
assert responses[0].content.parts[0].function_call.name == 'get_weather'
|
|
assert responses[0].content.parts[0].function_call.id == 'call_123'
|
|
|
|
# Check 2nd response: full args for first update
|
|
assert responses[1].partial is True
|
|
assert responses[1].content.parts[0].function_call.args == {
|
|
'location': 'London'
|
|
}
|
|
|
|
# Check 3rd response: full args for second update (merged)
|
|
assert responses[2].partial is True
|
|
assert responses[2].content.parts[0].function_call.args == {'country': 'UK'}
|
|
|
|
# Check 4th response: last delta (empty)
|
|
assert responses[3].partial is True
|
|
assert responses[3].content.parts == []
|
|
|
|
# Check 5th response: final accumulated
|
|
assert responses[4].finish_reason == types.FinishReason.STOP
|
|
# Full accumulated args
|
|
assert responses[4].content.parts[0].function_call.args == {
|
|
'location': 'London',
|
|
'country': 'UK',
|
|
}
|
|
|
|
# Check metadata and usage
|
|
assert responses[4].model_version == 'gpt-3.5-turbo'
|
|
assert responses[4].custom_metadata['id'] == 'chatcmpl-123'
|
|
assert responses[4].custom_metadata['created'], 1234567890
|
|
assert responses[4].custom_metadata['object'], 'chat.completion.chunk'
|
|
assert responses[4].custom_metadata['service_tier'], 'default'
|
|
assert responses[4].usage_metadata is not None
|
|
assert responses[4].usage_metadata.prompt_token_count == 10
|
|
assert responses[4].usage_metadata.candidates_token_count == 20
|
|
assert responses[4].usage_metadata.total_token_count == 30
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_async_streaming_multiple_function_calls():
|
|
# Mock streaming response with multiple tool calls
|
|
local_client = CompletionsHTTPClient(base_url='https://localhost')
|
|
llm_request = LlmRequest(
|
|
model='apigee/test',
|
|
contents=[
|
|
types.Content(role='user', parts=[types.Part.from_text(text='hi')])
|
|
],
|
|
)
|
|
chunk_data_1 = {
|
|
'choices': [{
|
|
'index': 0,
|
|
'delta': {
|
|
'tool_calls': [
|
|
{
|
|
'index': 0,
|
|
'id': 'call_1',
|
|
'type': 'function',
|
|
'function': {'name': 'func_1', 'arguments': ''},
|
|
},
|
|
{
|
|
'index': 1,
|
|
'id': 'call_2',
|
|
'type': 'function',
|
|
'function': {'name': 'func_2', 'arguments': ''},
|
|
},
|
|
]
|
|
},
|
|
'finish_reason': None,
|
|
}]
|
|
}
|
|
# the tool_call type is optional in chunk responses.
|
|
chunk_data_2 = {
|
|
'choices': [{
|
|
'index': 0,
|
|
'delta': {
|
|
'tool_calls': [
|
|
{'index': 0, 'function': {'arguments': '{"arg": 1}'}},
|
|
{'index': 1, 'function': {'arguments': '{"arg": 2}'}},
|
|
]
|
|
},
|
|
'finish_reason': None,
|
|
}]
|
|
}
|
|
chunk_data_3 = {
|
|
'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'tool_calls'}]
|
|
}
|
|
|
|
chunks = [
|
|
f'{json.dumps(chunk_data_1)}\n',
|
|
f'{json.dumps(chunk_data_2)}\n',
|
|
f'{json.dumps(chunk_data_3)}\n',
|
|
]
|
|
|
|
async def mock_aiter_lines():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.aiter_lines.return_value = mock_aiter_lines()
|
|
mock_response.status_code = 200
|
|
|
|
mock_stream_ctx = mock.AsyncMock()
|
|
mock_stream_ctx.__aenter__.return_value = mock_response
|
|
|
|
with mock.patch.object(
|
|
httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
|
|
):
|
|
responses = [
|
|
r
|
|
async for r in local_client.generate_content_async(
|
|
llm_request, stream=True
|
|
)
|
|
]
|
|
|
|
assert len(responses) == 4
|
|
parts = responses[-1].content.parts
|
|
assert len(parts) == 2
|
|
|
|
assert parts[0].function_call.name == 'func_1'
|
|
assert parts[0].function_call.args == {'arg': 1}
|
|
assert parts[0].function_call.id == 'call_1'
|
|
|
|
assert parts[1].function_call.name == 'func_2'
|
|
assert parts[1].function_call.args == {'arg': 2}
|
|
|
|
assert parts[1].function_call.id == 'call_2'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
('chunks', 'expected_response_count'),
|
|
[
|
|
(
|
|
[
|
|
'\n',
|
|
' \n',
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "Hello"}, "finish_reason": null}]}\n'
|
|
),
|
|
],
|
|
1,
|
|
),
|
|
(
|
|
[
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "Hello"}, "finish_reason": null}]}\n'
|
|
),
|
|
'[DONE]\n',
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "World"}, "finish_reason": "stop"}]}\n'
|
|
),
|
|
],
|
|
1, # Should stop after [DONE]
|
|
),
|
|
(
|
|
[
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "Hello"}, "finish_reason": null}]}\n'
|
|
),
|
|
' [DONE] \n',
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "World"}, "finish_reason": "stop"}]}\n'
|
|
),
|
|
],
|
|
1, # Should stop after [DONE]
|
|
),
|
|
(
|
|
[
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "Hello"}, "finish_reason": null}]}\n'
|
|
),
|
|
'data: [DONE]\n',
|
|
(
|
|
'data: {"choices": [{"index": 0, "delta": {"content":'
|
|
' "World"}, "finish_reason": "stop"}]}\n'
|
|
),
|
|
],
|
|
1, # Should stop after [DONE]
|
|
),
|
|
],
|
|
)
|
|
async def test_generate_content_async_streaming_parse_lines(
|
|
chunks, expected_response_count
|
|
):
|
|
local_client = CompletionsHTTPClient(base_url='https://localhost')
|
|
llm_request = LlmRequest(
|
|
model='apigee/test',
|
|
contents=[
|
|
types.Content(role='user', parts=[types.Part.from_text(text='hi')])
|
|
],
|
|
)
|
|
|
|
async def mock_aiter_lines():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.aiter_lines.return_value = mock_aiter_lines()
|
|
mock_response.status_code = 200
|
|
|
|
mock_stream_ctx = mock.AsyncMock()
|
|
mock_stream_ctx.__aenter__.return_value = mock_response
|
|
|
|
with mock.patch.object(
|
|
httpx.AsyncClient, 'stream', return_value=mock_stream_ctx
|
|
):
|
|
responses = [
|
|
r
|
|
async for r in local_client.generate_content_async(
|
|
llm_request, stream=True
|
|
)
|
|
]
|
|
assert len(responses) == expected_response_count
|
|
assert responses[0].content.parts[0].text == 'Hello'
|
|
|
|
|
|
def test_process_chunk_with_refusal_streaming():
|
|
handler = ChatCompletionsResponseHandler()
|
|
|
|
chunk1 = {
|
|
'choices': [{
|
|
'delta': {
|
|
'role': 'assistant',
|
|
'content': 'Hello',
|
|
},
|
|
'index': 0,
|
|
}]
|
|
}
|
|
responses1 = list(handler.process_chunk(chunk1))
|
|
assert len(responses1) == 1
|
|
assert responses1[0].content.parts[0].text == 'Hello'
|
|
|
|
chunk2 = {
|
|
'choices': [{
|
|
'delta': {
|
|
'refusal': 'I refuse',
|
|
},
|
|
'index': 0,
|
|
}]
|
|
}
|
|
responses2 = list(handler.process_chunk(chunk2))
|
|
assert len(responses2) == 1
|
|
assert responses2[0].content.parts[0].text == '\n[[REFUSAL]]: I refuse'
|
|
|
|
chunk3 = {
|
|
'choices': [{
|
|
'delta': {
|
|
'refusal': ' to answer',
|
|
},
|
|
'index': 0,
|
|
}]
|
|
}
|
|
responses3 = list(handler.process_chunk(chunk3))
|
|
assert len(responses3) == 1
|
|
assert responses3[0].content.parts[0].text == ' to answer'
|
|
|
|
chunk4 = {
|
|
'choices': [{
|
|
'delta': {},
|
|
'finish_reason': 'stop',
|
|
'index': 0,
|
|
}]
|
|
}
|
|
responses4 = list(handler.process_chunk(chunk4))
|
|
assert len(responses4) == 2
|
|
final_response = responses4[1]
|
|
assert final_response.finish_reason == types.FinishReason.STOP
|
|
assert (
|
|
final_response.content.parts[0].text
|
|
== 'Hello\n[[REFUSAL]]: I refuse to answer'
|
|
)
|