complete restructure
This commit is contained in:
@@ -0,0 +1,419 @@
|
||||
"""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"
|
||||
Reference in New Issue
Block a user