Files
mnemosyne/mnemosyne/llm_manager/tests/test_models.py
Robert Helewka 99bdb4ac92 Add Themis application with custom widgets, views, and utilities
- Implemented custom form widgets for date, time, and datetime fields with DaisyUI styling.
- Created utility functions for formatting dates, times, and numbers according to user preferences.
- Developed views for profile settings, API key management, and notifications, including health check endpoints.
- Added URL configurations for Themis tests and main application routes.
- Established test cases for custom widgets to ensure proper functionality and integration.
- Defined project metadata and dependencies in pyproject.toml for package management.
2026-03-21 02:00:18 +00:00

237 lines
7.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_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())
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", "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_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()