151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
"""Tests for the llama_wrapper module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import tempfile
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
|
|
from vibe_bot.config import (
|
|
CHAT_ENDPOINT,
|
|
CHAT_ENDPOINT_KEY,
|
|
CHAT_MODEL,
|
|
EMBEDDING_ENDPOINT,
|
|
EMBEDDING_ENDPOINT_KEY,
|
|
IMAGE_EDIT_ENDPOINT,
|
|
IMAGE_EDIT_ENDPOINT_KEY,
|
|
IMAGE_GEN_ENDPOINT,
|
|
IMAGE_GEN_ENDPOINT_KEY,
|
|
)
|
|
from vibe_bot.llama_wrapper import (
|
|
chat_completion,
|
|
chat_completion_instruct,
|
|
embedding,
|
|
image_edit,
|
|
image_generation,
|
|
)
|
|
|
|
TEMPDIR = Path(tempfile.mkdtemp())
|
|
|
|
|
|
def test_chat_completion_think() -> None:
|
|
"""Test chat completion with think model."""
|
|
chat_completion(
|
|
system_prompt="You are a helpful assistant.",
|
|
user_prompt="Tell me about Everquest",
|
|
openai_url=CHAT_ENDPOINT,
|
|
openai_api_key=CHAT_ENDPOINT_KEY,
|
|
model=CHAT_MODEL,
|
|
max_tokens=100,
|
|
)
|
|
|
|
|
|
def test_chat_completion_instruct() -> None:
|
|
"""Test chat completion with instruct model."""
|
|
chat_completion_instruct(
|
|
system_prompt="You are a helpful assistant.",
|
|
user_prompt="Tell me about Everquest",
|
|
openai_url=CHAT_ENDPOINT,
|
|
openai_api_key=CHAT_ENDPOINT_KEY,
|
|
model=CHAT_MODEL,
|
|
max_tokens=100,
|
|
)
|
|
|
|
|
|
def test_image_generation() -> None:
|
|
"""Test image generation endpoint."""
|
|
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
|
|
mock_response = MagicMock()
|
|
mock_data = MagicMock()
|
|
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
|
|
mock_response.data = [mock_data]
|
|
mock_openai.return_value.images.generate.return_value = mock_response
|
|
result = image_generation(
|
|
prompt="Generate an image of a horse",
|
|
openai_url=IMAGE_GEN_ENDPOINT,
|
|
openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
|
|
)
|
|
assert result == base64.b64encode(b"fake image data").decode()
|
|
|
|
|
|
def test_image_edit() -> None:
|
|
"""Test image edit endpoint."""
|
|
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
|
|
mock_response = MagicMock()
|
|
mock_data = MagicMock()
|
|
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
|
|
mock_response.data = [mock_data]
|
|
mock_openai.return_value.images.edit.return_value = mock_response
|
|
result = image_edit(
|
|
image=BytesIO(b"fake image"),
|
|
prompt="Paint the words 'horse' on the horse.",
|
|
openai_url=IMAGE_EDIT_ENDPOINT,
|
|
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
|
|
)
|
|
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 test_embeddings() -> None:
|
|
"""Test 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
|
|
|
|
with patch("vibe_bot.llama_wrapper.requests.post", side_effect=mock_post):
|
|
result1 = embedding(
|
|
"this is a horse",
|
|
openai_url=EMBEDDING_ENDPOINT,
|
|
openai_api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
result2 = embedding(
|
|
"this is a horse also",
|
|
openai_url=EMBEDDING_ENDPOINT,
|
|
openai_api_key=EMBEDDING_ENDPOINT_KEY,
|
|
model="embed",
|
|
)
|
|
result3 = embedding(
|
|
"this is a donkey",
|
|
openai_url=EMBEDDING_ENDPOINT,
|
|
openai_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
|