e408b7e072
* feat: support litellm-sdk for reranker endpoint * feat: support litellm-sdk for reranker endpoint * fix: make litellm SDK cohere test fixture async function-scoped * fix: store litellm module reference during initialization to avoid import issues * feat: add LiteLLM SDK embeddings support - Add LiteLLMSDKEmbeddings class for direct API access without proxy - Support multiple providers: Cohere, OpenAI, Together AI, HuggingFace, Voyage AI - Automatic dimension detection via test embedding - Provider-specific API key mapping - Batch processing support (configurable batch size) - Comprehensive test coverage (17 unit tests) - Update documentation with configuration examples Implements embeddings in same PR as reranker per user request * fix: correct config mocking in embeddings factory tests - Mock get_config() from its source module (hindsight_api.config) - Fixes factory tests that were returning LocalSTEmbeddings instead of LiteLLMSDKEmbeddings - All 17 unit tests now passing * fix: skip Cohere integration tests when API key is invalid - Catch initialization errors and skip tests instead of failing - Prevents CI failures when COHERE_API_KEY is set but invalid - Integration tests now properly skip when authentication fails * fix: skip Cohere reranker integration tests when API key is invalid - Add same error handling as embeddings tests - Prevents CI failures when COHERE_API_KEY is set but invalid - Tests now properly skip when authentication fails * Revert "fix: skip Cohere reranker integration tests when API key is invalid" This reverts commit 655dacaffb25851ff48e202b4379fc8332a66df7. * Revert "fix: skip Cohere integration tests when API key is invalid" This reverts commit 5d00548e39faa3da6b427ace89589a216816f9e7. * fix: pass API key directly to litellm SDK functions - Add api_key parameter to arerank(), rerank(), aembedding(), and embedding() calls - Prevents authentication issues in multi-process environments (pytest-xdist) - More reliable than relying solely on environment variables - Update test assertions to expect api_key parameter * feat: pass api_base parameter to litellm SDK calls and remove hasattr check * fix: raise errors instead of silently returning 0.0 scores * refactor: pass API keys directly in kwargs instead of setting env vars
387 lines
14 KiB
Python
387 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",
|
|
)
|
|
|
|
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",
|
|
)
|
|
|
|
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])
|