Robert Helewka
634845fee0
Implement hybrid search pipeline combining vector, fulltext, and graph search across Neo4j, with cross-attention reranking via Synesis (Qwen3-VL-Reranker-2B) `/v1/rerank` endpoint. - Add SearchService with vector, fulltext, and graph search strategies - Add SynesisRerankerClient for multimodal reranking via HTTP API - Add search API endpoint (POST /search/) with filtering by library, collection, and library_type - Add SearchRequest/Response serializers and image search results - Add "nonfiction" to library_type choices - Consolidate reranker stack from two models to single Synesis service - Handle image analysis_status as "skipped" when analysis is unavailable - Add comprehensive tests for search pipeline and reranker client
292 lines
9.7 KiB
Python
292 lines
9.7 KiB
Python
"""
|
|
Tests for LLM Manager models: LLMApi, LLMModel, LLMUsage.
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
|
|
from django.contrib.auth import get_user_model
|
|
from django.test import TestCase
|
|
|
|
from llm_manager.models import LLMApi, LLMModel, LLMUsage
|
|
|
|
User = get_user_model()
|
|
|
|
|
|
class LLMApiModelTest(TestCase):
|
|
"""Tests for the LLMApi model."""
|
|
|
|
def setUp(self):
|
|
self.user = User.objects.create_user(username="testuser", password="testpass123")
|
|
self.api = LLMApi.objects.create(
|
|
name="Test API",
|
|
api_type="openai",
|
|
base_url="https://api.example.com/v1",
|
|
is_active=True,
|
|
created_by=self.user,
|
|
)
|
|
|
|
def test_str(self):
|
|
self.assertEqual(str(self.api), "Test API (openai)")
|
|
|
|
def test_default_values(self):
|
|
self.assertTrue(self.api.is_active)
|
|
self.assertTrue(self.api.supports_streaming)
|
|
self.assertEqual(self.api.timeout_seconds, 60)
|
|
self.assertEqual(self.api.max_retries, 3)
|
|
self.assertEqual(self.api.last_test_status, "pending")
|
|
|
|
def test_uuid_primary_key(self):
|
|
self.assertIsNotNone(self.api.pk)
|
|
self.assertEqual(len(str(self.api.pk)), 36) # UUID format
|
|
|
|
def test_unique_name(self):
|
|
with self.assertRaises(Exception):
|
|
LLMApi.objects.create(
|
|
name="Test API",
|
|
api_type="ollama",
|
|
base_url="http://localhost:11434",
|
|
)
|
|
|
|
|
|
class LLMApiEncryptionTest(TestCase):
|
|
"""Tests for API key encryption."""
|
|
|
|
def test_api_key_encrypted_at_rest(self):
|
|
"""API key should be encrypted in the database."""
|
|
api = LLMApi.objects.create(
|
|
name="Encrypted Test",
|
|
api_type="openai",
|
|
base_url="https://api.example.com/v1",
|
|
api_key="sk-test-secret-key-12345",
|
|
)
|
|
# Re-fetch from database
|
|
api_fresh = LLMApi.objects.get(pk=api.pk)
|
|
self.assertEqual(api_fresh.api_key, "sk-test-secret-key-12345")
|
|
|
|
def test_blank_api_key(self):
|
|
api = LLMApi.objects.create(
|
|
name="No Key",
|
|
api_type="ollama",
|
|
base_url="http://localhost:11434",
|
|
api_key="",
|
|
)
|
|
api_fresh = LLMApi.objects.get(pk=api.pk)
|
|
self.assertEqual(api_fresh.api_key, "")
|
|
|
|
|
|
class LLMModelModelTest(TestCase):
|
|
"""Tests for the LLMModel model."""
|
|
|
|
def setUp(self):
|
|
self.api = LLMApi.objects.create(
|
|
name="Test API",
|
|
api_type="openai",
|
|
base_url="https://api.example.com/v1",
|
|
)
|
|
self.model = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="gpt-4o",
|
|
display_name="GPT-4o",
|
|
model_type="chat",
|
|
context_window=128000,
|
|
max_output_tokens=16384,
|
|
input_cost_per_1k=Decimal("0.0025"),
|
|
output_cost_per_1k=Decimal("0.01"),
|
|
)
|
|
|
|
def test_str(self):
|
|
self.assertEqual(str(self.model), "Test API: gpt-4o")
|
|
|
|
def test_unique_together(self):
|
|
"""Model name must be unique per API."""
|
|
with self.assertRaises(Exception):
|
|
LLMModel.objects.create(
|
|
api=self.api,
|
|
name="gpt-4o",
|
|
model_type="chat",
|
|
context_window=8192,
|
|
)
|
|
|
|
def test_model_types(self):
|
|
"""All model types should be creatable."""
|
|
for mtype in ["embedding", "vision", "audio", "reranker", "multimodal_embed"]:
|
|
m = LLMModel.objects.create(
|
|
api=self.api,
|
|
name=f"test-{mtype}",
|
|
model_type=mtype,
|
|
context_window=8192,
|
|
)
|
|
self.assertEqual(m.model_type, mtype)
|
|
|
|
def test_mnemosyne_fields(self):
|
|
"""Mnemosyne-specific fields: supports_multimodal, vector_dimensions."""
|
|
embed = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="text-embedding-3-large",
|
|
model_type="embedding",
|
|
context_window=8191,
|
|
vector_dimensions=3072,
|
|
supports_multimodal=False,
|
|
)
|
|
self.assertEqual(embed.vector_dimensions, 3072)
|
|
self.assertFalse(embed.supports_multimodal)
|
|
|
|
def test_get_system_embedding_model(self):
|
|
embed = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="embed-model",
|
|
model_type="embedding",
|
|
context_window=8191,
|
|
is_system_embedding_model=True,
|
|
)
|
|
result = LLMModel.get_system_embedding_model()
|
|
self.assertEqual(result.pk, embed.pk)
|
|
|
|
def test_get_system_chat_model(self):
|
|
self.model.is_system_chat_model = True
|
|
self.model.save()
|
|
result = LLMModel.get_system_chat_model()
|
|
self.assertEqual(result.pk, self.model.pk)
|
|
|
|
def test_get_system_reranker_model(self):
|
|
reranker = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="reranker-model",
|
|
model_type="reranker",
|
|
context_window=8192,
|
|
is_system_reranker_model=True,
|
|
)
|
|
result = LLMModel.get_system_reranker_model()
|
|
self.assertEqual(result.pk, reranker.pk)
|
|
|
|
def test_get_system_vision_model_with_vision_type(self):
|
|
vision = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="vision-model",
|
|
model_type="vision",
|
|
context_window=8192,
|
|
is_system_vision_model=True,
|
|
)
|
|
result = LLMModel.get_system_vision_model()
|
|
self.assertEqual(result.pk, vision.pk)
|
|
|
|
def test_get_system_vision_model_with_chat_type(self):
|
|
"""Vision-capable chat models can serve as system vision model."""
|
|
self.model.is_system_vision_model = True
|
|
self.model.save()
|
|
result = LLMModel.get_system_vision_model()
|
|
self.assertEqual(result.pk, self.model.pk)
|
|
|
|
def test_get_system_vision_model_excludes_embedding_type(self):
|
|
"""Embedding models should not be returned as vision model."""
|
|
embed = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="embed-only",
|
|
model_type="embedding",
|
|
context_window=8191,
|
|
is_system_vision_model=True,
|
|
)
|
|
result = LLMModel.get_system_vision_model()
|
|
self.assertIsNone(result)
|
|
|
|
def test_get_system_vision_model_excludes_inactive(self):
|
|
LLMModel.objects.create(
|
|
api=self.api,
|
|
name="inactive-vision",
|
|
model_type="vision",
|
|
context_window=8192,
|
|
is_system_vision_model=True,
|
|
is_active=False,
|
|
)
|
|
result = LLMModel.get_system_vision_model()
|
|
self.assertIsNone(result)
|
|
|
|
def test_get_system_model_returns_none(self):
|
|
"""Returns None when no system model is configured."""
|
|
self.assertIsNone(LLMModel.get_system_embedding_model())
|
|
self.assertIsNone(LLMModel.get_system_chat_model())
|
|
self.assertIsNone(LLMModel.get_system_reranker_model())
|
|
self.assertIsNone(LLMModel.get_system_vision_model())
|
|
|
|
|
|
class LLMUsageModelTest(TestCase):
|
|
"""Tests for the LLMUsage model."""
|
|
|
|
def setUp(self):
|
|
self.user = User.objects.create_user(username="testuser", password="testpass123")
|
|
self.api = LLMApi.objects.create(
|
|
name="Test API",
|
|
api_type="openai",
|
|
base_url="https://api.example.com/v1",
|
|
)
|
|
self.model = LLMModel.objects.create(
|
|
api=self.api,
|
|
name="gpt-4o",
|
|
model_type="chat",
|
|
context_window=128000,
|
|
input_cost_per_1k=Decimal("0.0025"),
|
|
output_cost_per_1k=Decimal("0.01"),
|
|
)
|
|
|
|
def test_cost_calculation(self):
|
|
"""Total cost is auto-calculated on save."""
|
|
usage = LLMUsage.objects.create(
|
|
user=self.user,
|
|
model=self.model,
|
|
input_tokens=1000,
|
|
output_tokens=500,
|
|
purpose="other",
|
|
)
|
|
# 1000/1000 * 0.0025 + 500/1000 * 0.01 = 0.0025 + 0.005 = 0.0075
|
|
self.assertAlmostEqual(float(usage.total_cost), 0.0075, places=4)
|
|
|
|
def test_cost_with_cached_tokens(self):
|
|
self.model.cached_cost_per_1k = Decimal("0.00125")
|
|
self.model.save()
|
|
usage = LLMUsage.objects.create(
|
|
user=self.user,
|
|
model=self.model,
|
|
input_tokens=1000,
|
|
output_tokens=500,
|
|
cached_tokens=2000,
|
|
purpose="responder",
|
|
)
|
|
# 0.0025 + 0.005 + 2000/1000 * 0.00125 = 0.0025 + 0.005 + 0.0025 = 0.01
|
|
self.assertAlmostEqual(float(usage.total_cost), 0.01, places=4)
|
|
|
|
def test_purpose_choices(self):
|
|
for purpose in ["responder", "reviewer", "embeddings", "search", "reranking", "multimodal_embed", "vision_analysis", "other"]:
|
|
usage = LLMUsage.objects.create(
|
|
user=self.user,
|
|
model=self.model,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
purpose=purpose,
|
|
)
|
|
self.assertEqual(usage.purpose, purpose)
|
|
|
|
def test_vision_analysis_purpose(self):
|
|
"""Vision analysis usage can be tracked."""
|
|
usage = LLMUsage.objects.create(
|
|
user=self.user,
|
|
model=self.model,
|
|
input_tokens=500,
|
|
output_tokens=200,
|
|
purpose="vision_analysis",
|
|
)
|
|
self.assertEqual(usage.purpose, "vision_analysis")
|
|
self.assertGreater(usage.total_cost, 0)
|
|
|
|
def test_protect_model_delete(self):
|
|
"""Deleting a model with usage records should raise ProtectedError."""
|
|
LLMUsage.objects.create(
|
|
user=self.user,
|
|
model=self.model,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
)
|
|
from django.db.models import ProtectedError
|
|
|
|
with self.assertRaises(ProtectedError):
|
|
self.model.delete()
|