wehub-resource-sync
76d991c447
Auto Update PR / update-prs (push) Has been cancelled
CI / format-check (push) Has been cancelled
CI / test (3.10) (push) Has been cancelled
CI / test (3.11) (push) Has been cancelled
CI / test (3.12) (push) Has been cancelled
CI / live-api-tests (push) Has been cancelled
CI / plugin-integration-test (push) Has been cancelled
CI / ollama-integration-test (push) Has been cancelled
CI / test-fork-pr (push) Has been cancelled
1061 lines
35 KiB
Python
1061 lines
35 KiB
Python
# Copyright 2025 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.
|
|
|
|
"""Tests for inference module.
|
|
|
|
Note: This file contains test helper classes that intentionally have
|
|
few public methods and define attributes outside __init__. These
|
|
pylint warnings are expected for test fixtures.
|
|
"""
|
|
# pylint: disable=attribute-defined-outside-init
|
|
|
|
from unittest import mock
|
|
|
|
from absl.testing import absltest
|
|
from absl.testing import parameterized
|
|
|
|
from langextract import exceptions
|
|
from langextract.core import base_model
|
|
from langextract.core import data
|
|
from langextract.core import types
|
|
from langextract.providers import gemini
|
|
from langextract.providers import ollama
|
|
from langextract.providers import openai
|
|
|
|
|
|
class TestBaseLanguageModel(absltest.TestCase):
|
|
|
|
def test_merge_kwargs_with_none(self):
|
|
"""Test merge_kwargs handles None runtime_kwargs."""
|
|
|
|
class TestModel(base_model.BaseLanguageModel): # pylint: disable=too-few-public-methods
|
|
|
|
def infer(self, batch_prompts, **kwargs):
|
|
return iter([])
|
|
|
|
model = TestModel()
|
|
model._extra_kwargs = {"a": 1, "b": 2}
|
|
|
|
result = model.merge_kwargs(None)
|
|
self.assertEqual(
|
|
{"a": 1, "b": 2},
|
|
result,
|
|
"merge_kwargs(None) should return stored kwargs unchanged",
|
|
)
|
|
|
|
result = model.merge_kwargs({})
|
|
self.assertEqual(
|
|
{"a": 1, "b": 2},
|
|
result,
|
|
"merge_kwargs({}) should return stored kwargs unchanged",
|
|
)
|
|
|
|
result = model.merge_kwargs({"b": 3, "c": 4})
|
|
self.assertEqual(
|
|
{"a": 1, "b": 3, "c": 4},
|
|
result,
|
|
"Runtime kwargs should override stored kwargs and add new keys",
|
|
)
|
|
|
|
def test_merge_kwargs_without_extra_kwargs(self):
|
|
"""Test merge_kwargs when _extra_kwargs doesn't exist."""
|
|
|
|
class TestModel(base_model.BaseLanguageModel): # pylint: disable=too-few-public-methods
|
|
|
|
def infer(self, batch_prompts, **kwargs):
|
|
return iter([])
|
|
|
|
model = TestModel()
|
|
# Intentionally not setting _extra_kwargs to test fallback behavior
|
|
|
|
result = model.merge_kwargs({"a": 1})
|
|
self.assertEqual(
|
|
{"a": 1},
|
|
result,
|
|
"merge_kwargs should work even without _extra_kwargs attribute",
|
|
)
|
|
|
|
|
|
class TestOllamaLanguageModel(absltest.TestCase):
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer(self, mock_ollama_query):
|
|
|
|
# Real gemma2 response structure from Ollama API for validation
|
|
gemma_response = {
|
|
"model": "gemma2:latest",
|
|
"created_at": "2025-01-23T22:37:08.579440841Z",
|
|
"response": "{'bus' : '**autóbusz**'} \n\n\n \n",
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
"context": [
|
|
106,
|
|
1645,
|
|
108,
|
|
1841,
|
|
603,
|
|
1986,
|
|
575,
|
|
59672,
|
|
235336,
|
|
107,
|
|
108,
|
|
106,
|
|
2516,
|
|
108,
|
|
9766,
|
|
6710,
|
|
235281,
|
|
865,
|
|
664,
|
|
688,
|
|
7958,
|
|
235360,
|
|
6710,
|
|
235306,
|
|
688,
|
|
12990,
|
|
235248,
|
|
110,
|
|
139,
|
|
108,
|
|
],
|
|
"total_duration": 24038204381,
|
|
"load_duration": 21551375738,
|
|
"prompt_eval_count": 15,
|
|
"prompt_eval_duration": 633000000,
|
|
"eval_count": 17,
|
|
"eval_duration": 1848000000,
|
|
}
|
|
mock_ollama_query.return_value = gemma_response
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="gemma2:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
)
|
|
batch_prompts = ["What is bus in Hungarian?"]
|
|
results = list(model.infer(batch_prompts))
|
|
|
|
mock_ollama_query.assert_called_once_with(
|
|
prompt="What is bus in Hungarian?",
|
|
model="gemma2:latest",
|
|
structured_output_format="json",
|
|
model_url="http://localhost:11434",
|
|
think=False,
|
|
)
|
|
expected_results = [[
|
|
types.ScoredOutput(
|
|
score=1.0, output="{'bus' : '**autóbusz**'} \n\n\n \n"
|
|
)
|
|
]]
|
|
self.assertEqual(results, expected_results)
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer_prefers_response_over_thinking(self, mock_ollama_query):
|
|
"""Test Ollama inference ignores reasoning when final output is present."""
|
|
thinking_response = {
|
|
"model": "deepseek-r1:latest",
|
|
"created_at": "2025-01-23T22:37:08.579440841Z",
|
|
"response": "{'bus' : '**autóbusz**'} \n\n\n \n",
|
|
"thinking": "The prompt asks for a Hungarian translation of bus.",
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
"context": [106, 1645, 108],
|
|
"total_duration": 24038204381,
|
|
"load_duration": 21551375738,
|
|
"prompt_eval_count": 15,
|
|
"prompt_eval_duration": 633000000,
|
|
"eval_count": 17,
|
|
"eval_duration": 1848000000,
|
|
}
|
|
mock_ollama_query.return_value = thinking_response
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="deepseek-r1:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
)
|
|
batch_prompts = ["What is bus in Hungarian?"]
|
|
results = list(model.infer(batch_prompts))
|
|
|
|
mock_ollama_query.assert_called_once_with(
|
|
prompt="What is bus in Hungarian?",
|
|
model="deepseek-r1:latest",
|
|
structured_output_format="json",
|
|
model_url="http://localhost:11434",
|
|
think=False,
|
|
)
|
|
expected_results = [[
|
|
types.ScoredOutput(
|
|
score=1.0, output="{'bus' : '**autóbusz**'} \n\n\n \n"
|
|
)
|
|
]]
|
|
self.assertEqual(results, expected_results)
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer_raises_on_empty_response_with_thinking(
|
|
self, mock_ollama_query
|
|
):
|
|
"""Test Ollama inference does not treat reasoning as final output."""
|
|
mock_ollama_query.return_value = {
|
|
"model": "deepseek-r1:latest",
|
|
"created_at": "2025-01-23T22:37:08.579440841Z",
|
|
"response": "",
|
|
"thinking": "The prompt asks for a Hungarian translation of bus.",
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
}
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="deepseek-r1:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceRuntimeError, "think=False"
|
|
):
|
|
list(model.infer(["What is bus in Hungarian?"]))
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer_raises_when_response_missing(self, mock_ollama_query):
|
|
"""Test Ollama inference requires final generated text."""
|
|
mock_ollama_query.return_value = {
|
|
"model": "deepseek-r1:latest",
|
|
"created_at": "2025-01-23T22:37:08.579440841Z",
|
|
"done": True,
|
|
"done_reason": "stop",
|
|
}
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="deepseek-r1:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceRuntimeError, "response.*field"
|
|
):
|
|
list(model.infer(["What is bus in Hungarian?"]))
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer_preserves_explicit_think(self, mock_ollama_query):
|
|
"""Test user-provided think setting is passed through."""
|
|
mock_ollama_query.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="deepseek-r1:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
)
|
|
|
|
list(model.infer(["Test prompt"], think=True))
|
|
|
|
mock_ollama_query.assert_called_once_with(
|
|
prompt="Test prompt",
|
|
model="deepseek-r1:latest",
|
|
structured_output_format="json",
|
|
model_url="http://localhost:11434",
|
|
think=True,
|
|
)
|
|
|
|
@mock.patch("langextract.providers.ollama.OllamaLanguageModel._ollama_query")
|
|
def test_ollama_infer_default_think_does_not_mutate_stored_kwargs(
|
|
self, mock_ollama_query
|
|
):
|
|
"""Test default think setting is per request only."""
|
|
mock_ollama_query.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="deepseek-r1:latest",
|
|
model_url="http://localhost:11434",
|
|
structured_output_format="json",
|
|
temperature=0.1,
|
|
)
|
|
stored_kwargs = model.merge_kwargs()
|
|
|
|
list(model.infer(["Test prompt"]))
|
|
|
|
self.assertEqual(model.merge_kwargs(), stored_kwargs)
|
|
self.assertNotIn("think", model.merge_kwargs())
|
|
|
|
def test_ollama_gpt_oss_model_matching(self):
|
|
for model_id in ("gpt-oss", "gpt-oss:20b", "GPT-OSS:20B"):
|
|
with self.subTest(model_id=model_id):
|
|
self.assertTrue(ollama._is_gpt_oss_model(model_id))
|
|
|
|
for model_id in (
|
|
"gpt-oss-120b",
|
|
"not-gpt-oss:20b",
|
|
"openai/gpt-oss:20b",
|
|
"gpt-oss:",
|
|
):
|
|
with self.subTest(model_id=model_id):
|
|
self.assertFalse(ollama._is_gpt_oss_model(model_id))
|
|
|
|
def test_ollama_chat_empty_content_with_thinking_raises(self):
|
|
with self.assertRaisesRegex(exceptions.InferenceRuntimeError, "think=True"):
|
|
ollama.OllamaLanguageModel._extract_chat_response_text(
|
|
{"message": {"content": "", "thinking": "reasoning"}}
|
|
)
|
|
|
|
def test_ollama_chat_missing_content_raises(self):
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceRuntimeError, "message.content"
|
|
):
|
|
ollama.OllamaLanguageModel._extract_chat_response_text({"done": True})
|
|
|
|
@mock.patch.object(ollama.requests, "post", autospec=True)
|
|
def test_ollama_gpt_oss_yaml_uses_generate_path(self, mock_post):
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": "extractions: []",
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="gpt-oss:20b",
|
|
model_url="http://localhost:11434",
|
|
format_type=types.FormatType.YAML,
|
|
)
|
|
|
|
results = list(model.infer(["Test prompt"]))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
payload = call_args.kwargs["json"]
|
|
self.assertEqual(call_args.args[0], "http://localhost:11434/api/generate")
|
|
self.assertEqual(payload["format"], "yaml")
|
|
self.assertNotIn("messages", payload)
|
|
self.assertEqual(
|
|
results,
|
|
[[types.ScoredOutput(score=1.0, output="extractions: []")]],
|
|
)
|
|
|
|
@mock.patch.object(ollama.requests, "post", autospec=True)
|
|
def test_ollama_gpt_oss_uses_chat_without_native_json(self, mock_post):
|
|
"""GPT-OSS avoids native JSON mode, which conflicts with Harmony format."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"message": {"content": '{"extractions": []}'},
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="gpt-oss:20b",
|
|
model_url="http://localhost:11434",
|
|
)
|
|
|
|
results = list(model.infer(["Test prompt"], temperature=0.0))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
self.assertEqual(call_args.args[0], "http://localhost:11434/api/chat")
|
|
self.assertDictEqual(
|
|
call_args.kwargs["json"],
|
|
{
|
|
"model": "gpt-oss:20b",
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"Output a single JSON object matching the requested "
|
|
"extraction format. Do not include code fences, prose, "
|
|
"or reasoning."
|
|
),
|
|
},
|
|
{"role": "user", "content": "Test prompt"},
|
|
],
|
|
"stream": False,
|
|
"options": {
|
|
"keep_alive": 300,
|
|
"temperature": 0.0,
|
|
"num_ctx": 2048,
|
|
},
|
|
"keep_alive": 300,
|
|
"think": False,
|
|
},
|
|
)
|
|
self.assertEqual(
|
|
results,
|
|
[[types.ScoredOutput(score=1.0, output='{"extractions": []}')]],
|
|
)
|
|
|
|
@mock.patch.object(ollama.requests, "post", autospec=True)
|
|
def test_ollama_extra_kwargs_passed_to_api(self, mock_post):
|
|
"""Verify extra kwargs like timeout and keep_alive are passed to the API."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
timeout=300,
|
|
keep_alive=600,
|
|
num_threads=8,
|
|
)
|
|
|
|
prompts = ["Test prompt"]
|
|
list(model.infer(prompts))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
json_payload = call_args.kwargs["json"]
|
|
|
|
self.assertEqual(call_args.args[0], "http://localhost:11434/api/generate")
|
|
self.assertEqual(json_payload["format"], "json")
|
|
self.assertNotIn("messages", json_payload)
|
|
self.assertEqual(json_payload["keep_alive"], 600)
|
|
self.assertIs(json_payload["think"], False)
|
|
self.assertNotIn("think", json_payload["options"])
|
|
self.assertEqual(json_payload["options"]["keep_alive"], 600)
|
|
self.assertEqual(json_payload["options"]["num_thread"], 8)
|
|
# timeout is passed to requests.post, not in the JSON payload
|
|
self.assertEqual(call_args.kwargs["timeout"], 300)
|
|
|
|
@mock.patch("requests.post")
|
|
def test_ollama_stop_and_top_p_passthrough(self, mock_post):
|
|
"""Verify stop and top_p parameters are passed to Ollama API."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
top_p=0.9,
|
|
stop=["\\n\\n", "END"],
|
|
)
|
|
|
|
prompts = ["Test prompt"]
|
|
list(model.infer(prompts))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
json_payload = call_args.kwargs["json"]
|
|
|
|
# Ollama expects 'stop' at top level, not in options
|
|
self.assertEqual(json_payload["stop"], ["\\n\\n", "END"])
|
|
self.assertEqual(json_payload["options"]["top_p"], 0.9)
|
|
|
|
@mock.patch("requests.post")
|
|
def test_ollama_defaults_when_unspecified(self, mock_post):
|
|
"""Verify Ollama uses correct defaults when parameters are not specified."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(model_id="test-model")
|
|
|
|
prompts = ["Test prompt"]
|
|
list(model.infer(prompts))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
json_payload = call_args.kwargs["json"]
|
|
|
|
self.assertEqual(json_payload["keep_alive"], 300)
|
|
self.assertEqual(json_payload["options"]["temperature"], 0.1)
|
|
self.assertEqual(json_payload["options"]["keep_alive"], 300)
|
|
self.assertEqual(json_payload["options"]["num_ctx"], 2048)
|
|
self.assertEqual(call_args.kwargs["timeout"], 120)
|
|
|
|
@mock.patch("requests.post")
|
|
def test_ollama_runtime_kwargs_override_stored(self, mock_post):
|
|
"""Verify runtime kwargs override stored kwargs."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
temperature=0.5,
|
|
keep_alive=300,
|
|
)
|
|
|
|
prompts = ["Test prompt"]
|
|
list(model.infer(prompts, temperature=0.8, keep_alive=600))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
json_payload = call_args.kwargs["json"]
|
|
|
|
self.assertEqual(json_payload["options"]["temperature"], 0.8)
|
|
self.assertEqual(json_payload["options"]["keep_alive"], 600)
|
|
|
|
@mock.patch("requests.post")
|
|
def test_ollama_temperature_zero(self, mock_post):
|
|
"""Test that temperature=0.0 is properly passed to Ollama."""
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"response": '{"test": "value"}',
|
|
"done": True,
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
temperature=0.0,
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
mock_post.assert_called_once()
|
|
call_args = mock_post.call_args
|
|
json_payload = call_args.kwargs["json"]
|
|
|
|
self.assertEqual(json_payload["options"]["temperature"], 0.0)
|
|
|
|
def test_ollama_default_timeout(self):
|
|
"""Test that default timeout is used when not specified."""
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
model_url="http://localhost:11434",
|
|
)
|
|
|
|
mock_response = mock.Mock(spec=["status_code", "json"])
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {"response": "test output"}
|
|
|
|
with mock.patch.object(
|
|
model._requests, "post", return_value=mock_response
|
|
) as mock_post:
|
|
model._ollama_query(prompt="test prompt")
|
|
|
|
mock_post.assert_called_once()
|
|
call_kwargs = mock_post.call_args[1]
|
|
self.assertEqual(
|
|
120,
|
|
call_kwargs["timeout"],
|
|
"Should use default timeout of 120 seconds",
|
|
)
|
|
|
|
def test_ollama_timeout_through_infer(self):
|
|
"""Test that timeout flows correctly through the infer() method."""
|
|
model = ollama.OllamaLanguageModel(
|
|
model_id="test-model",
|
|
model_url="http://localhost:11434",
|
|
timeout=60,
|
|
)
|
|
|
|
mock_response = mock.Mock(spec=["status_code", "json"])
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {"response": "test output"}
|
|
|
|
with mock.patch.object(
|
|
model._requests, "post", return_value=mock_response
|
|
) as mock_post:
|
|
list(model.infer(["test prompt"]))
|
|
|
|
mock_post.assert_called_once()
|
|
call_kwargs = mock_post.call_args[1]
|
|
self.assertEqual(
|
|
60,
|
|
call_kwargs["timeout"],
|
|
"Timeout from constructor should flow through infer()",
|
|
)
|
|
|
|
|
|
class TestGeminiLanguageModel(absltest.TestCase):
|
|
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_allowlist_filtering(self, mock_client_class):
|
|
"""Test that only allow-listed keys are passed through."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.text = '{"result": "test"}'
|
|
mock_client.models.generate_content.return_value = mock_response
|
|
|
|
model = gemini.GeminiLanguageModel(
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test-key",
|
|
# Allow-listed parameters
|
|
tools=["tool1", "tool2"],
|
|
stop_sequences=["\n\n"],
|
|
system_instruction="Be helpful",
|
|
# Unknown parameters to test filtering
|
|
unknown_param="should_be_ignored",
|
|
another_unknown="also_ignored",
|
|
)
|
|
|
|
expected_extra_kwargs = {
|
|
"tools": ["tool1", "tool2"],
|
|
"stop_sequences": ["\n\n"],
|
|
"system_instruction": "Be helpful",
|
|
}
|
|
self.assertEqual(
|
|
expected_extra_kwargs,
|
|
model._extra_kwargs,
|
|
"Only allow-listed kwargs should be stored in _extra_kwargs",
|
|
)
|
|
|
|
prompts = ["Test prompt"]
|
|
list(model.infer(prompts))
|
|
|
|
mock_client.models.generate_content.assert_called_once()
|
|
call_args = mock_client.models.generate_content.call_args
|
|
config = call_args.kwargs["config"]
|
|
|
|
for key in ["tools", "stop_sequences", "system_instruction"]:
|
|
self.assertIn(key, config, f"Expected {key} to be in API config")
|
|
self.assertEqual(
|
|
expected_extra_kwargs[key],
|
|
config[key],
|
|
f"Config value for {key} should match what was provided",
|
|
)
|
|
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_runtime_kwargs_filtered(self, mock_client_class):
|
|
"""Test that runtime kwargs are also filtered by allow-list."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.text = '{"result": "test"}'
|
|
mock_client.models.generate_content.return_value = mock_response
|
|
|
|
model = gemini.GeminiLanguageModel(
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test-key",
|
|
)
|
|
|
|
prompts = ["Test prompt"]
|
|
list(
|
|
model.infer(
|
|
prompts,
|
|
candidate_count=2,
|
|
safety_settings={"HARM_CATEGORY_DANGEROUS": "BLOCK_NONE"},
|
|
unknown_runtime_param="ignored",
|
|
)
|
|
)
|
|
|
|
call_args = mock_client.models.generate_content.call_args
|
|
config = call_args.kwargs["config"]
|
|
|
|
self.assertEqual(
|
|
2,
|
|
config.get("candidate_count"),
|
|
"candidate_count should be passed through to API",
|
|
)
|
|
self.assertEqual(
|
|
{"HARM_CATEGORY_DANGEROUS": "BLOCK_NONE"},
|
|
config.get("safety_settings"),
|
|
"safety_settings should be passed through to API",
|
|
)
|
|
self.assertNotIn(
|
|
"unknown_runtime_param", config, "Unknown kwargs should be filtered out"
|
|
)
|
|
|
|
def test_gemini_requires_auth_config(self):
|
|
"""Test that Gemini requires either API key or Vertex AI config."""
|
|
with self.assertRaises(exceptions.InferenceConfigError) as cm:
|
|
gemini.GeminiLanguageModel()
|
|
|
|
self.assertIn("Gemini models require either", str(cm.exception))
|
|
self.assertIn("API key", str(cm.exception))
|
|
self.assertIn("Vertex AI", str(cm.exception))
|
|
|
|
def test_gemini_vertexai_requires_project_and_location(self):
|
|
"""Test that Vertex AI mode requires both project and location."""
|
|
with self.assertRaises(exceptions.InferenceConfigError) as cm:
|
|
gemini.GeminiLanguageModel(vertexai=True)
|
|
|
|
self.assertIn("requires both project and location", str(cm.exception))
|
|
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_vertexai_initialization(self, mock_client_class):
|
|
"""Test successful initialization with Vertex AI config."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
model = gemini.GeminiLanguageModel(
|
|
vertexai=True, project="test-project", location="us-central1"
|
|
)
|
|
|
|
self.assertIsNone(model.api_key)
|
|
self.assertTrue(model.vertexai)
|
|
self.assertEqual(model.project, "test-project")
|
|
self.assertEqual(model.location, "us-central1")
|
|
mock_client_class.assert_called_once_with(
|
|
api_key=None,
|
|
vertexai=True,
|
|
credentials=None,
|
|
project="test-project",
|
|
location="us-central1",
|
|
http_options=None,
|
|
)
|
|
|
|
@mock.patch("absl.logging.warning")
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_warns_when_both_auth_provided(
|
|
self, mock_client_class, mock_warning
|
|
):
|
|
"""Test that warning is logged when both API key and Vertex AI are provided."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
gemini.GeminiLanguageModel(
|
|
api_key="test-key",
|
|
vertexai=True,
|
|
project="test-project",
|
|
location="us-central1",
|
|
)
|
|
|
|
mock_warning.assert_called_once()
|
|
warning_msg = mock_warning.call_args[0][0]
|
|
self.assertIn("Both API key and Vertex AI", warning_msg)
|
|
self.assertIn("API key will take precedence", warning_msg)
|
|
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_vertexai_with_http_options(self, mock_client_class):
|
|
"""Test that http_options are passed to genai.Client for VPC endpoints."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
http_options = {"base_url": "https://custom-vpc.p.googleapis.com"}
|
|
model = gemini.GeminiLanguageModel(
|
|
vertexai=True,
|
|
project="test-project",
|
|
location="us-central1",
|
|
http_options=http_options,
|
|
)
|
|
|
|
self.assertEqual(model.http_options, http_options)
|
|
mock_client_class.assert_called_once_with(
|
|
api_key=None,
|
|
vertexai=True,
|
|
credentials=None,
|
|
project="test-project",
|
|
location="us-central1",
|
|
http_options=http_options,
|
|
)
|
|
|
|
|
|
class TestOpenAILanguageModelInference(parameterized.TestCase):
|
|
|
|
@parameterized.named_parameters(
|
|
("without", "test-api-key", None, "gpt-4o-mini", 0.5),
|
|
("with", "test-api-key", "http://127.0.0.1:9001/v1", "gpt-4o-mini", 0.5),
|
|
)
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_infer_with_parameters(
|
|
self, api_key, base_url, model_id, temperature, mock_openai_class
|
|
):
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"name": "John", "age": 30}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
model_id=model_id,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
temperature=temperature,
|
|
)
|
|
|
|
batch_prompts = ["Extract name and age from: John is 30 years old"]
|
|
results = list(model.infer(batch_prompts))
|
|
|
|
# JSON format adds a system message; only explicitly set params are passed
|
|
mock_client.chat.completions.create.assert_called_once()
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(call_args.kwargs["model"], "gpt-4o-mini")
|
|
self.assertEqual(call_args.kwargs["temperature"], temperature)
|
|
self.assertEqual(call_args.kwargs["n"], 1)
|
|
self.assertEqual(len(call_args.kwargs["messages"]), 2)
|
|
self.assertEqual(call_args.kwargs["messages"][0]["role"], "system")
|
|
self.assertEqual(call_args.kwargs["messages"][1]["role"], "user")
|
|
|
|
expected_results = [
|
|
[types.ScoredOutput(score=1.0, output='{"name": "John", "age": 30}')]
|
|
]
|
|
self.assertEqual(results, expected_results)
|
|
|
|
|
|
class TestOpenAILanguageModel(absltest.TestCase):
|
|
|
|
def test_openai_parse_output_json(self):
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key", format_type=data.FormatType.JSON
|
|
)
|
|
|
|
output = '{"key": "value", "number": 42}'
|
|
parsed = model.parse_output(output)
|
|
self.assertEqual(parsed, {"key": "value", "number": 42})
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
model.parse_output("invalid json")
|
|
self.assertIn("Failed to parse output as JSON", str(context.exception))
|
|
|
|
def test_openai_parse_output_yaml(self):
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key", format_type=data.FormatType.YAML
|
|
)
|
|
|
|
output = "key: value\nnumber: 42"
|
|
parsed = model.parse_output(output)
|
|
self.assertEqual(parsed, {"key": "value", "number": 42})
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
model.parse_output("invalid: yaml: bad")
|
|
self.assertIn("Failed to parse output as YAML", str(context.exception))
|
|
|
|
def test_openai_no_api_key_raises_error(self):
|
|
with self.assertRaises(exceptions.InferenceConfigError) as context:
|
|
openai.OpenAILanguageModel(api_key=None)
|
|
self.assertEqual(str(context.exception), "API key not provided.")
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_extra_kwargs_passed(self, mock_openai_class):
|
|
"""Test that extra kwargs are passed to OpenAI API."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
frequency_penalty=0.5,
|
|
presence_penalty=0.7,
|
|
seed=42,
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(call_args.kwargs["frequency_penalty"], 0.5)
|
|
self.assertEqual(call_args.kwargs["presence_penalty"], 0.7)
|
|
self.assertEqual(call_args.kwargs["seed"], 42)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_runtime_kwargs_override(self, mock_openai_class):
|
|
"""Test that runtime kwargs override stored kwargs."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
temperature=0.5,
|
|
seed=123,
|
|
)
|
|
|
|
list(model.infer(["test prompt"], temperature=0.8, seed=456))
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(call_args.kwargs["temperature"], 0.8)
|
|
self.assertEqual(call_args.kwargs["seed"], 456)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_json_response_format(self, mock_openai_class):
|
|
"""Test that JSON format adds response_format parameter."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key", format_type=data.FormatType.JSON
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
mock_client.chat.completions.create.assert_called_once()
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(
|
|
call_args.kwargs["response_format"], {"type": "json_object"}
|
|
)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_temperature_zero(self, mock_openai_class):
|
|
"""Verify temperature=0.0 is properly passed to the API."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(api_key="test-key", temperature=0.0)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
mock_client.chat.completions.create.assert_called_once()
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(call_args.kwargs["temperature"], 0.0)
|
|
self.assertEqual(call_args.kwargs["model"], "gpt-4o-mini")
|
|
self.assertEqual(call_args.kwargs["n"], 1)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_temperature_none_not_sent(self, mock_openai_class):
|
|
"""Test that temperature=None is not sent to the API."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
# Test with temperature=None in model init
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
temperature=None,
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertNotIn("temperature", call_args.kwargs)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_none_values_filtered(self, mock_openai_class):
|
|
"""Test that None values are not passed to the API."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
top_p=0.9,
|
|
)
|
|
|
|
list(model.infer(["test prompt"], top_p=None, seed=None))
|
|
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertNotIn("top_p", call_args.kwargs)
|
|
self.assertNotIn("seed", call_args.kwargs)
|
|
|
|
@mock.patch("openai.OpenAI")
|
|
def test_openai_no_system_message_when_not_json_yaml(self, mock_openai_class):
|
|
"""Test that no system message is sent when format_type is not JSON/YAML."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content="test output"))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
format_type=None,
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
messages = call_args.kwargs["messages"]
|
|
|
|
self.assertEqual(len(messages), 1)
|
|
self.assertEqual(messages[0]["role"], "user")
|
|
self.assertEqual(messages[0]["content"], "test prompt")
|
|
|
|
@mock.patch("openai.OpenAI", autospec=True)
|
|
def test_openai_reasoning_effort_passed_directly(self, mock_openai_class):
|
|
"""reasoning_effort is passed as a top-level API parameter."""
|
|
mock_client = mock.Mock()
|
|
mock_openai_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.choices = [
|
|
mock.Mock(message=mock.Mock(content='{"result": "test"}'))
|
|
]
|
|
mock_client.chat.completions.create.return_value = mock_response
|
|
|
|
model = openai.OpenAILanguageModel(
|
|
api_key="test-key",
|
|
reasoning_effort="low",
|
|
)
|
|
|
|
list(model.infer(["test prompt"]))
|
|
|
|
call_args = mock_client.chat.completions.create.call_args
|
|
self.assertEqual(call_args.kwargs["reasoning_effort"], "low")
|
|
self.assertNotIn("reasoning", call_args.kwargs)
|
|
|
|
@mock.patch("google.genai.Client")
|
|
def test_gemini_none_values_filtered(self, mock_client_class):
|
|
"""Test that None values are not passed to Gemini API."""
|
|
mock_client = mock.Mock()
|
|
mock_client_class.return_value = mock_client
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.text = '{"result": "test"}'
|
|
mock_client.models.generate_content.return_value = mock_response
|
|
|
|
model = gemini.GeminiLanguageModel(
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test-key",
|
|
)
|
|
|
|
list(model.infer(["test prompt"], candidate_count=None))
|
|
|
|
call_args = mock_client.models.generate_content.call_args
|
|
config = call_args.kwargs["config"]
|
|
|
|
self.assertNotIn("candidate_count", config)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main()
|