420 lines
13 KiB
Python
420 lines
13 KiB
Python
"""Tests for the llm_client module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
from typing import Any, cast
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from vibe_bot.config import (
|
|
CHAT_MODEL,
|
|
EMBEDDING_ENDPOINT,
|
|
EMBEDDING_ENDPOINT_KEY,
|
|
)
|
|
from vibe_bot.llm_client import (
|
|
chat_complete,
|
|
chat_completion_instruct,
|
|
embedding,
|
|
image_edit,
|
|
image_generation,
|
|
)
|
|
|
|
|
|
@pytest.mark.live
|
|
def test_chat_complete_live() -> None:
|
|
"""Live call to the chat endpoint via the core async ``chat_complete``.
|
|
|
|
Unmocked: requires network access to the configured chat API. Ported from
|
|
the former ``test_chat_completion_think`` (its sync ``chat_completion``
|
|
wrapper was deleted).
|
|
"""
|
|
import asyncio
|
|
|
|
result = asyncio.run(
|
|
chat_complete(
|
|
[
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Tell me about Everquest"},
|
|
],
|
|
model=CHAT_MODEL,
|
|
max_tokens=100,
|
|
)
|
|
)
|
|
assert isinstance(result, str)
|
|
|
|
|
|
@pytest.mark.live
|
|
def test_chat_completion_instruct_live() -> None:
|
|
"""Live call to the chat endpoint via the async instruct adapter.
|
|
|
|
Unmocked: requires network access to the configured chat API.
|
|
"""
|
|
import asyncio
|
|
|
|
result = asyncio.run(
|
|
chat_completion_instruct(
|
|
system_prompt="You are a helpful assistant.",
|
|
user_prompt="Tell me about Everquest",
|
|
model=CHAT_MODEL,
|
|
max_tokens=100,
|
|
)
|
|
)
|
|
assert isinstance(result, str)
|
|
|
|
|
|
def test_image_generation() -> None:
|
|
"""Image generation returns the first b64 payload from the API."""
|
|
import asyncio
|
|
|
|
mock_client = MagicMock()
|
|
mock_data = MagicMock()
|
|
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
|
|
mock_response = MagicMock()
|
|
mock_response.data = [mock_data]
|
|
mock_client.images.generate = AsyncMock(return_value=mock_response)
|
|
|
|
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
|
|
result = asyncio.run(
|
|
image_generation(
|
|
prompt="Generate an image of a horse",
|
|
model="test-image-model",
|
|
)
|
|
)
|
|
assert result == base64.b64encode(b"fake image data").decode()
|
|
|
|
|
|
def test_image_generation_api_error_returns_empty() -> None:
|
|
"""A 4xx/5xx (APIStatusError) from the image API returns "" without raising."""
|
|
import asyncio
|
|
|
|
import openai
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.images.generate = AsyncMock(
|
|
side_effect=openai.APIStatusError(
|
|
"boom",
|
|
response=MagicMock(),
|
|
body=None,
|
|
)
|
|
)
|
|
|
|
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
|
|
result = asyncio.run(
|
|
image_generation(
|
|
prompt="Generate an image of a horse",
|
|
model="test-image-model",
|
|
)
|
|
)
|
|
assert result == ""
|
|
|
|
|
|
def test_image_edit() -> None:
|
|
"""Image edit returns the first b64 payload from the API."""
|
|
import asyncio
|
|
|
|
mock_client = MagicMock()
|
|
mock_data = MagicMock()
|
|
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
|
|
mock_response = MagicMock()
|
|
mock_response.data = [mock_data]
|
|
mock_client.images.edit = AsyncMock(return_value=mock_response)
|
|
|
|
with patch("vibe_bot.llm.images.get_image_edit_client", return_value=mock_client):
|
|
result = asyncio.run(
|
|
image_edit(
|
|
image=BytesIO(b"fake image"),
|
|
prompt="Paint the words 'horse' on the horse.",
|
|
model="test-image-edit-model",
|
|
)
|
|
)
|
|
assert result == base64.b64encode(b"fake edited image data").decode()
|
|
|
|
|
|
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
|
|
"""Calculate cosine similarity between two arrays.
|
|
|
|
Returns a value close to 1 for similar vectors,
|
|
close to 0 for orthogonal vectors,
|
|
and close to -1 for opposite vectors.
|
|
"""
|
|
a_arr, b_arr = np.array(a), np.array(b)
|
|
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
|
|
|
|
|
|
EMBEDDING_SIMILARITY_HIGH = 0.9
|
|
EMBEDDING_SIMILARITY_LOW = 0.5
|
|
|
|
|
|
def _mock_embedding_session(
|
|
post: MagicMock,
|
|
) -> MagicMock:
|
|
"""Build a mock requests.Session whose .post is ``post``."""
|
|
session = MagicMock()
|
|
session.post = post
|
|
return session
|
|
|
|
|
|
def test_embeddings() -> None:
|
|
"""Embedding similarity for similar and different texts."""
|
|
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
|
|
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
|
|
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
|
|
|
|
def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
|
|
json_data = kwargs.get("json", {})
|
|
text = json_data["input"][0]
|
|
if "horse" in text and "donkey" not in text and "also" not in text:
|
|
embedding_data = mock_horse_vec
|
|
elif "also" in text:
|
|
embedding_data = mock_horse_also_vec
|
|
else:
|
|
embedding_data = mock_donkey_vec
|
|
mock_resp = MagicMock()
|
|
mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
|
|
return mock_resp
|
|
|
|
session = _mock_embedding_session(MagicMock(side_effect=mock_post))
|
|
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
|
|
result1 = embedding(
|
|
"this is a horse",
|
|
url=EMBEDDING_ENDPOINT,
|
|
api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
result2 = embedding(
|
|
"this is a horse also",
|
|
url=EMBEDDING_ENDPOINT,
|
|
api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
result3 = embedding(
|
|
"this is a donkey",
|
|
url=EMBEDDING_ENDPOINT,
|
|
api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
|
|
assert similarity_1 > EMBEDDING_SIMILARITY_HIGH
|
|
|
|
similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
|
|
assert similarity_2 < EMBEDDING_SIMILARITY_LOW
|
|
|
|
|
|
def test_embedding_non_json_2xx_returns_empty() -> None:
|
|
"""A 2xx response with a non-JSON body must return [] without raising.
|
|
|
|
Regression test for ``resp.json()`` sitting outside the try block, so an
|
|
HTML error page (or any non-JSON 2xx body) raised JSONDecodeError out of
|
|
``embedding`` and, through it, out of ``get_conversation_context``.
|
|
"""
|
|
mock_resp = MagicMock()
|
|
mock_resp.raise_for_status.return_value = None
|
|
mock_resp.json.side_effect = ValueError("<html>rate limited</html>")
|
|
|
|
session = _mock_embedding_session(MagicMock(return_value=mock_resp))
|
|
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
|
|
result = embedding(
|
|
"this is a horse",
|
|
url=EMBEDDING_ENDPOINT,
|
|
api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
|
|
assert result == []
|
|
|
|
|
|
def test_chat_client_singleton_identity() -> None:
|
|
"""The shared chat client is built once and reused across calls."""
|
|
from vibe_bot import llm_client
|
|
|
|
client1 = llm_client.get_chat_client()
|
|
client2 = llm_client.get_chat_client()
|
|
assert client1 is client2
|
|
|
|
|
|
def test_image_gen_client_singleton_identity() -> None:
|
|
"""The shared image-generation client is built once, from a cold start."""
|
|
import vibe_bot.llm.images as images_mod
|
|
|
|
saved = images_mod._image_gen_client
|
|
images_mod._image_gen_client = None
|
|
try:
|
|
client1 = images_mod.get_image_gen_client()
|
|
client2 = images_mod.get_image_gen_client()
|
|
assert client1 is client2
|
|
finally:
|
|
images_mod._image_gen_client = saved
|
|
|
|
|
|
def test_image_edit_client_singleton_identity() -> None:
|
|
"""The shared image-edit client is built once, from a cold start."""
|
|
import vibe_bot.llm.images as images_mod
|
|
|
|
saved = images_mod._image_edit_client
|
|
images_mod._image_edit_client = None
|
|
try:
|
|
client1 = images_mod.get_image_edit_client()
|
|
client2 = images_mod.get_image_edit_client()
|
|
assert client1 is client2
|
|
finally:
|
|
images_mod._image_edit_client = saved
|
|
|
|
|
|
def test_flows_build_no_new_clients_or_sessions(
|
|
mock_ctx: MagicMock,
|
|
temp_db_path: str,
|
|
) -> None:
|
|
"""A full !doodlebob + chat turn constructs no new clients or sessions.
|
|
|
|
Every shared client and the embedding session are built once at
|
|
"startup"; running the whole image flow and a chat turn through the real
|
|
singletons (with only HTTP mocked) must not construct another
|
|
AsyncOpenAI client or requests.Session.
|
|
"""
|
|
import asyncio
|
|
|
|
import openai
|
|
import requests
|
|
|
|
from vibe_bot import llm_client
|
|
from vibe_bot.database import ChatDatabase
|
|
from vibe_bot.services.chat_service import ChatService
|
|
from vibe_bot.services.image_service import ImageService
|
|
|
|
# Startup: build every shared client and the embedding session once.
|
|
chat_client = llm_client.get_chat_client()
|
|
gen_client = llm_client.get_image_gen_client()
|
|
edit_client = llm_client.get_image_edit_client()
|
|
llm_client.get_embedding_session()
|
|
|
|
counts = {"async_openai": 0, "session": 0}
|
|
real_session = requests.Session
|
|
|
|
class CountingAsyncOpenAI(openai.AsyncOpenAI):
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
counts["async_openai"] += 1
|
|
super().__init__(**kwargs)
|
|
|
|
def counting_session() -> requests.Session:
|
|
counts["session"] += 1
|
|
return real_session()
|
|
|
|
# layout, image prompt, verify verdict, chat reply — in call order.
|
|
completions_create = AsyncMock(
|
|
side_effect=[
|
|
_make_response("square", None),
|
|
_make_response("a detailed prompt", None),
|
|
_make_response("PASS", None),
|
|
_make_response("a chat reply", None),
|
|
]
|
|
)
|
|
image_response = MagicMock()
|
|
image_response.data = [MagicMock(b64_json=base64.b64encode(b"fake image").decode())]
|
|
images_generate = AsyncMock(return_value=image_response)
|
|
|
|
registry = MagicMock()
|
|
registry.to_openai_tools.return_value = []
|
|
|
|
db = ChatDatabase(db_path=temp_db_path)
|
|
|
|
with (
|
|
patch.object(openai, "AsyncOpenAI", CountingAsyncOpenAI),
|
|
patch.object(requests, "Session", counting_session),
|
|
patch.object(chat_client.chat.completions, "create", completions_create),
|
|
patch.object(gen_client.images, "generate", images_generate),
|
|
patch("vibe_bot.llm_client.embedding", return_value=[0.25] * 32),
|
|
):
|
|
asyncio.run(
|
|
ImageService(db, MagicMock()).generate(mock_ctx, message="a centaur")
|
|
)
|
|
asyncio.run(
|
|
ChatService(db, registry).handle(
|
|
mock_ctx,
|
|
bot_name="alfred",
|
|
message="hello",
|
|
system_prompt="you are a butler",
|
|
response_prefix="alfred response",
|
|
)
|
|
)
|
|
|
|
assert counts["async_openai"] == 0
|
|
assert counts["session"] == 0
|
|
assert llm_client.get_chat_client() is chat_client
|
|
assert llm_client.get_image_gen_client() is gen_client
|
|
assert llm_client.get_image_edit_client() is edit_client
|
|
|
|
|
|
def _make_response(content: str | None, tool_calls: list[object] | None) -> MagicMock:
|
|
"""Build a mock chat completion response with the given message fields."""
|
|
message = MagicMock()
|
|
message.content = content
|
|
message.tool_calls = tool_calls
|
|
return MagicMock(choices=[MagicMock(message=message)])
|
|
|
|
|
|
def test_chat_complete_skips_non_function_tool_call() -> None:
|
|
"""A tool call that is not of type 'function' is skipped, not executed."""
|
|
import asyncio
|
|
|
|
from vibe_bot.llm_client import chat_complete
|
|
|
|
called = {"n": 0}
|
|
|
|
def tool_executor(name: str, args: dict[str, str]) -> str:
|
|
called["n"] += 1
|
|
return f"executed:{name}"
|
|
|
|
custom_tool_call = MagicMock()
|
|
custom_tool_call.type = "custom"
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.chat.completions.create = AsyncMock(
|
|
side_effect=[
|
|
_make_response(content=None, tool_calls=[custom_tool_call]),
|
|
_make_response(content="final answer", tool_calls=None),
|
|
]
|
|
)
|
|
|
|
with patch("vibe_bot.llm.chat.get_chat_client", return_value=mock_client):
|
|
result = asyncio.run(
|
|
chat_complete(
|
|
[{"role": "user", "content": "hi"}],
|
|
model="m",
|
|
max_tokens=10,
|
|
tool_executor=tool_executor,
|
|
)
|
|
)
|
|
|
|
assert result == "final answer"
|
|
assert called["n"] == 0
|
|
|
|
|
|
def test_tool_registry_dispatch_and_unknown_tool() -> None:
|
|
"""The registry renders schemas, dispatches known tools, and names unknowns."""
|
|
from vibe_bot.llm_client import ToolRegistry
|
|
|
|
def echo_tool(name: str, args: dict[str, str], **kwargs: object) -> str:
|
|
return f"echo:{args.get('text', '')}"
|
|
|
|
registry = ToolRegistry()
|
|
registry.register(
|
|
"echo",
|
|
"Echoes the text argument back.",
|
|
{"type": "object", "properties": {"text": {"type": "string"}}},
|
|
echo_tool,
|
|
)
|
|
|
|
tools = registry.to_openai_tools()
|
|
first = tools[0]
|
|
assert first["type"] == "function"
|
|
function_def = cast("dict[str, object]", first["function"])
|
|
assert function_def["name"] == "echo"
|
|
assert function_def["description"] == "Echoes the text argument back."
|
|
|
|
assert registry.execute("echo", {"text": "hi"}) == "echo:hi"
|
|
assert registry.execute("does_not_exist", {}) == "Unknown tool: does_not_exist"
|