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)
392 lines
14 KiB
Python
392 lines
14 KiB
Python
"""
|
|
Tests for LiteLLMSDKCrossEncoder.
|
|
|
|
Tests the LiteLLM SDK-based cross-encoder implementation for reranking.
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from hindsight_api.engine.cross_encoder import LiteLLMSDKCrossEncoder, create_cross_encoder_from_env
|
|
|
|
|
|
class TestLiteLLMSDKCrossEncoder:
|
|
"""Test suite for LiteLLMSDKCrossEncoder class."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialization_success(self):
|
|
"""Test successful initialization with valid config."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="deepinfra/Qwen3-reranker-8B",
|
|
)
|
|
|
|
assert encoder.provider_name == "litellm-sdk"
|
|
assert encoder.api_key == "test_key"
|
|
assert encoder.model == "deepinfra/Qwen3-reranker-8B"
|
|
assert encoder._initialized is False
|
|
|
|
# Mock the litellm import
|
|
mock_litellm = MagicMock()
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
assert encoder._initialized is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialization_missing_package(self):
|
|
"""Test initialization fails when litellm package is missing."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
with patch.dict("sys.modules", {"litellm": None}):
|
|
with pytest.raises(ImportError, match="litellm is required"):
|
|
await encoder.initialize()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initialization_idempotent(self):
|
|
"""Test that calling initialize() multiple times is safe."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
mock_litellm = MagicMock()
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
assert encoder._initialized is True
|
|
|
|
# Second call should be no-op
|
|
await encoder.initialize()
|
|
assert encoder._initialized is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_predict_single_query(self):
|
|
"""Test prediction with a single query and multiple documents."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="deepinfra/Qwen3-reranker-8B",
|
|
)
|
|
|
|
# Create mock response with results as TypedDicts
|
|
mock_response = MagicMock()
|
|
mock_response.results = [
|
|
{"index": 0, "relevance_score": 0.9},
|
|
{"index": 1, "relevance_score": 0.7},
|
|
{"index": 2, "relevance_score": 0.5},
|
|
]
|
|
|
|
mock_litellm = MagicMock()
|
|
mock_litellm.arerank = AsyncMock(return_value=mock_response)
|
|
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
|
|
pairs = [
|
|
("What is Python?", "Python is a programming language"),
|
|
("What is Python?", "Python is a snake"),
|
|
("What is Python?", "Python is a British comedy group"),
|
|
]
|
|
|
|
scores = await encoder.predict(pairs)
|
|
|
|
assert len(scores) == 3
|
|
assert scores == [0.9, 0.7, 0.5]
|
|
|
|
# Verify arerank was called correctly
|
|
mock_litellm.arerank.assert_called_once()
|
|
call_args = mock_litellm.arerank.call_args
|
|
assert call_args.kwargs["model"] == "deepinfra/Qwen3-reranker-8B"
|
|
assert call_args.kwargs["query"] == "What is Python?"
|
|
assert len(call_args.kwargs["documents"]) == 3
|
|
assert call_args.kwargs["api_key"] == "test_key"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_predict_multiple_queries(self):
|
|
"""Test prediction with multiple different queries (grouped efficiently)."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
# First query response
|
|
mock_response1 = MagicMock()
|
|
mock_response1.results = [
|
|
{"index": 0, "relevance_score": 0.9},
|
|
{"index": 1, "relevance_score": 0.7},
|
|
]
|
|
|
|
# Second query response
|
|
mock_response2 = MagicMock()
|
|
mock_response2.results = [
|
|
{"index": 0, "relevance_score": 0.8},
|
|
]
|
|
|
|
mock_litellm = MagicMock()
|
|
mock_litellm.arerank = AsyncMock(side_effect=[mock_response1, mock_response2])
|
|
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
|
|
pairs = [
|
|
("What is Python?", "Python is a programming language"),
|
|
("What is Python?", "Python is a snake"),
|
|
("What is Java?", "Java is a programming language"),
|
|
]
|
|
|
|
scores = await encoder.predict(pairs)
|
|
|
|
assert len(scores) == 3
|
|
assert scores[0] == 0.9 # First query, first doc
|
|
assert scores[1] == 0.7 # First query, second doc
|
|
assert scores[2] == 0.8 # Second query, first doc
|
|
|
|
# Verify arerank was called twice (once per unique query)
|
|
assert mock_litellm.arerank.call_count == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_predict_empty_pairs(self):
|
|
"""Test prediction with empty input."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
mock_litellm = MagicMock()
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
scores = await encoder.predict([])
|
|
assert scores == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_predict_not_initialized(self):
|
|
"""Test that predict fails if encoder not initialized."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
pairs = [("query", "document")]
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
await encoder.predict(pairs)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_predict_error_handling(self):
|
|
"""Test that errors during prediction are raised."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
|
|
# Mock litellm to raise an error
|
|
mock_litellm = MagicMock()
|
|
mock_litellm.arerank = AsyncMock(side_effect=Exception("API Error"))
|
|
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
|
|
pairs = [
|
|
("What is Python?", "Python is a programming language"),
|
|
]
|
|
|
|
# Should raise the exception
|
|
with pytest.raises(Exception, match="API Error"):
|
|
await encoder.predict(pairs)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_api_base(self):
|
|
"""Test that custom API base URL is passed to rerank calls."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="cohere/rerank-english-v3.0",
|
|
api_base="https://custom.api.example.com",
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.results = [
|
|
{"index": 0, "relevance_score": 0.9},
|
|
]
|
|
|
|
mock_litellm = MagicMock()
|
|
mock_litellm.arerank = AsyncMock(return_value=mock_response)
|
|
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
|
|
# Test that api_base is passed to arerank
|
|
pairs = [("query", "document")]
|
|
scores = await encoder.predict(pairs)
|
|
|
|
assert scores == [0.9]
|
|
mock_litellm.arerank.assert_called_once()
|
|
call_args = mock_litellm.arerank.call_args
|
|
assert call_args.kwargs["api_base"] == "https://custom.api.example.com"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_response_with_direct_score_list(self):
|
|
"""Test handling of response format with direct score list."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key="test_key",
|
|
model="some-provider/model",
|
|
)
|
|
|
|
# Mock litellm to return direct list of scores
|
|
mock_litellm = MagicMock()
|
|
mock_litellm.arerank = AsyncMock(return_value=[0.9, 0.7, 0.5])
|
|
|
|
with patch.dict("sys.modules", {"litellm": mock_litellm}):
|
|
await encoder.initialize()
|
|
|
|
pairs = [
|
|
("query", "doc1"),
|
|
("query", "doc2"),
|
|
("query", "doc3"),
|
|
]
|
|
|
|
scores = await encoder.predict(pairs)
|
|
|
|
assert scores == [0.9, 0.7, 0.5]
|
|
|
|
|
|
class TestFactoryFunction:
|
|
"""Test suite for create_cross_encoder_from_env factory function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_litellm_sdk_from_env(self):
|
|
"""Test creating LiteLLM SDK cross-encoder from environment variables."""
|
|
env_vars = {
|
|
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B",
|
|
}
|
|
|
|
with patch.dict(os.environ, env_vars, clear=False):
|
|
# Need to reload config to pick up env vars
|
|
from hindsight_api.config import HindsightConfig
|
|
|
|
config = HindsightConfig.from_env()
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=config):
|
|
encoder = create_cross_encoder_from_env()
|
|
|
|
assert isinstance(encoder, LiteLLMSDKCrossEncoder)
|
|
assert encoder.api_key == "test_key"
|
|
assert encoder.model == "deepinfra/Qwen3-reranker-8B"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_litellm_sdk_missing_api_key(self):
|
|
"""Test that factory raises error when API key is missing."""
|
|
env_vars = {
|
|
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "deepinfra/Qwen3-reranker-8B",
|
|
}
|
|
|
|
with patch.dict(os.environ, env_vars, clear=False):
|
|
# Remove API key if set
|
|
if "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY" in os.environ:
|
|
del os.environ["HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"]
|
|
|
|
from hindsight_api.config import HindsightConfig
|
|
|
|
config = HindsightConfig.from_env()
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=config):
|
|
with pytest.raises(ValueError, match="HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY is required"):
|
|
create_cross_encoder_from_env()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_litellm_sdk_with_custom_api_base(self):
|
|
"""Test creating LiteLLM SDK cross-encoder with custom API base."""
|
|
env_vars = {
|
|
"HINDSIGHT_API_RERANKER_PROVIDER": "litellm-sdk",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY": "test_key",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL": "cohere/rerank-english-v3.0",
|
|
"HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE": "https://custom.api.example.com",
|
|
}
|
|
|
|
with patch.dict(os.environ, env_vars, clear=False):
|
|
from hindsight_api.config import HindsightConfig
|
|
|
|
config = HindsightConfig.from_env()
|
|
|
|
with patch("hindsight_api.config.get_config", return_value=config):
|
|
encoder = create_cross_encoder_from_env()
|
|
|
|
assert isinstance(encoder, LiteLLMSDKCrossEncoder)
|
|
assert encoder.api_base == "https://custom.api.example.com"
|
|
|
|
|
|
class TestLiteLLMSDKCohereCrossEncoder:
|
|
"""Tests for LiteLLM SDK calling Cohere (runs in CI with COHERE_API_KEY)."""
|
|
|
|
@pytest.fixture
|
|
async def litellm_cohere_cross_encoder(self):
|
|
"""Create LiteLLM SDK cross-encoder instance for Cohere."""
|
|
if not os.environ.get("COHERE_API_KEY"):
|
|
pytest.skip("Cohere API key not available (set COHERE_API_KEY)")
|
|
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key=os.environ["COHERE_API_KEY"],
|
|
model="cohere/rerank-english-v3.0",
|
|
)
|
|
await encoder.initialize()
|
|
return encoder
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_sdk_cohere_initialization(self, litellm_cohere_cross_encoder):
|
|
"""Test that LiteLLM SDK Cohere cross-encoder initializes correctly."""
|
|
assert litellm_cohere_cross_encoder.provider_name == "litellm-sdk"
|
|
assert litellm_cohere_cross_encoder.model == "cohere/rerank-english-v3.0"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_litellm_sdk_cohere_predict(self, litellm_cohere_cross_encoder):
|
|
"""Test that LiteLLM SDK can call Cohere rerank API."""
|
|
pairs = [
|
|
("What is the capital of France?", "Paris is the capital of France."),
|
|
("What is the capital of France?", "The Eiffel Tower is in Paris."),
|
|
("What is the capital of France?", "Python is a programming language."),
|
|
]
|
|
scores = await litellm_cohere_cross_encoder.predict(pairs)
|
|
|
|
assert len(scores) == 3
|
|
assert all(isinstance(s, float) for s in scores)
|
|
# The first result should be most relevant
|
|
assert scores[0] > scores[2], "Direct answer should score higher than unrelated text"
|
|
# All scores should be in valid range
|
|
assert all(0.0 <= score <= 1.0 for score in scores)
|
|
|
|
|
|
class TestIntegration:
|
|
"""Integration tests with real API (optional - requires API keys)."""
|
|
|
|
@pytest.mark.skipif(
|
|
not os.environ.get("DEEPINFRA_API_KEY"),
|
|
reason="DEEPINFRA_API_KEY not set - skipping integration test",
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_real_deepinfra_api(self):
|
|
"""Test with real DeepInfra API (requires DEEPINFRA_API_KEY env var)."""
|
|
encoder = LiteLLMSDKCrossEncoder(
|
|
api_key=os.environ["DEEPINFRA_API_KEY"],
|
|
model="deepinfra/Qwen3-reranker-8B",
|
|
)
|
|
|
|
await encoder.initialize()
|
|
|
|
pairs = [
|
|
("What is Python?", "Python is a high-level programming language"),
|
|
("What is Python?", "Python is a species of snake"),
|
|
("What is Python?", "Python is unrelated text about cars"),
|
|
]
|
|
|
|
scores = await encoder.predict(pairs)
|
|
|
|
# First doc should have highest score (most relevant)
|
|
assert len(scores) == 3
|
|
assert scores[0] > scores[1]
|
|
assert scores[1] > scores[2]
|
|
assert all(0.0 <= score <= 1.0 for score in scores)
|