wehub-resource-sync
336 lines
12 KiB
Python
336 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from google.genai import types
|
|
|
|
from livekit.agents import llm
|
|
from livekit.agents.llm import ChatContext, function_tool
|
|
from livekit.agents.types import APIConnectOptions
|
|
from livekit.plugins.google.llm import LLM, LLMStream
|
|
from livekit.plugins.google.realtime.realtime_api import RealtimeModel, RealtimeSession
|
|
|
|
pytestmark = pytest.mark.plugin("google")
|
|
|
|
|
|
@pytest.fixture
|
|
def llm_stream():
|
|
mock_llm = MagicMock()
|
|
mock_llm._thought_signatures = {}
|
|
|
|
with patch.object(LLMStream, "__init__", lambda self, *a, **kw: None):
|
|
stream = LLMStream.__new__(LLMStream)
|
|
stream._llm = mock_llm
|
|
stream._model = "gemini-2.0-flash"
|
|
return stream
|
|
|
|
|
|
class TestParsePartFunctionCall:
|
|
def test_function_call_with_text_returns_none_content(self, llm_stream: LLMStream):
|
|
part = types.Part(
|
|
function_call=types.FunctionCall(name="get_weather", args={"city": "Paris"}),
|
|
text="get_weather",
|
|
)
|
|
|
|
chunk = llm_stream._parse_part("test-id", part)
|
|
|
|
assert chunk is not None
|
|
assert chunk.delta.content is None
|
|
assert chunk.delta.tool_calls is not None
|
|
assert len(chunk.delta.tool_calls) == 1
|
|
tool_call = chunk.delta.tool_calls[0]
|
|
assert tool_call.name == "get_weather"
|
|
assert '"city": "Paris"' in tool_call.arguments
|
|
|
|
def test_function_call_without_text_returns_none_content(self, llm_stream: LLMStream):
|
|
part = types.Part(
|
|
function_call=types.FunctionCall(name="get_weather", args={"city": "Paris"}),
|
|
)
|
|
|
|
chunk = llm_stream._parse_part("test-id", part)
|
|
|
|
assert chunk is not None
|
|
assert chunk.delta.content is None
|
|
assert len(chunk.delta.tool_calls) == 1
|
|
|
|
def test_text_only_part_returns_text_content(self, llm_stream: LLMStream):
|
|
part = types.Part(text="Hello world")
|
|
|
|
chunk = llm_stream._parse_part("test-id", part)
|
|
|
|
assert chunk is not None
|
|
assert chunk.delta.content == "Hello world"
|
|
assert not chunk.delta.tool_calls
|
|
|
|
def test_empty_text_part_returns_none(self, llm_stream: LLMStream):
|
|
part = types.Part(text="")
|
|
|
|
chunk = llm_stream._parse_part("test-id", part)
|
|
|
|
assert chunk is None
|
|
|
|
|
|
class TestCachedContentOption:
|
|
"""Verify the ``cached_content`` constructor option propagates from
|
|
``LLM.__init__`` through ``_LLMOptions`` and into the keyword
|
|
arguments handed to ``GenerateContentConfig`` for every request.
|
|
|
|
The propagation tests are ``async def`` because ``LLM.chat()`` builds
|
|
an ``LLMStream`` whose constructor schedules a metrics-monitoring
|
|
task on the running event loop. A sync test would raise
|
|
``RuntimeError: no running event loop`` before reaching the
|
|
assertion.
|
|
"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cached_content_propagates_to_extra_kwargs(self) -> None:
|
|
llm = LLM(model="gemini-2.5-flash", api_key="test", cached_content="cachedContents/abc123")
|
|
|
|
stream = llm.chat(chat_ctx=ChatContext.empty())
|
|
try:
|
|
assert stream._extra_kwargs.get("cached_content") == "cachedContents/abc123"
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cached_content_omitted_when_not_set(self) -> None:
|
|
"""Backward compatibility: callers that don't pass
|
|
``cached_content`` must produce a request config without the
|
|
field, so existing behaviour is unchanged."""
|
|
llm = LLM(model="gemini-2.5-flash", api_key="test")
|
|
|
|
stream = llm.chat(chat_ctx=ChatContext.empty())
|
|
try:
|
|
assert "cached_content" not in stream._extra_kwargs
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
def test_cached_content_stored_on_opts(self) -> None:
|
|
llm = LLM(
|
|
model="gemini-2.5-flash",
|
|
api_key="test",
|
|
cached_content="projects/p/locations/us-central1/cachedContents/xyz",
|
|
)
|
|
|
|
assert llm._opts.cached_content == "projects/p/locations/us-central1/cachedContents/xyz"
|
|
|
|
|
|
class TestCachedContentRequestSuppression:
|
|
"""Gemini's API rejects ``generateContent`` requests that pass
|
|
``cached_content`` together with ``system_instruction``, ``tools``,
|
|
or ``tool_config`` — those fields belong inside the CachedContent
|
|
resource. The plugin therefore strips them off the outgoing request
|
|
whenever a cache is attached. These tests run the LLMStream against
|
|
a stubbed ``generate_content_stream`` and assert on the
|
|
``GenerateContentConfig`` it received.
|
|
"""
|
|
|
|
@staticmethod
|
|
async def _single_response_async_iter():
|
|
"""Emit one minimal-but-valid GenerateContentResponse so the
|
|
retry layer in livekit.agents.LLM doesn't treat the stream as
|
|
empty and re-issue the request three more times."""
|
|
yield types.GenerateContentResponse(
|
|
candidates=[
|
|
types.Candidate(
|
|
content=types.Content(
|
|
role="model",
|
|
parts=[types.Part(text="ok")],
|
|
),
|
|
finish_reason=types.FinishReason.STOP,
|
|
)
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def _patched_stream_capture(cls) -> tuple[AsyncMock, dict]:
|
|
captured: dict = {}
|
|
|
|
async def fake_stream(**kwargs):
|
|
captured["model"] = kwargs.get("model")
|
|
captured["contents"] = kwargs.get("contents")
|
|
captured["config"] = kwargs.get("config")
|
|
return cls._single_response_async_iter()
|
|
|
|
return AsyncMock(side_effect=fake_stream), captured
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_omits_system_instruction_when_cached_content_set(self) -> None:
|
|
"""With a cache attached, the outgoing request must carry
|
|
``system_instruction=None`` — the system prompt lives in the
|
|
cache resource and re-sending it would make Gemini return 400."""
|
|
llm = LLM(
|
|
model="gemini-2.5-flash",
|
|
api_key="test",
|
|
cached_content="cachedContents/abc123",
|
|
)
|
|
|
|
chat_ctx = ChatContext.empty()
|
|
chat_ctx.add_message(role="system", content="system prompt that lives in cache")
|
|
chat_ctx.add_message(role="user", content="hi")
|
|
|
|
fake, captured = self._patched_stream_capture()
|
|
with patch.object(llm._client.aio.models, "generate_content_stream", fake):
|
|
stream = llm.chat(chat_ctx=chat_ctx)
|
|
try:
|
|
async for _ in stream:
|
|
pass
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
config = captured["config"]
|
|
assert config.system_instruction is None
|
|
assert config.cached_content == "cachedContents/abc123"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_omits_tools_when_cached_content_set(self) -> None:
|
|
"""With a cache attached, the outgoing request must NOT include
|
|
``tools`` even if the LLMStream was constructed with function
|
|
tools — the tool schemas belong inside the cache resource."""
|
|
|
|
@function_tool
|
|
async def example_tool(query: str) -> str:
|
|
"""Look something up."""
|
|
return query
|
|
|
|
llm = LLM(
|
|
model="gemini-2.5-flash",
|
|
api_key="test",
|
|
cached_content="cachedContents/abc123",
|
|
)
|
|
|
|
fake, captured = self._patched_stream_capture()
|
|
with patch.object(llm._client.aio.models, "generate_content_stream", fake):
|
|
stream = llm.chat(chat_ctx=ChatContext.empty(), tools=[example_tool])
|
|
try:
|
|
async for _ in stream:
|
|
pass
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
config = captured["config"]
|
|
assert config.tools is None
|
|
assert config.tool_config is None
|
|
assert config.cached_content == "cachedContents/abc123"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_includes_system_instruction_and_tools_when_no_cache(self) -> None:
|
|
"""Backward compatibility: without ``cached_content``, the
|
|
request still carries ``system_instruction`` and ``tools`` as
|
|
before. Suppression is gated strictly on the cache being set."""
|
|
|
|
@function_tool
|
|
async def example_tool(query: str) -> str:
|
|
"""Look something up."""
|
|
return query
|
|
|
|
llm = LLM(model="gemini-2.5-flash", api_key="test")
|
|
|
|
chat_ctx = ChatContext.empty()
|
|
chat_ctx.add_message(role="system", content="system prompt sent on every request")
|
|
chat_ctx.add_message(role="user", content="hi")
|
|
|
|
fake, captured = self._patched_stream_capture()
|
|
with patch.object(llm._client.aio.models, "generate_content_stream", fake):
|
|
stream = llm.chat(chat_ctx=chat_ctx, tools=[example_tool])
|
|
try:
|
|
async for _ in stream:
|
|
pass
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
config = captured["config"]
|
|
assert config.system_instruction is not None
|
|
assert config.tools is not None and len(config.tools) >= 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_merges_timeout_into_caller_http_options(self) -> None:
|
|
caller_http_options = types.HttpOptions(headers={"X-Vertex-Test": "1"})
|
|
llm = LLM(
|
|
model="gemini-2.5-flash",
|
|
api_key="test",
|
|
http_options=caller_http_options,
|
|
)
|
|
|
|
fake, captured = self._patched_stream_capture()
|
|
with patch.object(llm._client.aio.models, "generate_content_stream", fake):
|
|
stream = llm.chat(
|
|
chat_ctx=ChatContext.empty(),
|
|
conn_options=APIConnectOptions(timeout=7.5),
|
|
)
|
|
try:
|
|
async for _ in stream:
|
|
pass
|
|
finally:
|
|
await stream.aclose()
|
|
|
|
config = captured["config"]
|
|
assert config.http_options.timeout == 7500
|
|
assert config.http_options.headers["X-Vertex-Test"] == "1"
|
|
assert "livekit-agents/" in config.http_options.headers["x-goog-api-client"]
|
|
assert caller_http_options.timeout is None
|
|
assert caller_http_options.headers == {"X-Vertex-Test": "1"}
|
|
|
|
|
|
class TestMediaResolution:
|
|
def test_llm_media_resolution_is_passed_to_stream_kwargs(self):
|
|
model = LLM(
|
|
api_key="test-api-key",
|
|
media_resolution=types.MediaResolution.MEDIA_RESOLUTION_LOW,
|
|
)
|
|
|
|
with patch.object(
|
|
LLMStream,
|
|
"__init__",
|
|
lambda self, *a, **kw: self.__dict__.update(_extra_kwargs=kw["extra_kwargs"]),
|
|
):
|
|
stream = model.chat(chat_ctx=llm.ChatContext.empty())
|
|
|
|
assert (
|
|
stream._extra_kwargs["media_resolution"] == types.MediaResolution.MEDIA_RESOLUTION_LOW
|
|
)
|
|
|
|
def test_llm_media_resolution_is_omitted_by_default(self):
|
|
model = LLM(api_key="test-api-key")
|
|
|
|
with patch.object(
|
|
LLMStream,
|
|
"__init__",
|
|
lambda self, *a, **kw: self.__dict__.update(_extra_kwargs=kw["extra_kwargs"]),
|
|
):
|
|
stream = model.chat(chat_ctx=llm.ChatContext.empty())
|
|
|
|
assert "media_resolution" not in stream._extra_kwargs
|
|
|
|
def test_realtime_media_resolution_is_passed_to_connect_config(self):
|
|
model = RealtimeModel(
|
|
api_key="test-api-key",
|
|
media_resolution=types.MediaResolution.MEDIA_RESOLUTION_LOW,
|
|
)
|
|
session = RealtimeSession.__new__(RealtimeSession)
|
|
session._opts = model._opts
|
|
session._tools = llm.ToolContext.empty()
|
|
session._realtime_model = model
|
|
session._session_resumption_handle = None
|
|
|
|
config = session._build_connect_config()
|
|
|
|
assert config.generation_config
|
|
assert (
|
|
config.generation_config.media_resolution == types.MediaResolution.MEDIA_RESOLUTION_LOW
|
|
)
|
|
|
|
def test_realtime_media_resolution_is_unset_by_default(self):
|
|
model = RealtimeModel(api_key="test-api-key")
|
|
session = RealtimeSession.__new__(RealtimeSession)
|
|
session._opts = model._opts
|
|
session._tools = llm.ToolContext.empty()
|
|
session._realtime_model = model
|
|
session._session_resumption_handle = None
|
|
|
|
config = session._build_connect_config()
|
|
|
|
assert config.generation_config
|
|
assert config.generation_config.media_resolution is None
|