15ea23d5d6
* feat: introduce hindsight-api-slim and hindsight-all-slim packages Closes #552 - Move all source code from hindsight-api/ to new hindsight-api-slim/ - hindsight-api-slim has heavy ML deps (torch, sentence-transformers, transformers, einops, flashrank, mlx, mlx-lm, safetensors) and pg0-embedded as optional extras: [local-ml], [embedded-db], [all] - hindsight-api becomes a zero-code meta-package depending on hindsight-api-slim[all] for full backward compatibility - Add hindsight-all-slim meta-package: hindsight-api-slim + client + embed - hindsight-all updated to depend on hindsight-api-slim[all] - pg0.py: lazy-import pg0 with clear ImportError pointing to [embedded-db] - Dockerfile: replace sed hack with proper uv sync --extra flags - Update release.yml, test.yml, lint.sh, release.sh, CLAUDE.md and all path references throughout the repo * refactor: rename hindsight/ directory to hindsight-all/ * docs: document hindsight-api-slim and hindsight-all-slim package variants Add package variants table and extras explanation to installation.md * docs: remove emojis from installation.md, use professional tone * docs: link Docker slim variant to pip package variants section * docs: consolidate Docker image variants into single table * ci: fix working-directory paths after package restructure - Replace all hindsight-api → hindsight-api-slim in test.yml - Replace hindsight → hindsight-all in test.yml - Add --extra embedded-db to test-embed API install step * ci: add local-ml and embedded-db extras to API sync steps These extras were previously implicit in the old hindsight-api package (which bundled everything). Now that hindsight-api-slim uses optional extras, we must explicitly request local-ml and embedded-db in CI. * ci: add API install step with embedded-db to test-embed smoke test The smoke test starts hindsight-api as a daemon, which requires pg0-embedded. Add a dedicated install step for hindsight-api-slim with embedded-db extra so the daemon can start successfully. * ci: remove --no-install-project when using optional extras When --no-install-project is combined with --extra, the optional deps are not installed because extras require the project to be active. Remove --no-install-project from steps that need local-ml or embedded-db. * ci: fix ordering of uv sync steps to preserve optional extras When uv sync runs for a different workspace member, it removes optional extras installed for other members. Fix by always running extra-requiring API sync last, after other workspace member syncs. Also remove --no-install-project from embedded-db sync in test-embed, as --no-install-project prevents optional extras from being active. * ci: add local-ml extra to test-embed API install for smoke test The smoke test starts the full API server which needs sentence-transformers for local embeddings (default provider). Add local-ml extra to the install. * ci: simplify extras with --all-extras and add slim pip smoke test - Replace explicit --extra local-ml --extra embedded-db with --all-extras for cleaner, more maintainable sync steps - Add test-pip-slim job: tests hindsight-api-slim[embedded-db] without local ML models, using Cohere for embeddings/reranking (mirrors Docker slim smoke test approach) * ci: simplify slim smoke test to health check only (mirrors Docker test)
389 lines
14 KiB
Python
389 lines
14 KiB
Python
"""
|
|
Tests for LiteLLM SDK embeddings implementation.
|
|
|
|
These tests cover:
|
|
1. Initialization (success, missing package, missing API key, idempotent)
|
|
2. Encode (single text, multiple texts, batching, error handling)
|
|
3. Provider-specific configuration (Cohere, OpenAI, etc.)
|
|
4. Factory function (create from env, validation errors)
|
|
5. Dimension detection
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from hindsight_api.config import (
|
|
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
|
ENV_EMBEDDINGS_LITELLM_SDK_MODEL,
|
|
ENV_EMBEDDINGS_PROVIDER,
|
|
HindsightConfig,
|
|
)
|
|
from hindsight_api.engine.embeddings import LiteLLMSDKEmbeddings, create_embeddings_from_env
|
|
|
|
|
|
class TestLiteLLMSDKEmbeddings:
|
|
"""Unit tests for LiteLLMSDKEmbeddings with mocked litellm responses."""
|
|
|
|
@pytest.fixture
|
|
def mock_litellm(self):
|
|
"""Mock litellm module."""
|
|
mock = MagicMock()
|
|
|
|
# Mock aembedding (async) for initialization
|
|
mock_response = MagicMock()
|
|
mock_response.data = [{"embedding": [0.1] * 768, "index": 0}]
|
|
mock.aembedding = AsyncMock(return_value=mock_response)
|
|
|
|
# Mock embedding (sync) for encode
|
|
mock_sync_response = MagicMock()
|
|
mock_sync_response.data = [
|
|
{"embedding": [0.1] * 768, "index": 0},
|
|
{"embedding": [0.2] * 768, "index": 1},
|
|
]
|
|
mock.embedding = MagicMock(return_value=mock_sync_response)
|
|
|
|
return mock
|
|
|
|
@pytest.fixture
|
|
async def embeddings(self, mock_litellm):
|
|
"""Create initialized LiteLLMSDKEmbeddings instance."""
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
# Manually set the mock (simulating successful initialization)
|
|
emb._litellm = mock_litellm
|
|
emb._dimension = 768
|
|
return emb
|
|
|
|
async def test_initialization_success(self, mock_litellm):
|
|
"""Test successful initialization."""
|
|
with patch("builtins.__import__", side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args)):
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
|
|
assert emb._litellm is None
|
|
assert emb._dimension is None
|
|
|
|
await emb.initialize()
|
|
|
|
assert emb._litellm is not None
|
|
assert emb._dimension == 768
|
|
|
|
# Verify test embedding was called
|
|
mock_litellm.aembedding.assert_called_once_with(
|
|
model="cohere/embed-english-v3.0",
|
|
input=["test"],
|
|
api_key="test_key",
|
|
encoding_format="float",
|
|
)
|
|
|
|
async def test_initialization_missing_package(self):
|
|
"""Test initialization fails gracefully when litellm is not installed."""
|
|
def mock_import(name, *args):
|
|
if name == "litellm":
|
|
raise ImportError("No module named 'litellm'")
|
|
return __import__(name, *args)
|
|
|
|
with patch("builtins.__import__", side_effect=mock_import):
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
|
|
with pytest.raises(ImportError, match="litellm is required"):
|
|
await emb.initialize()
|
|
|
|
async def test_initialization_idempotent(self, embeddings, mock_litellm):
|
|
"""Test that calling initialize() multiple times is safe."""
|
|
# embeddings._litellm is already set in fixture
|
|
assert embeddings._litellm is not None
|
|
|
|
# Call again
|
|
await embeddings.initialize()
|
|
|
|
# Should still have same litellm instance
|
|
assert embeddings._litellm is not None
|
|
|
|
async def test_encode_single_text(self, embeddings, mock_litellm):
|
|
"""Test encoding a single text."""
|
|
# Set up mock response
|
|
mock_litellm.embedding.return_value.data = [
|
|
{"embedding": [0.5] * 768, "index": 0},
|
|
]
|
|
|
|
result = embeddings.encode(["Hello world"])
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert len(result[0]) == 768
|
|
assert all(isinstance(x, float) for x in result[0])
|
|
assert all(abs(x - 0.5) < 0.001 for x in result[0])
|
|
|
|
# Verify call
|
|
mock_litellm.embedding.assert_called_once_with(
|
|
model="cohere/embed-english-v3.0",
|
|
input=["Hello world"],
|
|
api_key="test_key",
|
|
encoding_format="float",
|
|
)
|
|
|
|
async def test_encode_multiple_texts(self, embeddings, mock_litellm):
|
|
"""Test encoding multiple texts."""
|
|
# Set up mock response
|
|
mock_litellm.embedding.return_value.data = [
|
|
{"embedding": [0.1] * 768, "index": 0},
|
|
{"embedding": [0.2] * 768, "index": 1},
|
|
{"embedding": [0.3] * 768, "index": 2},
|
|
]
|
|
|
|
texts = ["First text", "Second text", "Third text"]
|
|
result = embeddings.encode(texts)
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
assert len(result[0]) == 768
|
|
assert len(result[1]) == 768
|
|
assert len(result[2]) == 768
|
|
assert all(abs(x - 0.1) < 0.001 for x in result[0])
|
|
assert all(abs(x - 0.2) < 0.001 for x in result[1])
|
|
assert all(abs(x - 0.3) < 0.001 for x in result[2])
|
|
|
|
async def test_encode_batching(self, embeddings, mock_litellm):
|
|
"""Test that large inputs are batched correctly."""
|
|
# Create embeddings with small batch size
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=2, # Small batch for testing
|
|
timeout=60.0,
|
|
)
|
|
emb._litellm = mock_litellm
|
|
emb._initialized = True
|
|
emb._dimension = 768
|
|
|
|
# Mock responses for each batch
|
|
def mock_embedding_side_effect(model, input, **kwargs):
|
|
mock_response = MagicMock()
|
|
mock_response.data = [
|
|
{"embedding": [float(i)] * 768, "index": i} for i in range(len(input))
|
|
]
|
|
return mock_response
|
|
|
|
mock_litellm.embedding.side_effect = mock_embedding_side_effect
|
|
|
|
# Encode 5 texts (should create 3 batches: 2, 2, 1)
|
|
texts = [f"Text {i}" for i in range(5)]
|
|
result = emb.encode(texts)
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 5
|
|
assert all(len(embedding) == 768 for embedding in result)
|
|
|
|
# Verify batching: should be called 3 times
|
|
assert mock_litellm.embedding.call_count == 3
|
|
|
|
# Verify batch sizes
|
|
calls = mock_litellm.embedding.call_args_list
|
|
assert len(calls[0][1]["input"]) == 2 # First batch
|
|
assert len(calls[1][1]["input"]) == 2 # Second batch
|
|
assert len(calls[2][1]["input"]) == 1 # Third batch
|
|
|
|
async def test_encode_empty_list(self, embeddings):
|
|
"""Test encoding empty list returns empty list."""
|
|
result = embeddings.encode([])
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 0
|
|
|
|
async def test_encode_before_initialization(self, mock_litellm):
|
|
"""Test that encode raises error if not initialized."""
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
emb.encode(["test"])
|
|
|
|
async def test_encode_error_handling(self, embeddings, mock_litellm):
|
|
"""Test error handling during encoding."""
|
|
# Make embedding raise an error
|
|
mock_litellm.embedding.side_effect = Exception("API Error")
|
|
|
|
with pytest.raises(Exception, match="API Error"):
|
|
embeddings.encode(["test"])
|
|
|
|
async def test_dimension_property(self, embeddings):
|
|
"""Test dimension property."""
|
|
assert embeddings.dimension == 768
|
|
|
|
async def test_dimension_before_initialization(self, mock_litellm):
|
|
"""Test dimension raises error if not initialized."""
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
_ = emb.dimension
|
|
|
|
async def test_custom_api_base(self, mock_litellm):
|
|
"""Test custom API base URL is passed to embedding calls."""
|
|
with patch("builtins.__import__", side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args)):
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key="test_key",
|
|
model="cohere/embed-english-v3.0",
|
|
api_base="https://custom.api.com",
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
|
|
await emb.initialize()
|
|
|
|
# Verify api_base is set
|
|
assert emb.api_base == "https://custom.api.com"
|
|
|
|
# Verify api_base is passed to aembedding
|
|
mock_litellm.aembedding.assert_called_once()
|
|
call_args = mock_litellm.aembedding.call_args
|
|
assert call_args.kwargs["api_base"] == "https://custom.api.com"
|
|
|
|
# Test encode also passes api_base
|
|
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
|
|
emb.encode(["test"])
|
|
|
|
mock_litellm.embedding.assert_called_once()
|
|
call_args = mock_litellm.embedding.call_args
|
|
assert call_args.kwargs["api_base"] == "https://custom.api.com"
|
|
|
|
|
|
class TestLiteLLMSDKEmbeddingsFactory:
|
|
"""Test the factory function for creating LiteLLM SDK embeddings."""
|
|
|
|
def test_create_from_env_success(self, monkeypatch):
|
|
"""Test creating embeddings from environment variables."""
|
|
# Mock get_config() to return configured HindsightConfig
|
|
mock_config = MagicMock()
|
|
mock_config.embeddings_provider = "litellm-sdk"
|
|
mock_config.embeddings_litellm_sdk_api_key = "test_key"
|
|
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
|
mock_config.embeddings_litellm_sdk_api_base = None
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
|
embeddings = create_embeddings_from_env()
|
|
|
|
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
|
assert embeddings.api_key == "test_key"
|
|
assert embeddings.model == "cohere/embed-english-v3.0"
|
|
|
|
def test_create_from_env_missing_api_key(self, monkeypatch):
|
|
"""Test that missing API key raises error."""
|
|
# Mock get_config() with missing API key
|
|
mock_config = MagicMock()
|
|
mock_config.embeddings_provider = "litellm-sdk"
|
|
mock_config.embeddings_litellm_sdk_api_key = None # Missing key
|
|
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
|
with pytest.raises(ValueError, match=ENV_EMBEDDINGS_LITELLM_SDK_API_KEY):
|
|
create_embeddings_from_env()
|
|
|
|
def test_create_from_env_with_api_base(self, monkeypatch):
|
|
"""Test creating embeddings with custom API base."""
|
|
# Mock get_config() with custom API base
|
|
mock_config = MagicMock()
|
|
mock_config.embeddings_provider = "litellm-sdk"
|
|
mock_config.embeddings_litellm_sdk_api_key = "test_key"
|
|
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
|
mock_config.embeddings_litellm_sdk_api_base = "https://custom.api.com"
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
|
embeddings = create_embeddings_from_env()
|
|
|
|
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
|
assert embeddings.api_base == "https://custom.api.com"
|
|
|
|
|
|
class TestLiteLLMSDKCohereEmbeddings:
|
|
"""Integration tests calling real Cohere API (matches CI pattern)."""
|
|
|
|
@pytest.fixture
|
|
async def litellm_cohere_embeddings(self):
|
|
"""Create embeddings instance with real Cohere API key."""
|
|
if not os.environ.get("COHERE_API_KEY"):
|
|
pytest.skip("Cohere API key not available")
|
|
|
|
emb = LiteLLMSDKEmbeddings(
|
|
api_key=os.environ["COHERE_API_KEY"],
|
|
model="cohere/embed-english-v3.0",
|
|
api_base=None,
|
|
batch_size=100,
|
|
timeout=60.0,
|
|
)
|
|
await emb.initialize()
|
|
return emb
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_sdk_cohere_encode(self, litellm_cohere_embeddings):
|
|
"""Test real Cohere API call for embeddings."""
|
|
texts = [
|
|
"The quick brown fox jumps over the lazy dog",
|
|
"Machine learning is a subset of artificial intelligence",
|
|
"Python is a popular programming language",
|
|
]
|
|
|
|
result = litellm_cohere_embeddings.encode(texts)
|
|
|
|
# Verify result type and shape
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
assert all(len(embedding) > 0 for embedding in result)
|
|
assert all(isinstance(x, float) for x in result[0])
|
|
|
|
# Verify embeddings are not zeros (common API failure mode)
|
|
for i, embedding in enumerate(result):
|
|
assert not all(abs(x) < 0.0001 for x in embedding), f"Embedding {i} is all zeros"
|
|
|
|
# Verify embeddings are normalized (Cohere returns normalized vectors)
|
|
for i, embedding in enumerate(result):
|
|
norm = sum(x * x for x in embedding) ** 0.5
|
|
assert 0.9 < norm < 1.1, f"Embedding {i} norm {norm} is not close to 1.0"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_sdk_cohere_dimension(self, litellm_cohere_embeddings):
|
|
"""Test dimension detection with real Cohere API."""
|
|
dimension = litellm_cohere_embeddings.dimension
|
|
|
|
# Cohere embed-english-v3.0 has 1024 dimensions
|
|
assert dimension == 1024
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_sdk_cohere_single_text(self, litellm_cohere_embeddings):
|
|
"""Test encoding single text with real Cohere API."""
|
|
result = litellm_cohere_embeddings.encode(["Hello world"])
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert len(result[0]) == 1024
|
|
assert not all(abs(x) < 0.0001 for x in result[0])
|