wehub-resource-sync
97e91a83f3
Ruff / Ruff (push) Waiting to run
Test / Core Tests (push) Waiting to run
Test / Offline Coverage Tests (Python 3.10) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.11) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.12) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.13) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.9) (push) Waiting to run
Test / Full Coverage (Python 3.11) (push) Waiting to run
Test / Core Provider Tests (OpenAI) (push) Blocked by required conditions
Test / Core Provider Tests (Anthropic) (push) Blocked by required conditions
Test / Core Provider Tests (Google) (push) Blocked by required conditions
Test / Core Provider Tests (Other) (push) Blocked by required conditions
Test / Anthropic Tests (push) Blocked by required conditions
Test / Gemini Tests (push) Blocked by required conditions
Test / Google GenAI Tests (push) Blocked by required conditions
Test / Vertex AI Tests (push) Blocked by required conditions
Test / OpenAI Tests (push) Blocked by required conditions
Test / Writer Tests (push) Blocked by required conditions
Test / Auto Client Tests (push) Blocked by required conditions
ty / type-check (push) Waiting to run
720 lines
23 KiB
Python
720 lines
23 KiB
Python
"""Offline behavior coverage for the Google GenAI v2 provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import builtins
|
|
import importlib
|
|
import json
|
|
from collections.abc import AsyncGenerator, AsyncIterator, Iterator
|
|
from types import SimpleNamespace
|
|
from typing import Any, ClassVar, cast
|
|
|
|
import pytest
|
|
from google.genai import Client, types
|
|
from pydantic import BaseModel
|
|
|
|
from instructor.v2.core.errors import ClientError, ModeError
|
|
from instructor.v2.core.mode import Mode
|
|
from instructor.v2.core.providers import Provider
|
|
from instructor.v2.dsl.iterable import IterableBase
|
|
from instructor.v2.dsl.parallel import ParallelBase
|
|
from instructor.v2.dsl.partial import Partial, PartialBase
|
|
from instructor.v2.dsl.simple_type import AdapterBase
|
|
from instructor.v2.providers.genai import handlers, multimodal
|
|
from instructor.v2.providers.genai import client as genai_client
|
|
from tests.coverage._streams import async_items
|
|
|
|
|
|
class Answer(BaseModel):
|
|
answer: int
|
|
|
|
|
|
class RecordingModels:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict[str, Any]]] = []
|
|
|
|
def generate_content(self, **kwargs: Any) -> dict[str, Any]:
|
|
self.calls.append(("content", kwargs))
|
|
return {"kind": "content", **kwargs}
|
|
|
|
def generate_content_stream(self, **kwargs: Any) -> Iterator[dict[str, Any]]:
|
|
self.calls.append(("stream", kwargs))
|
|
yield {"kind": "stream", **kwargs}
|
|
|
|
|
|
class RecordingAsyncModels:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict[str, Any]]] = []
|
|
|
|
async def generate_content(self, **kwargs: Any) -> dict[str, Any]:
|
|
self.calls.append(("content", kwargs))
|
|
return {"kind": "content", **kwargs}
|
|
|
|
async def generate_content_stream(
|
|
self, **kwargs: Any
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
self.calls.append(("stream", kwargs))
|
|
|
|
async def stream() -> AsyncIterator[dict[str, Any]]:
|
|
yield {"kind": "stream", **kwargs}
|
|
|
|
return stream()
|
|
|
|
|
|
class RecordingClient:
|
|
def __init__(self) -> None:
|
|
self.models = RecordingModels()
|
|
self.aio = SimpleNamespace(models=RecordingAsyncModels())
|
|
|
|
|
|
def _identity_patch(**kwargs: Any) -> Any:
|
|
return kwargs["func"]
|
|
|
|
|
|
def test_from_genai_rejects_missing_sdk_and_wrong_client(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(genai_client, "Client", None)
|
|
with pytest.raises(ClientError, match="google-genai is not installed"):
|
|
genai_client.from_genai(cast(Client, object()))
|
|
|
|
monkeypatch.setattr(genai_client, "Client", RecordingClient)
|
|
with pytest.raises(ClientError, match="Got: object"):
|
|
genai_client.from_genai(cast(Client, object()))
|
|
|
|
|
|
def test_from_genai_rejects_an_unregistered_mode(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(genai_client, "Client", RecordingClient)
|
|
|
|
with pytest.raises(ModeError) as error:
|
|
genai_client.from_genai(cast(Client, RecordingClient()), mode=Mode.JSON_SCHEMA)
|
|
|
|
assert error.value.provider == Provider.GENAI.value
|
|
assert error.value.mode == Mode.JSON_SCHEMA.value
|
|
assert Mode.TOOLS.value in error.value.valid_modes
|
|
assert Mode.JSON.value in error.value.valid_modes
|
|
|
|
|
|
def test_from_genai_sync_wrapper_routes_content_and_stream(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(genai_client, "Client", RecordingClient)
|
|
monkeypatch.setattr(genai_client, "patch_v2", _identity_patch)
|
|
client = RecordingClient()
|
|
|
|
wrapped = genai_client.from_genai(
|
|
cast(Client, client), mode=Mode.TOOLS, model="gemini-default"
|
|
)
|
|
content = wrapped.create_fn(contents=["one"])
|
|
stream = list(
|
|
wrapped.create_fn(model="gemini-override", stream=True, contents=["two"])
|
|
)
|
|
|
|
assert content == {
|
|
"kind": "content",
|
|
"model": "gemini-default",
|
|
"contents": ["one"],
|
|
}
|
|
assert stream == [
|
|
{
|
|
"kind": "stream",
|
|
"model": "gemini-override",
|
|
"contents": ["two"],
|
|
}
|
|
]
|
|
assert client.models.calls == [
|
|
("content", {"model": "gemini-default", "contents": ["one"]}),
|
|
("stream", {"model": "gemini-override", "contents": ["two"]}),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_from_genai_async_wrapper_routes_content_and_stream(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(genai_client, "Client", RecordingClient)
|
|
monkeypatch.setattr(genai_client, "patch_v2", _identity_patch)
|
|
client = RecordingClient()
|
|
|
|
wrapped = genai_client.from_genai(
|
|
cast(Client, client), mode=Mode.JSON, use_async=True, model="gemini-default"
|
|
)
|
|
content = await wrapped.create_fn(contents=["one"])
|
|
stream = await wrapped.create_fn(
|
|
model="gemini-override", stream=True, contents=["two"]
|
|
)
|
|
chunks = [chunk async for chunk in stream]
|
|
|
|
assert content == {
|
|
"kind": "content",
|
|
"model": "gemini-default",
|
|
"contents": ["one"],
|
|
}
|
|
assert chunks == [
|
|
{
|
|
"kind": "stream",
|
|
"model": "gemini-override",
|
|
"contents": ["two"],
|
|
}
|
|
]
|
|
assert client.aio.models.calls == [
|
|
("content", {"model": "gemini-default", "contents": ["one"]}),
|
|
("stream", {"model": "gemini-override", "contents": ["two"]}),
|
|
]
|
|
|
|
|
|
def test_genai_client_module_handles_missing_optional_dependency(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
original_import = builtins.__import__
|
|
|
|
def blocked_import(
|
|
name: str,
|
|
globals: dict[str, Any] | None = None,
|
|
locals: dict[str, Any] | None = None,
|
|
fromlist: tuple[str, ...] = (),
|
|
level: int = 0,
|
|
) -> Any:
|
|
if name == "google.genai" and "Client" in fromlist:
|
|
raise ImportError("google-genai is unavailable")
|
|
return original_import(name, globals, locals, fromlist, level)
|
|
|
|
try:
|
|
with monkeypatch.context() as context:
|
|
context.setattr(builtins, "__import__", blocked_import)
|
|
reloaded = importlib.reload(genai_client)
|
|
assert reloaded.Client is None
|
|
with pytest.raises(ClientError, match="google-genai is not installed"):
|
|
reloaded.from_genai(object())
|
|
finally:
|
|
restored = importlib.reload(genai_client)
|
|
|
|
assert restored.Client is not None
|
|
|
|
|
|
def test_genai_package_handles_a_failed_client_import(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
package = importlib.import_module("instructor.v2.providers.genai")
|
|
original_import = builtins.__import__
|
|
|
|
def blocked_import(
|
|
name: str,
|
|
globals: dict[str, Any] | None = None,
|
|
locals: dict[str, Any] | None = None,
|
|
fromlist: tuple[str, ...] = (),
|
|
level: int = 0,
|
|
) -> Any:
|
|
if (
|
|
level == 1
|
|
and name == "client"
|
|
and globals is not None
|
|
and globals.get("__package__") == "instructor.v2.providers.genai"
|
|
):
|
|
raise ImportError("client dependency is unavailable")
|
|
return original_import(name, globals, locals, fromlist, level)
|
|
|
|
try:
|
|
with monkeypatch.context() as context:
|
|
context.setattr(builtins, "__import__", blocked_import)
|
|
assert importlib.reload(package).from_genai is None
|
|
finally:
|
|
restored = importlib.reload(package)
|
|
|
|
assert callable(restored.from_genai)
|
|
|
|
|
|
def test_reask_tools_preserves_model_turn_and_adds_function_error() -> None:
|
|
earlier_model_content = types.Content(
|
|
role="model", parts=[types.Part.from_text(text="Let me check that.")]
|
|
)
|
|
model_content = types.Content(
|
|
role="model",
|
|
parts=[types.Part.from_function_call(name="Answer", args={"answer": "bad"})],
|
|
)
|
|
response = SimpleNamespace(
|
|
candidates=[
|
|
SimpleNamespace(content=None),
|
|
SimpleNamespace(content=earlier_model_content),
|
|
SimpleNamespace(content=model_content),
|
|
]
|
|
)
|
|
original_contents = (
|
|
types.Content(role="user", parts=[types.Part.from_text(text="How many?")]),
|
|
)
|
|
|
|
result = handlers.reask_genai_tools(
|
|
{"contents": original_contents}, response, ValueError("answer must be an int")
|
|
)
|
|
|
|
assert result["contents"][0] is original_contents[0]
|
|
assert result["contents"][1] is model_content
|
|
tool_turn = result["contents"][2]
|
|
assert tool_turn.role == "tool"
|
|
assert tool_turn.parts[0].function_response.name == "Answer"
|
|
assert (
|
|
"answer must be an int"
|
|
in tool_turn.parts[0].function_response.response["error"]
|
|
)
|
|
assert len(original_contents) == 1
|
|
|
|
|
|
def test_reask_tools_without_function_call_adds_a_user_correction() -> None:
|
|
model_content = types.Content(
|
|
role="model", parts=[types.Part.from_text(text="The answer is many.")]
|
|
)
|
|
response = SimpleNamespace(candidates=[SimpleNamespace(content=model_content)])
|
|
|
|
result = handlers.reask_genai_tools(
|
|
{"contents": None}, response, ValueError("answer must be an int")
|
|
)
|
|
|
|
assert result["contents"][0] is model_content
|
|
correction = result["contents"][1]
|
|
assert correction.role == "user"
|
|
assert "answer must be an int" in correction.parts[0].text
|
|
assert "Recall the function correctly" in correction.parts[0].text
|
|
|
|
|
|
def test_parse_genai_tools_ignores_thought_parts_and_validates_function_call() -> None:
|
|
completion = types.GenerateContentResponse(
|
|
candidates=[
|
|
types.Candidate(
|
|
content=types.Content(
|
|
role="model",
|
|
parts=[
|
|
types.Part(text="Let me think", thought=True),
|
|
types.Part.from_function_call(
|
|
name="Answer", args={"answer": 7}
|
|
),
|
|
],
|
|
)
|
|
)
|
|
]
|
|
)
|
|
|
|
parsed = handlers.parse_genai_tools(Answer, completion, strict=True)
|
|
handled = handlers.GenAIToolsHandler(mode=Mode.TOOLS).parse_response(
|
|
completion, Answer, strict=True
|
|
)
|
|
|
|
assert parsed == Answer(answer=7)
|
|
assert handled == Answer(answer=7)
|
|
assert cast(Any, handled)._raw_response is completion
|
|
|
|
|
|
class MissingTextChunk:
|
|
def __init__(
|
|
self,
|
|
fallback_text: str | None,
|
|
text_error: type[Exception] = ValueError,
|
|
) -> None:
|
|
self.candidates = [
|
|
SimpleNamespace(
|
|
content=SimpleNamespace(parts=[SimpleNamespace(text=fallback_text)])
|
|
)
|
|
]
|
|
self.text_error = text_error
|
|
|
|
@property
|
|
def text(self) -> str:
|
|
raise self.text_error("text accessor failed")
|
|
|
|
|
|
def test_streaming_extractors_handle_tools_text_fallback_and_bad_chunks() -> None:
|
|
tool_chunk = SimpleNamespace(
|
|
candidates=[
|
|
SimpleNamespace(
|
|
content=SimpleNamespace(
|
|
parts=[
|
|
SimpleNamespace(
|
|
function_call=SimpleNamespace(args={"answer": 8})
|
|
)
|
|
]
|
|
)
|
|
)
|
|
]
|
|
)
|
|
tools_handler = handlers.GenAIToolsHandler(mode=Mode.TOOLS)
|
|
json_handler = handlers.GenAIStructuredOutputsHandler(mode=Mode.JSON)
|
|
|
|
assert list(tools_handler.extract_streaming_json([object(), tool_chunk])) == [
|
|
json.dumps({"answer": 8})
|
|
]
|
|
assert list(
|
|
json_handler.extract_streaming_json(
|
|
[object(), SimpleNamespace(text='{"answer": 9}'), MissingTextChunk("tail")]
|
|
)
|
|
) == ['{"answer": 9}', "tail"]
|
|
with pytest.raises(ValueError, match="text accessor failed"):
|
|
list(json_handler.extract_streaming_json([MissingTextChunk(None)]))
|
|
with pytest.raises(RuntimeError, match="text accessor failed"):
|
|
list(
|
|
json_handler.extract_streaming_json(
|
|
[MissingTextChunk("must not mask an internal error", RuntimeError)]
|
|
)
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_streaming_extractors_handle_tools_text_fallback_and_bad_chunks() -> (
|
|
None
|
|
):
|
|
tool_chunk = SimpleNamespace(
|
|
candidates=[
|
|
SimpleNamespace(
|
|
content=SimpleNamespace(
|
|
parts=[
|
|
SimpleNamespace(
|
|
function_call=SimpleNamespace(args={"answer": 10})
|
|
)
|
|
]
|
|
)
|
|
)
|
|
]
|
|
)
|
|
tools_handler = handlers.GenAIToolsHandler(mode=Mode.TOOLS)
|
|
json_handler = handlers.GenAIStructuredOutputsHandler(mode=Mode.JSON)
|
|
|
|
tool_chunks = [
|
|
item
|
|
async for item in tools_handler.extract_streaming_json_async(
|
|
async_items([object(), tool_chunk])
|
|
)
|
|
]
|
|
json_chunks = [
|
|
item
|
|
async for item in json_handler.extract_streaming_json_async(
|
|
async_items(
|
|
[
|
|
object(),
|
|
SimpleNamespace(text='{"answer": 11}'),
|
|
MissingTextChunk("tail"),
|
|
]
|
|
)
|
|
)
|
|
]
|
|
|
|
assert tool_chunks == [json.dumps({"answer": 10})]
|
|
assert json_chunks == ['{"answer": 11}', "tail"]
|
|
with pytest.raises(ValueError, match="text accessor failed"):
|
|
[
|
|
item
|
|
async for item in json_handler.extract_streaming_json_async(
|
|
async_items([MissingTextChunk(None)])
|
|
)
|
|
]
|
|
with pytest.raises(RuntimeError, match="text accessor failed"):
|
|
[
|
|
item
|
|
async for item in json_handler.extract_streaming_json_async(
|
|
async_items(
|
|
[MissingTextChunk("must not mask an internal error", RuntimeError)]
|
|
)
|
|
)
|
|
]
|
|
|
|
|
|
class StreamIterable(BaseModel, IterableBase):
|
|
tasks: list[Answer]
|
|
task_type: ClassVar[type[BaseModel] | None] = Answer
|
|
|
|
|
|
def test_handler_base_request_and_response_paths(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
base = handlers.GenAIHandlerBase(mode=Mode.JSON)
|
|
kwargs = {"keep": True}
|
|
|
|
assert base._extract_system_instruction({}) is None
|
|
assert base._wrap_streaming_model(None, stream=True) is None
|
|
wrapped_model = base._wrap_streaming_model(Answer, stream=True)
|
|
assert wrapped_model is not None
|
|
assert issubclass(wrapped_model, PartialBase)
|
|
assert base.handle_reask(kwargs, None, ValueError("bad")) == kwargs
|
|
assert base.handle_reask(kwargs, None, ValueError("bad")) is not kwargs
|
|
assert base.parse_response("raw", None) == "raw"
|
|
with pytest.raises(NotImplementedError):
|
|
base.prepare_request(Answer, {})
|
|
|
|
iterable_stream = [SimpleNamespace(text='[{"answer": 1}, {"answer": 2}]')]
|
|
partial_stream = [SimpleNamespace(text='{"answer":'), SimpleNamespace(text=" 3}")]
|
|
iterable = base.parse_response(iterable_stream, StreamIterable, stream=True)
|
|
partial = base.parse_response(partial_stream, Partial[Answer], stream=True)
|
|
assert list(iterable) == [Answer(answer=1), Answer(answer=2)]
|
|
assert isinstance(partial, list)
|
|
assert partial[-1] == Answer(answer=3)
|
|
|
|
iterable_result = StreamIterable(tasks=[Answer(answer=1), Answer(answer=2)])
|
|
monkeypatch.setattr(
|
|
handlers,
|
|
"parse_genai_structured_outputs",
|
|
lambda *_args, **_kwargs: iterable_result,
|
|
)
|
|
assert base.parse_response(object(), Answer) == [
|
|
Answer(answer=1),
|
|
Answer(answer=2),
|
|
]
|
|
|
|
parallel_result = [Answer(answer=3)]
|
|
monkeypatch.setattr(
|
|
handlers,
|
|
"parse_genai_structured_outputs",
|
|
lambda *_args, **_kwargs: parallel_result,
|
|
)
|
|
parallel_model = cast(type[BaseModel], ParallelBase(Answer))
|
|
assert base.parse_response(object(), parallel_model) is parallel_result
|
|
|
|
class AnswerAdapter(AdapterBase):
|
|
content: int
|
|
|
|
monkeypatch.setattr(
|
|
handlers,
|
|
"parse_genai_structured_outputs",
|
|
lambda *_args, **_kwargs: AnswerAdapter(content=4),
|
|
)
|
|
assert base.parse_response(object(), AnswerAdapter) == 4
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handler_base_async_streaming_response_path() -> None:
|
|
handler = handlers.GenAIStructuredOutputsHandler(mode=Mode.JSON)
|
|
stream = async_items([SimpleNamespace(text='[{"answer": 1}, {"answer": 2}]')])
|
|
|
|
result = cast(
|
|
AsyncGenerator[BaseModel, None],
|
|
handler.parse_response(stream, StreamIterable, stream=True, is_async=True),
|
|
)
|
|
|
|
assert [item async for item in result] == [Answer(answer=1), Answer(answer=2)]
|
|
|
|
|
|
def test_prepare_request_without_model_keeps_explicit_system_instruction() -> None:
|
|
handler = handlers.GenAIToolsHandler(mode=Mode.TOOLS)
|
|
|
|
model, request = handler.prepare_request(
|
|
None,
|
|
{
|
|
"system": "Keep the answer short.",
|
|
"messages": [{"role": "user", "content": "How many?"}],
|
|
"temperature": 0.2,
|
|
},
|
|
)
|
|
|
|
assert model is None
|
|
assert request["config"].system_instruction == "Keep the answer short."
|
|
assert request["contents"][0].parts[0].text == "How many?"
|
|
assert "system" not in request
|
|
assert "temperature" not in request
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("handler", "expected_mode"),
|
|
[
|
|
(handlers.GenAIToolsHandler(mode=Mode.TOOLS), "tools"),
|
|
(handlers.GenAIStructuredOutputsHandler(mode=Mode.JSON), "json"),
|
|
],
|
|
)
|
|
def test_prepare_request_merges_openai_generation_settings(
|
|
handler: handlers.GenAIHandlerBase, expected_mode: str
|
|
) -> None:
|
|
original = {
|
|
"messages": [
|
|
{"role": "system", "content": "Return a number."},
|
|
{"role": "user", "content": "How many?"},
|
|
],
|
|
"max_tokens": 32,
|
|
"temperature": 0.25,
|
|
"top_p": 0.8,
|
|
"n": 1,
|
|
"stop": ["END"],
|
|
"seed": 2,
|
|
"presence_penalty": 0.1,
|
|
"frequency_penalty": 0.2,
|
|
}
|
|
|
|
prepared_model, request = handler.prepare_request(Answer, original)
|
|
|
|
assert prepared_model is not None
|
|
assert prepared_model.__name__ == "Answer"
|
|
assert prepared_model.model_validate({"answer": 4}).model_dump() == {"answer": 4}
|
|
assert "messages" not in request
|
|
assert len(request["contents"]) == 1
|
|
assert request["contents"][0].role == "user"
|
|
assert request["contents"][0].parts[0].text == "How many?"
|
|
config = request["config"]
|
|
assert config.max_output_tokens == 32
|
|
assert config.temperature == 0.25
|
|
assert config.top_p == 0.8
|
|
assert config.candidate_count == 1
|
|
assert config.stop_sequences == ["END"]
|
|
assert config.seed == 2
|
|
assert config.presence_penalty == 0.1
|
|
assert config.frequency_penalty == 0.2
|
|
assert config.system_instruction.strip() == "Return a number."
|
|
if expected_mode == "tools":
|
|
assert config.tools[0].function_declarations[0].name == "Answer"
|
|
assert config.tool_config.function_calling_config.allowed_function_names == [
|
|
"Answer"
|
|
]
|
|
else:
|
|
assert config.response_mime_type == "application/json"
|
|
assert config.response_schema is prepared_model
|
|
assert original["max_tokens"] == 32
|
|
assert len(original["messages"]) == 2
|
|
|
|
|
|
def test_multimodal_types_reports_missing_optional_dependency(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
original_import = builtins.__import__
|
|
|
|
def blocked_import(
|
|
name: str,
|
|
globals: dict[str, Any] | None = None,
|
|
locals: dict[str, Any] | None = None,
|
|
fromlist: tuple[str, ...] = (),
|
|
level: int = 0,
|
|
) -> Any:
|
|
if name == "google.genai" and "types" in fromlist:
|
|
raise ImportError("google-genai is unavailable")
|
|
return original_import(name, globals, locals, fromlist, level)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", blocked_import)
|
|
with pytest.raises(ImportError, match="google-genai package is required"):
|
|
multimodal._types()
|
|
|
|
|
|
def test_multimodal_encodes_gcs_image_and_audio() -> None:
|
|
image = SimpleNamespace(
|
|
source="gs://bucket/image.png", media_type="image/png", data=b"png-bytes"
|
|
)
|
|
audio = SimpleNamespace(
|
|
source="recording.wav", media_type="audio/wav", data="d2F2LWJ5dGVz"
|
|
)
|
|
|
|
image_part = multimodal.image_to_genai(image)
|
|
audio_part = multimodal.audio_to_genai(audio)
|
|
|
|
assert image_part.inline_data.data == b"png-bytes"
|
|
assert image_part.inline_data.mime_type == "image/png"
|
|
assert audio_part.inline_data.data == b"wav-bytes"
|
|
assert audio_part.inline_data.mime_type == "audio/wav"
|
|
|
|
|
|
def _install_pdf_upload_client(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
active_after_poll: bool,
|
|
pending_name: str | None = "files/pending",
|
|
) -> tuple[list[tuple[str, str]], list[int]]:
|
|
pending = SimpleNamespace(
|
|
state=types.FileState.PROCESSING,
|
|
name=pending_name,
|
|
uri="gs://bucket/pending.pdf",
|
|
mime_type="application/pdf",
|
|
)
|
|
polled = SimpleNamespace(
|
|
state=(
|
|
types.FileState.ACTIVE if active_after_poll else types.FileState.PROCESSING
|
|
),
|
|
name="files/pending",
|
|
uri="gs://bucket/ready.pdf" if active_after_poll else pending.uri,
|
|
mime_type="application/pdf",
|
|
)
|
|
calls: list[tuple[str, str]] = []
|
|
|
|
class Files:
|
|
def upload(self, *, file: str) -> SimpleNamespace:
|
|
calls.append(("upload", file))
|
|
return pending
|
|
|
|
def get(self, *, name: str) -> SimpleNamespace:
|
|
calls.append(("get", name))
|
|
return polled
|
|
|
|
class Client:
|
|
def __init__(self) -> None:
|
|
self.files = Files()
|
|
|
|
sleeps: list[int] = []
|
|
monkeypatch.setattr("google.genai.Client", Client)
|
|
monkeypatch.setattr("time.sleep", sleeps.append)
|
|
return calls, sleeps
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("max_retries", "expected_calls", "expected_sleeps"),
|
|
[
|
|
(0, [("upload", "/tmp/pending.pdf")], []),
|
|
(
|
|
1,
|
|
[("upload", "/tmp/pending.pdf"), ("get", "files/pending")],
|
|
[3],
|
|
),
|
|
],
|
|
)
|
|
def test_upload_new_pdf_file_stops_at_the_retry_limit(
|
|
max_retries: int,
|
|
expected_calls: list[tuple[str, str]],
|
|
expected_sleeps: list[int],
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
calls, sleeps = _install_pdf_upload_client(monkeypatch, active_after_poll=False)
|
|
|
|
with pytest.raises(Exception, match="Max retries reached"):
|
|
multimodal.upload_new_pdf_file(
|
|
SimpleNamespace,
|
|
"/tmp/pending.pdf",
|
|
retry_delay=3,
|
|
max_retries=max_retries,
|
|
)
|
|
|
|
assert calls == expected_calls
|
|
assert sleeps == expected_sleeps
|
|
|
|
|
|
def test_upload_new_pdf_file_recovers_on_the_last_allowed_retry(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
calls, sleeps = _install_pdf_upload_client(monkeypatch, active_after_poll=True)
|
|
|
|
result = multimodal.upload_new_pdf_file(
|
|
SimpleNamespace, "/tmp/pending.pdf", retry_delay=3, max_retries=1
|
|
)
|
|
|
|
assert result == SimpleNamespace(
|
|
source="gs://bucket/ready.pdf", media_type="application/pdf", data=None
|
|
)
|
|
assert calls == [("upload", "/tmp/pending.pdf"), ("get", "files/pending")]
|
|
assert sleeps == [3]
|
|
|
|
|
|
def test_upload_new_pdf_file_rejects_a_pending_file_without_a_name(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
calls, sleeps = _install_pdf_upload_client(
|
|
monkeypatch, active_after_poll=False, pending_name=None
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Cannot poll a pending GenAI file"):
|
|
multimodal.upload_new_pdf_file(
|
|
SimpleNamespace, "/tmp/pending.pdf", retry_delay=3, max_retries=1
|
|
)
|
|
|
|
assert calls == [("upload", "/tmp/pending.pdf")]
|
|
assert sleeps == []
|
|
|
|
|
|
def test_extract_multimodal_content_preserves_uploaded_files() -> None:
|
|
uploaded = types.File(
|
|
name="files/ready", uri="gs://bucket/ready.pdf", mime_type="application/pdf"
|
|
)
|
|
|
|
result = multimodal.extract_multimodal_content([uploaded])
|
|
|
|
assert result == [uploaded]
|
|
assert result[0] is uploaded
|