From d47c8a28cce998f21e1752684ebf81680f1dc901 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 13 Jan 2026 16:55:28 +0100 Subject: [PATCH] feat: support litellm gateway (#154) --- hindsight-api/hindsight_api/config.py | 11 ++ .../hindsight_api/engine/cross_encoder.py | 122 +++++++++++++++- .../hindsight_api/engine/embeddings.py | 131 +++++++++++++++++- .../docs/developer/configuration.md | 27 +++- 4 files changed, 287 insertions(+), 4 deletions(-) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index d6c4f078..243fd24a 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -49,6 +49,12 @@ ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL" ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL" ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL" +# LiteLLM gateway configuration (for embeddings and reranker via LiteLLM proxy) +ENV_LITELLM_API_BASE = "HINDSIGHT_API_LITELLM_API_BASE" +ENV_LITELLM_API_KEY = "HINDSIGHT_API_LITELLM_API_KEY" +ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL" +ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL" + ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER" ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL" ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT" @@ -124,6 +130,11 @@ 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" +# LiteLLM defaults +DEFAULT_LITELLM_API_BASE = "http://localhost:4000" +DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small" +DEFAULT_RERANKER_LITELLM_MODEL = "cohere/rerank-english-v3.0" + DEFAULT_HOST = "0.0.0.0" DEFAULT_PORT = 8888 DEFAULT_LOG_LEVEL = "info" diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 16ef2fc1..3a90c09e 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -15,19 +15,24 @@ from concurrent.futures import ThreadPoolExecutor import httpx from ..config import ( + DEFAULT_LITELLM_API_BASE, DEFAULT_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_MODEL, + DEFAULT_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT, DEFAULT_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_PROVIDER, DEFAULT_RERANKER_TEI_BATCH_SIZE, DEFAULT_RERANKER_TEI_MAX_CONCURRENT, ENV_COHERE_API_KEY, + ENV_LITELLM_API_BASE, + ENV_LITELLM_API_KEY, ENV_RERANKER_COHERE_BASE_URL, ENV_RERANKER_COHERE_MODEL, ENV_RERANKER_FLASHRANK_CACHE_DIR, ENV_RERANKER_FLASHRANK_MODEL, + ENV_RERANKER_LITELLM_MODEL, ENV_RERANKER_LOCAL_MAX_CONCURRENT, ENV_RERANKER_LOCAL_MODEL, ENV_RERANKER_PROVIDER, @@ -651,6 +656,116 @@ class FlashRankCrossEncoder(CrossEncoderModel): return await loop.run_in_executor(FlashRankCrossEncoder._executor, self._predict_sync, pairs) +class LiteLLMCrossEncoder(CrossEncoderModel): + """ + LiteLLM cross-encoder implementation using LiteLLM proxy's /rerank endpoint. + + LiteLLM provides a unified interface for multiple reranking providers via + the Cohere-compatible /rerank endpoint. + See: https://docs.litellm.ai/docs/rerank + + Supported providers via LiteLLM: + - Cohere (rerank-english-v3.0, etc.) - prefix with cohere/ + - Together AI - prefix with together_ai/ + - Azure AI - prefix with azure_ai/ + - Jina AI - prefix with jina_ai/ + - AWS Bedrock - prefix with bedrock/ + - Voyage AI - prefix with voyage/ + """ + + def __init__( + self, + api_base: str = DEFAULT_LITELLM_API_BASE, + api_key: str | None = None, + model: str = DEFAULT_RERANKER_LITELLM_MODEL, + timeout: float = 60.0, + ): + """ + Initialize LiteLLM cross-encoder client. + + Args: + api_base: Base URL of the LiteLLM proxy (default: http://localhost:4000) + api_key: API key for the LiteLLM proxy (optional, depends on proxy config) + model: Reranking model name (default: cohere/rerank-english-v3.0) + Use provider prefix (e.g., cohere/, together_ai/, voyage/) + timeout: Request timeout in seconds (default: 60.0) + """ + self.api_base = api_base.rstrip("/") + 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 "litellm" + + async def initialize(self) -> None: + """Initialize the async HTTP client.""" + if self._async_client is not None: + return + + logger.info(f"Reranker: initializing LiteLLM provider at {self.api_base} with model {self.model}") + + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + self._async_client = httpx.AsyncClient(timeout=self.timeout, headers=headers) + logger.info("Reranker: LiteLLM provider initialized") + + async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: + """ + Score query-document pairs using the LiteLLM proxy's /rerank endpoint. + + 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 (LiteLLM rerank expects one query with multiple documents) + 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] + + # LiteLLM /rerank follows Cohere API format + response = await self._async_client.post( + f"{self.api_base}/rerank", + json={ + "model": self.model, + "query": query, + "documents": texts, + "top_n": len(texts), # Return all scores + }, + ) + response.raise_for_status() + result = response.json() + + # Map scores back to original positions + # Response format: {"results": [{"index": 0, "relevance_score": 0.9}, ...]} + for item in result.get("results", []): + original_idx = item["index"] + score = item.get("relevance_score", item.get("score", 0.0)) + all_scores[indices[original_idx]] = score + + return all_scores + + def create_cross_encoder_from_env() -> CrossEncoderModel: """ Create a CrossEncoderModel instance based on environment variables. @@ -687,9 +802,14 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL) cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR) return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir) + elif provider == "litellm": + api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE) + api_key = os.environ.get(ENV_LITELLM_API_KEY) + model = os.environ.get(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL) + return LiteLLMCrossEncoder(api_base=api_base, api_key=api_key, model=model) elif provider == "rrf": return RRFPassthroughCrossEncoder() else: raise ValueError( - f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'rrf'" + f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'flashrank', 'litellm', 'rrf'" ) diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index ee7a558c..066406fa 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -17,18 +17,23 @@ import httpx from ..config import ( DEFAULT_EMBEDDINGS_COHERE_MODEL, + DEFAULT_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_PROVIDER, + DEFAULT_LITELLM_API_BASE, ENV_COHERE_API_KEY, ENV_EMBEDDINGS_COHERE_BASE_URL, ENV_EMBEDDINGS_COHERE_MODEL, + ENV_EMBEDDINGS_LITELLM_MODEL, ENV_EMBEDDINGS_LOCAL_MODEL, ENV_EMBEDDINGS_OPENAI_API_KEY, ENV_EMBEDDINGS_OPENAI_BASE_URL, ENV_EMBEDDINGS_OPENAI_MODEL, ENV_EMBEDDINGS_PROVIDER, ENV_EMBEDDINGS_TEI_URL, + ENV_LITELLM_API_BASE, + ENV_LITELLM_API_KEY, ENV_LLM_API_KEY, ) @@ -549,6 +554,123 @@ class CohereEmbeddings(Embeddings): return all_embeddings +class LiteLLMEmbeddings(Embeddings): + """ + LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint. + + LiteLLM provides a unified interface for multiple embedding providers. + The proxy exposes an OpenAI-compatible /embeddings endpoint. + See: https://docs.litellm.ai/docs/embedding/supported_embedding + + Supported providers via LiteLLM: + - OpenAI (text-embedding-3-small, text-embedding-ada-002, etc.) + - Cohere (embed-english-v3.0, etc.) - prefix with cohere/ + - Vertex AI (textembedding-gecko, etc.) - prefix with vertex_ai/ + - HuggingFace, Mistral, Voyage AI, etc. + + The embedding dimension is auto-detected from the model at initialization. + """ + + def __init__( + self, + api_base: str = DEFAULT_LITELLM_API_BASE, + api_key: str | None = None, + model: str = DEFAULT_EMBEDDINGS_LITELLM_MODEL, + batch_size: int = 100, + timeout: float = 60.0, + ): + """ + Initialize LiteLLM embeddings client. + + Args: + api_base: Base URL of the LiteLLM proxy (default: http://localhost:4000) + api_key: API key for the LiteLLM proxy (optional, depends on proxy config) + model: Embedding model name (default: text-embedding-3-small) + Use provider prefix for non-OpenAI models (e.g., cohere/embed-english-v3.0) + batch_size: Maximum batch size for embedding requests (default: 100) + timeout: Request timeout in seconds (default: 60.0) + """ + self.api_base = api_base.rstrip("/") + self.api_key = api_key + self.model = model + self.batch_size = batch_size + self.timeout = timeout + self._client: httpx.Client | None = None + self._dimension: int | None = None + + @property + def provider_name(self) -> str: + return "litellm" + + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + return self._dimension + + async def initialize(self) -> None: + """Initialize the HTTP client and detect embedding dimension.""" + if self._client is not None: + return + + logger.info(f"Embeddings: initializing LiteLLM provider at {self.api_base} with model {self.model}") + + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + self._client = httpx.Client(timeout=self.timeout, headers=headers) + + # Do a test embedding to detect dimension + try: + response = self._client.post( + f"{self.api_base}/embeddings", + json={"model": self.model, "input": ["test"]}, + ) + response.raise_for_status() + result = response.json() + if result.get("data") and len(result["data"]) > 0: + self._dimension = len(result["data"][0]["embedding"]) + logger.info(f"Embeddings: LiteLLM provider initialized (model: {self.model}, dim: {self._dimension})") + except httpx.HTTPError as e: + raise RuntimeError(f"Failed to connect to LiteLLM proxy at {self.api_base}: {e}") + + def encode(self, texts: list[str]) -> list[list[float]]: + """ + Generate embeddings using the LiteLLM proxy. + + Args: + texts: List of text strings to encode + + Returns: + List of embedding vectors + """ + if self._client is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + + if not texts: + return [] + + all_embeddings = [] + + # Process in batches + for i in range(0, len(texts), self.batch_size): + batch = texts[i : i + self.batch_size] + + response = self._client.post( + f"{self.api_base}/embeddings", + json={"model": self.model, "input": batch}, + ) + response.raise_for_status() + result = response.json() + + # Sort by index to ensure correct order + batch_embeddings = sorted(result["data"], key=lambda x: x["index"]) + all_embeddings.extend([e["embedding"] for e in batch_embeddings]) + + return all_embeddings + + def create_embeddings_from_env() -> Embeddings: """ Create an Embeddings instance based on environment variables. @@ -587,5 +709,12 @@ def create_embeddings_from_env() -> Embeddings: model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL) base_url = os.environ.get(ENV_EMBEDDINGS_COHERE_BASE_URL) or None return CohereEmbeddings(api_key=api_key, model=model, base_url=base_url) + elif provider == "litellm": + api_base = os.environ.get(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE) + api_key = os.environ.get(ENV_LITELLM_API_KEY) + model = os.environ.get(ENV_EMBEDDINGS_LITELLM_MODEL, DEFAULT_EMBEDDINGS_LITELLM_MODEL) + return LiteLLMEmbeddings(api_base=api_base, api_key=api_key, model=model) else: - raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere'") + raise ValueError( + f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere', 'litellm'" + ) diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 0347a5b4..26d6e215 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -139,7 +139,7 @@ export HINDSIGHT_API_REFLECT_LLM_MODEL=llama-3.3-70b-versatile | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_EMBEDDINGS_PROVIDER` | Provider: `local`, `tei`, `openai`, or `cohere` | `local` | +| `HINDSIGHT_API_EMBEDDINGS_PROVIDER` | Provider: `local`, `tei`, `openai`, `cohere`, or `litellm` | `local` | | `HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL` | Model for local provider | `BAAI/bge-small-en-v1.5` | | `HINDSIGHT_API_EMBEDDINGS_TEI_URL` | TEI server URL | - | | `HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY` | OpenAI API key (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | @@ -148,6 +148,9 @@ export HINDSIGHT_API_REFLECT_LLM_MODEL=llama-3.3-70b-versatile | `HINDSIGHT_API_COHERE_API_KEY` | Cohere API key (shared for embeddings and reranker) | - | | `HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL` | Cohere embedding model | `embed-english-v3.0` | | `HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | +| `HINDSIGHT_API_LITELLM_API_BASE` | LiteLLM proxy base URL (shared for embeddings and reranker) | `http://localhost:4000` | +| `HINDSIGHT_API_LITELLM_API_KEY` | LiteLLM proxy API key (optional, depends on proxy config) | - | +| `HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL` | LiteLLM embedding model (use provider prefix, e.g., `cohere/embed-english-v3.0`) | `text-embedding-3-small` | ```bash # Local (default) - uses SentenceTransformers @@ -179,6 +182,12 @@ export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere export HINDSIGHT_API_COHERE_API_KEY=your-azure-api-key export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0 export HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com + +# LiteLLM proxy - unified gateway for multiple providers +export HINDSIGHT_API_EMBEDDINGS_PROVIDER=litellm +export HINDSIGHT_API_LITELLM_API_BASE=http://localhost:4000 +export HINDSIGHT_API_LITELLM_API_KEY=your-litellm-key # optional +export HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL=text-embedding-3-small # or cohere/embed-english-v3.0 ``` #### Embedding Dimensions @@ -201,7 +210,7 @@ Supported OpenAI embedding dimensions: | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_RERANKER_PROVIDER` | Provider: `local`, `tei`, or `cohere` | `local` | +| `HINDSIGHT_API_RERANKER_PROVIDER` | Provider: `local`, `tei`, `cohere`, `flashrank`, `litellm`, 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_TEI_URL` | TEI server URL | - | @@ -209,6 +218,7 @@ Supported OpenAI embedding dimensions: | `HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT` | Max concurrent TEI reranking requests | `8` | | `HINDSIGHT_API_RERANKER_COHERE_MODEL` | Cohere rerank model | `rerank-english-v3.0` | | `HINDSIGHT_API_RERANKER_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | +| `HINDSIGHT_API_RERANKER_LITELLM_MODEL` | LiteLLM rerank model (use provider prefix, e.g., `cohere/rerank-english-v3.0`) | `cohere/rerank-english-v3.0` | ```bash # Local (default) - uses SentenceTransformers CrossEncoder @@ -229,8 +239,21 @@ export HINDSIGHT_API_RERANKER_PROVIDER=cohere export HINDSIGHT_API_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 + +# LiteLLM proxy - unified gateway for multiple reranking providers +export HINDSIGHT_API_RERANKER_PROVIDER=litellm +export HINDSIGHT_API_LITELLM_API_BASE=http://localhost:4000 +export HINDSIGHT_API_LITELLM_API_KEY=your-litellm-key # optional +export HINDSIGHT_API_RERANKER_LITELLM_MODEL=cohere/rerank-english-v3.0 # or voyage/rerank-2, together_ai/... ``` +LiteLLM supports multiple reranking providers via the `/rerank` endpoint: +- Cohere (`cohere/rerank-english-v3.0`, `cohere/rerank-multilingual-v3.0`) +- Together AI (`together_ai/...`) +- Voyage AI (`voyage/rerank-2`) +- Jina AI (`jina_ai/...`) +- AWS Bedrock (`bedrock/...`) + ### Authentication By default, Hindsight runs without authentication. For production deployments, enable API key authentication using the built-in tenant extension: