diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index c9f8a0b6..76020f2a 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -218,6 +218,10 @@ ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES" ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL" ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR" +# ZeroEntropy configuration (reranker only) +ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY" +ENV_RERANKER_ZEROENTROPY_MODEL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL" + ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION" ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION" @@ -360,6 +364,8 @@ DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0" DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0" +DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2" + # Vector extension (pgvector, vchord, or pgvectorscale) DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale" @@ -605,6 +611,8 @@ class HindsightConfig: reranker_litellm_sdk_api_key: str | None reranker_litellm_sdk_model: str reranker_litellm_sdk_api_base: str | None + reranker_zeroentropy_api_key: str | None + reranker_zeroentropy_model: str # Server host: str @@ -979,6 +987,9 @@ class HindsightConfig: reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY), reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL), reranker_litellm_sdk_api_base=os.getenv(ENV_RERANKER_LITELLM_SDK_API_BASE) or None, + # ZeroEntropy reranker + reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY), + reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL), # Server host=os.getenv(ENV_HOST, DEFAULT_HOST), port=int(os.getenv(ENV_PORT, DEFAULT_PORT)), diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 75ad687f..6d6016eb 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -29,6 +29,7 @@ from ..config import ( DEFAULT_RERANKER_PROVIDER, DEFAULT_RERANKER_TEI_BATCH_SIZE, DEFAULT_RERANKER_TEI_MAX_CONCURRENT, + DEFAULT_RERANKER_ZEROENTROPY_MODEL, ENV_RERANKER_COHERE_API_KEY, ENV_RERANKER_COHERE_MODEL, ENV_RERANKER_FLASHRANK_CACHE_DIR, @@ -42,6 +43,7 @@ from ..config import ( ENV_RERANKER_TEI_BATCH_SIZE, ENV_RERANKER_TEI_MAX_CONCURRENT, ENV_RERANKER_TEI_URL, + ENV_RERANKER_ZEROENTROPY_API_KEY, ) logger = logging.getLogger(__name__) @@ -556,6 +558,104 @@ class CohereCrossEncoder(CrossEncoderModel): return all_scores +class ZeroEntropyCrossEncoder(CrossEncoderModel): + """ + ZeroEntropy cross-encoder implementation using the ZeroEntropy Rerank API. + + Supports zerank-2 (flagship) and zerank-2-small models. + See: https://docs.zeroentropy.dev/models + """ + + RERANK_URL = "https://api.zeroentropy.dev/models/rerank" + + def __init__( + self, + api_key: str, + model: str = DEFAULT_RERANKER_ZEROENTROPY_MODEL, + timeout: float = 60.0, + ): + """ + Initialize ZeroEntropy cross-encoder client. + + Args: + api_key: ZeroEntropy API key + model: ZeroEntropy rerank model name (default: zerank-2) + timeout: Request timeout in seconds (default: 60.0) + """ + self.api_key = api_key + self.model = model + self.timeout = timeout + self._async_client: httpx.AsyncClient | None = None + + @property + def provider_name(self) -> str: + return "zeroentropy" + + async def initialize(self) -> None: + """Initialize the async HTTP client.""" + if self._async_client is not None: + return + + logger.info(f"Reranker: initializing ZeroEntropy provider with model {self.model}") + self._async_client = httpx.AsyncClient( + timeout=self.timeout, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + ) + logger.info("Reranker: ZeroEntropy provider initialized") + + async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: + """ + Score query-document pairs using the ZeroEntropy Rerank API. + + Args: + pairs: List of (query, document) tuples to score + + Returns: + List of relevance scores + """ + if self._async_client is None: + raise RuntimeError("Reranker not initialized. Call initialize() first.") + + if not pairs: + return [] + + # Group pairs by query for efficient batching + query_groups: dict[str, list[tuple[int, str]]] = {} + for idx, (query, text) in enumerate(pairs): + if query not in query_groups: + query_groups[query] = [] + query_groups[query].append((idx, text)) + + all_scores = [0.0] * len(pairs) + + for query, indexed_texts in query_groups.items(): + texts = [text for _, text in indexed_texts] + indices = [idx for idx, _ in indexed_texts] + + response = await self._async_client.post( + self.RERANK_URL, + json={ + "model": self.model, + "query": query, + "documents": texts, + "top_n": len(texts), + }, + ) + response.raise_for_status() + result = response.json() + + # Map scores back to original positions + for item in result.get("results", []): + original_idx = item["index"] + score = item["relevance_score"] + all_scores[indices[original_idx]] = score + + return all_scores + + class RRFPassthroughCrossEncoder(CrossEncoderModel): """ Passthrough cross-encoder that preserves RRF scores without neural reranking. @@ -1010,9 +1110,19 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: model=config.reranker_litellm_sdk_model, api_base=config.reranker_litellm_sdk_api_base, ) + elif provider == "zeroentropy": + api_key = config.reranker_zeroentropy_api_key + if not api_key: + raise ValueError( + f"{ENV_RERANKER_ZEROENTROPY_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'zeroentropy'" + ) + return ZeroEntropyCrossEncoder( + api_key=api_key, + model=config.reranker_zeroentropy_model, + ) elif provider == "rrf": return RRFPassthroughCrossEncoder() else: raise ValueError( - f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'litellm', 'litellm-sdk', 'rrf'" + f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf'" ) diff --git a/hindsight-api/tests/test_custom_embedding_dimension.py b/hindsight-api/tests/test_custom_embedding_dimension.py index 8d382481..fcee083d 100644 --- a/hindsight-api/tests/test_custom_embedding_dimension.py +++ b/hindsight-api/tests/test_custom_embedding_dimension.py @@ -15,7 +15,7 @@ import pytest from sqlalchemy import create_engine, text from hindsight_api import MemoryEngine, RequestContext -from hindsight_api.engine.cross_encoder import CohereCrossEncoder, LocalSTCrossEncoder +from hindsight_api.engine.cross_encoder import CohereCrossEncoder, LocalSTCrossEncoder, ZeroEntropyCrossEncoder from hindsight_api.engine.embeddings import CohereEmbeddings, LocalSTEmbeddings, OpenAIEmbeddings from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer from hindsight_api.engine.task_backend import SyncTaskBackend @@ -98,9 +98,7 @@ def get_row_count(db_url: str, schema: str = "public") -> int: """Get the number of rows with embeddings in memory_units.""" engine = create_engine(db_url) with engine.connect() as conn: - return conn.execute( - text(f"SELECT COUNT(*) FROM {schema}.memory_units WHERE embedding IS NOT NULL") - ).scalar() + return conn.execute(text(f"SELECT COUNT(*) FROM {schema}.memory_units WHERE embedding IS NOT NULL")).scalar() def insert_test_embedding(db_url: str, schema: str, dimension: int): @@ -610,3 +608,59 @@ class TestCohereIntegration: await memory.close() except Exception: pass + + +# ============================================================================= +# ZeroEntropy Reranker Tests +# ============================================================================= + + +def has_zeroentropy_api_key() -> bool: + """Check if ZeroEntropy API key is available.""" + return bool(os.environ.get("ZEROENTROPY_API_KEY")) + + +def get_zeroentropy_api_key() -> str: + """Get ZeroEntropy API key from environment.""" + return os.environ.get("ZEROENTROPY_API_KEY", "") + + +@pytest.fixture(scope="module") +def zeroentropy_cross_encoder(): + """Create ZeroEntropy cross-encoder instance.""" + if not has_zeroentropy_api_key(): + pytest.skip("ZeroEntropy API key not available (set ZEROENTROPY_API_KEY)") + + cross_encoder = ZeroEntropyCrossEncoder( + api_key=get_zeroentropy_api_key(), + model="zerank-2", + ) + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(cross_encoder.initialize()) + finally: + loop.close() + return cross_encoder + + +class TestZeroEntropyCrossEncoder: + """Tests for ZeroEntropy cross-encoder/reranker.""" + + def test_zeroentropy_cross_encoder_initialization(self, zeroentropy_cross_encoder): + """Test that ZeroEntropy cross-encoder initializes correctly.""" + assert zeroentropy_cross_encoder.provider_name == "zeroentropy" + + @pytest.mark.asyncio + async def test_zeroentropy_cross_encoder_predict(self, zeroentropy_cross_encoder): + """Test that ZeroEntropy cross-encoder can score pairs.""" + 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 zeroentropy_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" diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 35e8d3cf..2531780d 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -408,7 +408,7 @@ Supported OpenAI embedding dimensions: | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_RERANKER_PROVIDER` | Provider: `local`, `tei`, `cohere`, `flashrank`, `litellm`, `litellm-sdk`, or `rrf` | `local` | +| `HINDSIGHT_API_RERANKER_PROVIDER` | Provider: `local`, `tei`, `cohere`, `zeroentropy`, `flashrank`, `litellm`, `litellm-sdk`, or `rrf` | `local` | | `HINDSIGHT_API_RERANKER_LOCAL_MODEL` | Model for local provider | `cross-encoder/ms-marco-MiniLM-L-6-v2` | | `HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT` | Max concurrent local reranking (prevents CPU thrashing under load) | `4` | | `HINDSIGHT_API_RERANKER_LOCAL_TRUST_REMOTE_CODE` | Allow loading models with custom code (security risk, disabled by default) | `false` | @@ -424,6 +424,8 @@ Supported OpenAI embedding dimensions: | `HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY` | LiteLLM **SDK** API key for direct reranking (no proxy needed) | - | | `HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL` | LiteLLM SDK rerank model (e.g., `deepinfra/Qwen3-reranker-8B`) | `cohere/rerank-english-v3.0` | | `HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE` | Custom API base URL for LiteLLM SDK (optional) | - | +| `HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY` | ZeroEntropy API key for reranking | - | +| `HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL` | ZeroEntropy rerank model (`zerank-2`, `zerank-2-small`) | `zerank-2` | | `HINDSIGHT_API_RERANKER_FLASHRANK_MODEL` | FlashRank model for fast CPU-based reranking | `ms-marco-MiniLM-L-12-v2` | | `HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR` | Cache directory for FlashRank models | System default | @@ -453,6 +455,11 @@ export HINDSIGHT_API_RERANKER_COHERE_API_KEY=your-azure-api-key export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 export HINDSIGHT_API_RERANKER_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com +# ZeroEntropy - cloud-based reranking (state-of-the-art accuracy) +export HINDSIGHT_API_RERANKER_PROVIDER=zeroentropy +export HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY=your-api-key +export HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL=zerank-2 # or zerank-2-small + # LiteLLM proxy - unified gateway for multiple reranking providers (requires running LiteLLM proxy server) export HINDSIGHT_API_RERANKER_PROVIDER=litellm export HINDSIGHT_API_RERANKER_LITELLM_API_BASE=http://localhost:4000