feat: support litellm gateway

* feat: support custom url for openai embeddings & cohere

* feat: support custom url for openai embeddings & cohere

hindsight-api/hindsight_api/config.py

@@ -41,10 +41,19 @@ ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
 ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
 ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
 ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
+ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
 
 ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
 ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
+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"
@@ -64,6 +73,7 @@ ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
 ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
 ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
 ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
+ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
 ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
 ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
 
@@ -120,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"
@@ -128,6 +143,7 @@ DEFAULT_MCP_ENABLED = True
 DEFAULT_GRAPH_RETRIEVER = "link_expansion"  # Options: "link_expansion", "mpfp", "bfs"
 DEFAULT_MPFP_TOP_K_NEIGHBORS = 20  # Fan-out limit per node in MPFP graph traversal
 DEFAULT_RECALL_MAX_CONCURRENT = 32  # Max concurrent recall operations per worker
+DEFAULT_RECALL_CONNECTION_BUDGET = 4  # Max concurrent DB connections per recall operation
 DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
 
 # Observation thresholds
@@ -222,6 +238,8 @@ class HindsightConfig:
     embeddings_provider: str
     embeddings_local_model: str
     embeddings_tei_url: str | None
+    embeddings_openai_base_url: str | None
+    embeddings_cohere_base_url: str | None
 
     # Reranker
     reranker_provider: str
@@ -230,6 +248,7 @@ class HindsightConfig:
     reranker_tei_batch_size: int
     reranker_tei_max_concurrent: int
     reranker_max_candidates: int
+    reranker_cohere_base_url: str | None
 
     # Server
     host: str
@@ -241,6 +260,7 @@ class HindsightConfig:
     graph_retriever: str
     mpfp_top_k_neighbors: int
     recall_max_concurrent: int
+    recall_connection_budget: int
 
     # Observation thresholds
     observation_min_facts: int
@@ -297,6 +317,8 @@ class HindsightConfig:
             embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
             embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
             embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
+            embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None,
+            embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None,
             # Reranker
             reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
             reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
@@ -306,6 +328,7 @@ class HindsightConfig:
                 os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT))
             ),
             reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
+            reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None,
             # Server
             host=os.getenv(ENV_HOST, DEFAULT_HOST),
             port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
@@ -315,6 +338,9 @@ class HindsightConfig:
             graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
             mpfp_top_k_neighbors=int(os.getenv(ENV_MPFP_TOP_K_NEIGHBORS, str(DEFAULT_MPFP_TOP_K_NEIGHBORS))),
             recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))),
+            recall_connection_budget=int(
+                os.getenv(ENV_RECALL_CONNECTION_BUDGET, str(DEFAULT_RECALL_CONNECTION_BUDGET))
+            ),
             # Optimization flags
             skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
             lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",

hindsight-api/hindsight_api/engine/cross_encoder.py

@@ -15,18 +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,
@@ -392,6 +398,7 @@ class CohereCrossEncoder(CrossEncoderModel):
         self,
         api_key: str,
         model: str = DEFAULT_RERANKER_COHERE_MODEL,
+        base_url: str | None = None,
         timeout: float = 60.0,
     ):
         """
@@ -400,10 +407,12 @@ class CohereCrossEncoder(CrossEncoderModel):
         Args:
             api_key: Cohere API key
             model: Cohere rerank model name (default: rerank-english-v3.0)
+            base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint)
             timeout: Request timeout in seconds (default: 60.0)
         """
         self.api_key = api_key
         self.model = model
+        self.base_url = base_url
         self.timeout = timeout
         self._client = None
 
@@ -421,8 +430,14 @@ class CohereCrossEncoder(CrossEncoderModel):
         except ImportError:
             raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
 
-        logger.info(f"Reranker: initializing Cohere provider with model {self.model}")
-        self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
+        base_url_msg = f" at {self.base_url}" if self.base_url else ""
+        logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}")
+
+        # Build client kwargs, only including base_url if set (for Azure or custom endpoints)
+        client_kwargs = {"api_key": self.api_key, "timeout": self.timeout}
+        if self.base_url:
+            client_kwargs["base_url"] = self.base_url
+        self._client = cohere.Client(**client_kwargs)
         logger.info("Reranker: Cohere provider initialized")
 
     async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
@@ -641,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.
@@ -671,14 +796,20 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
         if not api_key:
             raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
         model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL)
-        return CohereCrossEncoder(api_key=api_key, model=model)
+        base_url = os.environ.get(ENV_RERANKER_COHERE_BASE_URL) or None
+        return CohereCrossEncoder(api_key=api_key, model=model, base_url=base_url)
     elif provider == "flashrank":
         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'"
         )

hindsight-api/hindsight_api/engine/db_budget.py

@@ -0,0 +1,284 @@
+"""
+Database connection budget management.
+
+Limits concurrent database connections per operation to prevent
+a single operation (e.g., recall with parallel queries) from
+exhausting the connection pool.
+"""
+
+import asyncio
+import logging
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, AsyncIterator
+
+if TYPE_CHECKING:
+    import asyncpg
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OperationBudget:
+    """
+    Tracks connection budget for a single operation.
+
+    Each operation gets a semaphore limiting its concurrent connections.
+    """
+
+    operation_id: str
+    max_connections: int
+    semaphore: asyncio.Semaphore = field(init=False)
+    active_count: int = field(default=0, init=False)
+
+    def __post_init__(self):
+        self.semaphore = asyncio.Semaphore(self.max_connections)
+
+
+class ConnectionBudgetManager:
+    """
+    Manages per-operation connection budgets.
+
+    Usage:
+        manager = ConnectionBudgetManager(default_budget=4)
+
+        # Start an operation
+        async with manager.operation(max_connections=2) as op:
+            # Acquire connections within the budget
+            async with op.acquire(pool) as conn:
+                await conn.fetch(...)
+
+            # Multiple connections respect the budget
+            async with op.acquire(pool) as conn1, op.acquire(pool) as conn2:
+                # At most 2 concurrent connections for this operation
+                ...
+    """
+
+    def __init__(self, default_budget: int = 4):
+        """
+        Initialize the budget manager.
+
+        Args:
+            default_budget: Default max connections per operation
+        """
+        self.default_budget = default_budget
+        self._operations: dict[str, OperationBudget] = {}
+        self._lock = asyncio.Lock()
+
+    @asynccontextmanager
+    async def operation(
+        self,
+        max_connections: int | None = None,
+        operation_id: str | None = None,
+    ) -> AsyncIterator["BudgetedOperation"]:
+        """
+        Create a budgeted operation context.
+
+        Args:
+            max_connections: Max concurrent connections for this operation.
+                           Defaults to manager's default_budget.
+            operation_id: Optional custom operation ID. Auto-generated if not provided.
+
+        Yields:
+            BudgetedOperation context for acquiring connections
+        """
+        op_id = operation_id or f"op-{uuid.uuid4().hex[:12]}"
+        budget = max_connections or self.default_budget
+
+        async with self._lock:
+            if op_id in self._operations:
+                raise ValueError(f"Operation {op_id} already exists")
+            self._operations[op_id] = OperationBudget(op_id, budget)
+
+        try:
+            yield BudgetedOperation(self, op_id)
+        finally:
+            async with self._lock:
+                self._operations.pop(op_id, None)
+
+    def _get_budget(self, operation_id: str) -> OperationBudget:
+        """Get budget for an operation (internal use)."""
+        budget = self._operations.get(operation_id)
+        if not budget:
+            raise ValueError(f"Operation {operation_id} not found")
+        return budget
+
+
+class BudgetedOperation:
+    """
+    A single operation with connection budget.
+
+    Provides methods to acquire connections within the budget.
+    """
+
+    def __init__(self, manager: ConnectionBudgetManager, operation_id: str):
+        self._manager = manager
+        self.operation_id = operation_id
+
+    @property
+    def budget(self) -> OperationBudget:
+        """Get the budget for this operation."""
+        return self._manager._get_budget(self.operation_id)
+
+    @asynccontextmanager
+    async def acquire(self, pool: "asyncpg.Pool") -> AsyncIterator["asyncpg.Connection"]:
+        """
+        Acquire a connection within the operation's budget.
+
+        Blocks if the operation has reached its connection limit.
+
+        Args:
+            pool: asyncpg connection pool
+
+        Yields:
+            Database connection
+        """
+        budget = self.budget
+        async with budget.semaphore:
+            budget.active_count += 1
+            conn = await pool.acquire()
+            try:
+                yield conn
+            finally:
+                budget.active_count -= 1
+                await pool.release(conn)
+
+    def wrap_pool(self, pool: "asyncpg.Pool") -> "BudgetedPool":
+        """
+        Wrap a pool with this operation's budget.
+
+        The returned BudgetedPool can be passed to functions expecting a pool,
+        and all acquire() calls will be limited by this operation's budget.
+
+        Args:
+            pool: asyncpg connection pool to wrap
+
+        Returns:
+            BudgetedPool that limits connections to this operation's budget
+        """
+        return BudgetedPool(pool, self)
+
+    async def acquire_many(
+        self,
+        pool: "asyncpg.Pool",
+        count: int,
+    ) -> AsyncIterator[list["asyncpg.Connection"]]:
+        """
+        Acquire multiple connections within the budget.
+
+        Note: This acquires connections sequentially to respect the budget.
+        For parallel acquisition, use multiple acquire() calls with asyncio.gather().
+
+        Args:
+            pool: asyncpg connection pool
+            count: Number of connections to acquire
+
+        Yields:
+            List of database connections
+        """
+        connections = []
+        try:
+            for _ in range(count):
+                conn = await pool.acquire()
+                connections.append(conn)
+            yield connections
+        finally:
+            for conn in connections:
+                await pool.release(conn)
+
+
+# Global default manager instance
+_default_manager: ConnectionBudgetManager | None = None
+
+
+def get_budget_manager(default_budget: int = 4) -> ConnectionBudgetManager:
+    """
+    Get or create the global budget manager.
+
+    Args:
+        default_budget: Default max connections per operation
+
+    Returns:
+        Global ConnectionBudgetManager instance
+    """
+    global _default_manager
+    if _default_manager is None:
+        _default_manager = ConnectionBudgetManager(default_budget=default_budget)
+    return _default_manager
+
+
+@asynccontextmanager
+async def budgeted_operation(
+    max_connections: int | None = None,
+    operation_id: str | None = None,
+    default_budget: int = 4,
+) -> AsyncIterator[BudgetedOperation]:
+    """
+    Convenience function to create a budgeted operation.
+
+    Args:
+        max_connections: Max concurrent connections for this operation
+        operation_id: Optional custom operation ID
+        default_budget: Default budget if manager not yet created
+
+    Yields:
+        BudgetedOperation context
+
+    Example:
+        async with budgeted_operation(max_connections=2) as op:
+            async with op.acquire(pool) as conn:
+                await conn.fetch(...)
+    """
+    manager = get_budget_manager(default_budget)
+    async with manager.operation(max_connections, operation_id) as op:
+        yield op
+
+
+class BudgetedPool:
+    """
+    A pool wrapper that limits concurrent connection acquisitions.
+
+    This can be passed to functions expecting a pool, and acquire()
+    calls will be limited by the budget semaphore.
+
+    Usage:
+        async with budgeted_operation(max_connections=4) as op:
+            budgeted_pool = op.wrap_pool(pool)
+            # Pass budgeted_pool to functions that expect a pool
+            await some_function(budgeted_pool, ...)
+    """
+
+    def __init__(self, pool: "asyncpg.Pool", operation: BudgetedOperation):
+        self._pool = pool
+        self._operation = operation
+
+    async def acquire(self) -> "asyncpg.Connection":
+        """
+        Acquire a connection within the budget.
+
+        Note: Caller must release the connection when done.
+        Prefer using as context manager via acquire_with_retry or op.acquire().
+        """
+        budget = self._operation.budget
+        await budget.semaphore.acquire()
+        budget.active_count += 1
+        try:
+            return await self._pool.acquire()
+        except Exception:
+            budget.active_count -= 1
+            budget.semaphore.release()
+            raise
+
+    async def release(self, conn: "asyncpg.Connection") -> None:
+        """Release a connection back to the pool."""
+        budget = self._operation.budget
+        try:
+            await self._pool.release(conn)
+        finally:
+            budget.active_count -= 1
+            budget.semaphore.release()
+
+    def __getattr__(self, name):
+        """Proxy other attributes to the underlying pool."""
+        return getattr(self._pool, name)

hindsight-api/hindsight_api/engine/embeddings.py

@@ -17,16 +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,
 )
 
@@ -322,6 +329,7 @@ class OpenAIEmbeddings(Embeddings):
         self,
         api_key: str,
         model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
+        base_url: str | None = None,
         batch_size: int = 100,
         max_retries: int = 3,
     ):
@@ -331,11 +339,13 @@ class OpenAIEmbeddings(Embeddings):
         Args:
             api_key: OpenAI API key
             model: OpenAI embedding model name (default: text-embedding-3-small)
+            base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
             batch_size: Maximum batch size for embedding requests (default: 100)
             max_retries: Maximum number of retries for failed requests (default: 3)
         """
         self.api_key = api_key
         self.model = model
+        self.base_url = base_url
         self.batch_size = batch_size
         self.max_retries = max_retries
         self._client = None
@@ -361,8 +371,14 @@ class OpenAIEmbeddings(Embeddings):
         except ImportError:
             raise ImportError("openai is required for OpenAIEmbeddings. Install it with: pip install openai")
 
-        logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}")
-        self._client = OpenAI(api_key=self.api_key, max_retries=self.max_retries)
+        base_url_msg = f" at {self.base_url}" if self.base_url else ""
+        logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}{base_url_msg}")
+
+        # Build client kwargs, only including base_url if set (for Azure or custom endpoints)
+        client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries}
+        if self.base_url:
+            client_kwargs["base_url"] = self.base_url
+        self._client = OpenAI(**client_kwargs)
 
         # Try to get dimension from known models, otherwise do a test embedding
         if self.model in self.MODEL_DIMENSIONS:
@@ -435,6 +451,7 @@ class CohereEmbeddings(Embeddings):
         self,
         api_key: str,
         model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL,
+        base_url: str | None = None,
         batch_size: int = 96,
         timeout: float = 60.0,
         input_type: str = "search_document",
@@ -445,6 +462,7 @@ class CohereEmbeddings(Embeddings):
         Args:
             api_key: Cohere API key
             model: Cohere embedding model name (default: embed-english-v3.0)
+            base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint)
             batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit)
             timeout: Request timeout in seconds (default: 60.0)
             input_type: Input type for embeddings (default: search_document).
@@ -452,6 +470,7 @@ class CohereEmbeddings(Embeddings):
         """
         self.api_key = api_key
         self.model = model
+        self.base_url = base_url
         self.batch_size = batch_size
         self.timeout = timeout
         self.input_type = input_type
@@ -478,8 +497,14 @@ class CohereEmbeddings(Embeddings):
         except ImportError:
             raise ImportError("cohere is required for CohereEmbeddings. Install it with: pip install cohere")
 
-        logger.info(f"Embeddings: initializing Cohere provider with model {self.model}")
-        self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
+        base_url_msg = f" at {self.base_url}" if self.base_url else ""
+        logger.info(f"Embeddings: initializing Cohere provider with model {self.model}{base_url_msg}")
+
+        # Build client kwargs, only including base_url if set (for Azure or custom endpoints)
+        client_kwargs = {"api_key": self.api_key, "timeout": self.timeout}
+        if self.base_url:
+            client_kwargs["base_url"] = self.base_url
+        self._client = cohere.Client(**client_kwargs)
 
         # Try to get dimension from known models, otherwise do a test embedding
         if self.model in self.MODEL_DIMENSIONS:
@@ -529,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.
@@ -558,12 +700,21 @@ def create_embeddings_from_env() -> Embeddings:
                 f"when {ENV_EMBEDDINGS_PROVIDER} is 'openai'"
             )
         model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
-        return OpenAIEmbeddings(api_key=api_key, model=model)
+        base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None
+        return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url)
     elif provider == "cohere":
         api_key = os.environ.get(ENV_COHERE_API_KEY)
         if not api_key:
             raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'")
         model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL)
-        return CohereEmbeddings(api_key=api_key, model=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'"
+        )

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
 
 from ..config import get_config
 from ..metrics import get_metrics_collector
+from .db_budget import budgeted_operation
 
 # Context variable for current schema (async-safe, per-task isolation)
 _current_schema: contextvars.ContextVar[str] = contextvars.ContextVar("current_schema", default="public")
@@ -1609,38 +1610,40 @@ class MemoryEngine(MemoryEngineInterface):
                 tracer.record_query_embedding(query_embedding)
                 tracer.add_phase_metric("generate_query_embedding", step_duration)
 
-            # Step 2: N*4-Way Parallel Retrieval (N fact types × 4 retrieval methods)
+            # Step 2: Optimized parallel retrieval using batched queries
+            # - Semantic + BM25 combined in 1 CTE query for ALL fact types
+            # - Graph runs per fact type (complex traversal)
+            # - Temporal runs per fact type (if constraint detected)
             step_start = time.time()
             query_embedding_str = str(query_embedding)
 
-            from .search.retrieval import get_default_graph_retriever, retrieve_parallel
+            from .search.retrieval import (
+                get_default_graph_retriever,
+                retrieve_all_fact_types_parallel,
+            )
 
             # Track each retrieval start time
             retrieval_start = time.time()
 
-            # Temporal extraction now runs IN PARALLEL with other retrievals inside retrieve_parallel
-            # This prevents slow dateparser from blocking semantic/BM25/graph retrieval
-            tc_duration = 0.0  # Will be tracked inside temporal retrieval timing
-
-            # Run retrieval for each fact type in parallel
-            # Each retrieve_parallel uses ~4 connections, so 3 fact types = ~12 concurrent connections
-            retrieval_tasks = [
-                retrieve_parallel(
-                    pool,
+            # Run optimized retrieval with connection budget
+            config = get_config()
+            async with budgeted_operation(
+                max_connections=config.recall_connection_budget,
+                operation_id=f"recall-{recall_id}",
+            ) as op:
+                budgeted_pool = op.wrap_pool(pool)
+                parallel_start = time.time()
+                multi_result = await retrieve_all_fact_types_parallel(
+                    budgeted_pool,
                     query,
                     query_embedding_str,
                     bank_id,
-                    ft,
+                    fact_type,  # Pass all fact types at once
                     thinking_budget,
                     question_date,
                     self.query_analyzer,
-                    temporal_constraint=None,  # Extracted in parallel inside retrieve_parallel
                 )
-                for ft in fact_type
-            ]
-            parallel_start = time.time()
-            all_retrievals = await asyncio.gather(*retrieval_tasks)
-            parallel_duration = time.time() - parallel_start
+                parallel_duration = time.time() - parallel_start
 
             # Combine all results from all fact types and aggregate timings
             semantic_results = []
@@ -1657,12 +1660,15 @@ class MemoryEngine(MemoryEngineInterface):
             all_mpfp_timings = []
 
             detected_temporal_constraint = None
-            max_conn_wait = 0.0
-            for idx, retrieval_result in enumerate(all_retrievals):
+            max_conn_wait = multi_result.max_conn_wait
+            for ft in fact_type:
+                retrieval_result = multi_result.results_by_fact_type.get(ft)
+                if not retrieval_result:
+                    continue
+
                 # Log fact types in this retrieval batch
-                ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
                 logger.debug(
-                    f"[RECALL {recall_id}] Fact type '{ft_name}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}"
+                    f"[RECALL {recall_id}] Fact type '{ft}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}"
                 )
 
                 semantic_results.extend(retrieval_result.semantic)
@@ -1678,8 +1684,6 @@ class MemoryEngine(MemoryEngineInterface):
                     detected_temporal_constraint = retrieval_result.temporal_constraint
                 # Collect MPFP timings
                 all_mpfp_timings.extend(retrieval_result.mpfp_timings)
-                # Track max connection wait
-                max_conn_wait = max(max_conn_wait, retrieval_result.max_conn_wait)
 
             # If no temporal results from any fact type, set to None
             if not temporal_results:
@@ -1711,10 +1715,8 @@ class MemoryEngine(MemoryEngineInterface):
                 temporal_count = len(temporal_results) if temporal_results else 0
                 timing_parts.append(f"temporal={temporal_count}({aggregated_timings['temporal']:.3f}s)")
                 temporal_info = f" | temporal_range={start_dt.strftime('%Y-%m-%d')} to {end_dt.strftime('%Y-%m-%d')}"
-            # Only tc is sequential setup now (adjacency loads in parallel with retrieval)
-            setup_info = f", tc={tc_duration:.3f}s" if tc_duration > 0.01 else ""
             log_buffer.append(
-                f"  [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{setup_info}{temporal_info}"
+                f"  [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{temporal_info}"
             )
 
             # Log graph retriever timing breakdown if available
@@ -1751,8 +1753,10 @@ class MemoryEngine(MemoryEngineInterface):
                     return [(r.id, r.__dict__) for r in results]
 
                 # Add retrieval results per fact type (to show parallel execution in UI)
-                for idx, rr in enumerate(all_retrievals):
-                    ft_name = fact_type[idx] if idx < len(fact_type) else "unknown"
+                for ft_name in fact_type:
+                    rr = multi_result.results_by_fact_type.get(ft_name)
+                    if not rr:
+                        continue
 
                     # Add semantic retrieval results for this fact type
                     tracer.add_retrieval_results(

hindsight-api/hindsight_api/engine/search/retrieval.py

@@ -39,6 +39,18 @@ class ParallelRetrievalResult:
     max_conn_wait: float = 0.0  # Maximum connection acquisition wait time across all methods
 
 
+@dataclass
+class MultiFactTypeRetrievalResult:
+    """Result from retrieval across all fact types."""
+
+    # Results per fact type
+    results_by_fact_type: dict[str, ParallelRetrievalResult]
+    # Aggregate timings
+    timings: dict[str, float] = field(default_factory=dict)
+    # Max connection wait across all operations
+    max_conn_wait: float = 0.0
+
+
 # Default graph retriever instance (can be overridden)
 _default_graph_retriever: GraphRetriever | None = None
 
@@ -158,6 +170,356 @@ async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, lim
     return [RetrievalResult.from_db_row(dict(r)) for r in results]
 
 
+async def retrieve_semantic_bm25_combined(
+    conn,
+    query_emb_str: str,
+    query_text: str,
+    bank_id: str,
+    fact_types: list[str],
+    limit: int,
+) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
+    """
+    Combined semantic + BM25 retrieval for multiple fact types in a single query.
+
+    Uses CTEs with window functions to get top-N results per fact type per method,
+    all in one database round-trip.
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        query_text: Query text for BM25
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        limit: Maximum results per method per fact type
+
+    Returns:
+        Dict mapping fact_type -> (semantic_results, bm25_results)
+    """
+    import re
+
+    # Sanitize query text for BM25 (same as retrieve_bm25)
+    sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
+    tokens = [token for token in sanitized_text.split() if token]
+
+    # If no valid tokens for BM25, just run semantic
+    if not tokens:
+        results = await conn.fetch(
+            f"""
+            WITH semantic_ranked AS (
+                SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                       1 - (embedding <=> $1::vector) AS similarity,
+                       NULL::float AS bm25_score,
+                       'semantic' AS source,
+                       ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
+                FROM {fq_table("memory_units")}
+                WHERE bank_id = $2
+                  AND embedding IS NOT NULL
+                  AND fact_type = ANY($3)
+                  AND (1 - (embedding <=> $1::vector)) >= 0.3
+            )
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM semantic_ranked
+            WHERE rn <= $4
+            """,
+            query_emb_str,
+            bank_id,
+            fact_types,
+            limit,
+        )
+        # Group by fact_type
+        result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {
+            ft: ([], []) for ft in fact_types
+        }
+        for r in results:
+            row = dict(r)
+            ft = row.get("fact_type")
+            row.pop("source", None)
+            if ft in result_dict:
+                result_dict[ft][0].append(RetrievalResult.from_db_row(row))
+        return result_dict
+
+    query_tsquery = " | ".join(tokens)
+
+    # Combined CTE query for both semantic and BM25 across all fact types
+    # Uses window functions to limit per fact_type per method
+    results = await conn.fetch(
+        f"""
+        WITH semantic_ranked AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   1 - (embedding <=> $1::vector) AS similarity,
+                   NULL::float AS bm25_score,
+                   'semantic' AS source,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND embedding IS NOT NULL
+              AND fact_type = ANY($3)
+              AND (1 - (embedding <=> $1::vector)) >= 0.3
+        ),
+        bm25_ranked AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   NULL::float AS similarity,
+                   ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score,
+                   'bm25' AS source,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND fact_type = ANY($3)
+              AND search_vector @@ to_tsquery('english', $5)
+        ),
+        semantic AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM semantic_ranked WHERE rn <= $4
+        ),
+        bm25 AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   similarity, bm25_score, source
+            FROM bm25_ranked WHERE rn <= $4
+        )
+        SELECT * FROM semantic
+        UNION ALL
+        SELECT * FROM bm25
+        """,
+        query_emb_str,
+        bank_id,
+        fact_types,
+        limit,
+        query_tsquery,
+    )
+
+    # Group results by fact_type and source
+    result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
+    for r in results:
+        row = dict(r)
+        source = row.pop("source", None)
+        ft = row.get("fact_type")
+        if ft in result_dict:
+            if source == "semantic":
+                result_dict[ft][0].append(RetrievalResult.from_db_row(row))
+            else:
+                result_dict[ft][1].append(RetrievalResult.from_db_row(row))
+
+    return result_dict
+
+
+async def retrieve_temporal_combined(
+    conn,
+    query_emb_str: str,
+    bank_id: str,
+    fact_types: list[str],
+    start_date: datetime,
+    end_date: datetime,
+    budget: int,
+    semantic_threshold: float = 0.1,
+) -> dict[str, list[RetrievalResult]]:
+    """
+    Temporal retrieval for multiple fact types in a single query.
+
+    Batches the entry point query using window functions to get top-N per fact type,
+    then runs spreading for each fact type.
+
+    Args:
+        conn: Database connection
+        query_emb_str: Query embedding as string
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        start_date: Start of time range
+        end_date: End of time range
+        budget: Node budget for spreading per fact type
+        semantic_threshold: Minimum semantic similarity to include
+
+    Returns:
+        Dict mapping fact_type -> list of RetrievalResult
+    """
+    from ..memory_engine import fq_table
+
+    # Ensure dates are timezone-aware
+    if start_date.tzinfo is None:
+        start_date = start_date.replace(tzinfo=UTC)
+    if end_date.tzinfo is None:
+        end_date = end_date.replace(tzinfo=UTC)
+
+    # Batch query: Get entry points for ALL fact types at once with window function
+    entry_points = await conn.fetch(
+        f"""
+        WITH ranked_entries AS (
+            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
+                   1 - (embedding <=> $1::vector) AS similarity,
+                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $2
+              AND fact_type = ANY($3)
+              AND embedding IS NOT NULL
+              AND (
+                  (occurred_start IS NOT NULL AND occurred_end IS NOT NULL
+                   AND occurred_start <= $5 AND occurred_end >= $4)
+                  OR
+                  (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
+                  OR
+                  (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
+                  OR
+                  (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
+              )
+              AND (1 - (embedding <=> $1::vector)) >= $6
+        )
+        SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id, similarity
+        FROM ranked_entries
+        WHERE rn <= 10
+        """,
+        query_emb_str,
+        bank_id,
+        fact_types,
+        start_date,
+        end_date,
+        semantic_threshold,
+    )
+
+    if not entry_points:
+        return {ft: [] for ft in fact_types}
+
+    # Group entry points by fact type
+    entries_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
+    for ep in entry_points:
+        ft = ep["fact_type"]
+        if ft in entries_by_ft:
+            entries_by_ft[ft].append(ep)
+
+    # Calculate shared temporal parameters
+    total_days = (end_date - start_date).total_seconds() / 86400
+    mid_date = start_date + (end_date - start_date) / 2
+
+    # Process each fact type (spreading needs to stay per fact type due to link filtering)
+    results_by_ft: dict[str, list[RetrievalResult]] = {}
+
+    for ft in fact_types:
+        ft_entry_points = entries_by_ft.get(ft, [])
+        if not ft_entry_points:
+            results_by_ft[ft] = []
+            continue
+
+        results = []
+        visited = set()
+        node_scores = {}
+
+        # Process entry points
+        for ep in ft_entry_points:
+            unit_id = str(ep["id"])
+            visited.add(unit_id)
+
+            # Calculate temporal proximity
+            best_date = None
+            if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
+                best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
+            elif ep["occurred_start"] is not None:
+                best_date = ep["occurred_start"]
+            elif ep["occurred_end"] is not None:
+                best_date = ep["occurred_end"]
+            elif ep["mentioned_at"] is not None:
+                best_date = ep["mentioned_at"]
+
+            if best_date:
+                days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
+                temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
+            else:
+                temporal_proximity = 0.5
+
+            ep_result = RetrievalResult.from_db_row(dict(ep))
+            ep_result.temporal_score = temporal_proximity
+            ep_result.temporal_proximity = temporal_proximity
+            results.append(ep_result)
+            node_scores[unit_id] = (ep["similarity"], 1.0)
+
+        # Spreading through temporal links (same as single-fact-type version)
+        frontier = list(node_scores.keys())
+        budget_remaining = budget - len(ft_entry_points)
+        batch_size = 20
+
+        while frontier and budget_remaining > 0:
+            batch_ids = frontier[:batch_size]
+            frontier = frontier[batch_size:]
+
+            neighbors = await conn.fetch(
+                f"""
+                SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
+                       ml.weight, ml.link_type, ml.from_unit_id,
+                       1 - (mu.embedding <=> $1::vector) AS similarity
+                FROM {fq_table("memory_links")} ml
+                JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
+                WHERE ml.from_unit_id = ANY($2::uuid[])
+                  AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
+                  AND ml.weight >= 0.1
+                  AND mu.fact_type = $3
+                  AND mu.embedding IS NOT NULL
+                  AND (1 - (mu.embedding <=> $1::vector)) >= $4
+                ORDER BY ml.weight DESC
+                LIMIT $5
+                """,
+                query_emb_str,
+                batch_ids,
+                ft,
+                semantic_threshold,
+                batch_size * 10,
+            )
+
+            for n in neighbors:
+                neighbor_id = str(n["id"])
+                if neighbor_id in visited:
+                    continue
+
+                visited.add(neighbor_id)
+                budget_remaining -= 1
+
+                parent_id = str(n["from_unit_id"])
+                _, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
+
+                neighbor_best_date = None
+                if n["occurred_start"] is not None and n["occurred_end"] is not None:
+                    neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
+                elif n["occurred_start"] is not None:
+                    neighbor_best_date = n["occurred_start"]
+                elif n["occurred_end"] is not None:
+                    neighbor_best_date = n["occurred_end"]
+                elif n["mentioned_at"] is not None:
+                    neighbor_best_date = n["mentioned_at"]
+
+                if neighbor_best_date:
+                    days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
+                    neighbor_temporal_proximity = (
+                        1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
+                    )
+                else:
+                    neighbor_temporal_proximity = 0.3
+
+                link_type = n["link_type"]
+                if link_type in ("causes", "caused_by"):
+                    causal_boost = 2.0
+                elif link_type in ("enables", "prevents"):
+                    causal_boost = 1.5
+                else:
+                    causal_boost = 1.0
+
+                propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
+                combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
+
+                neighbor_result = RetrievalResult.from_db_row(dict(n))
+                neighbor_result.temporal_score = combined_temporal
+                neighbor_result.temporal_proximity = neighbor_temporal_proximity
+                results.append(neighbor_result)
+
+                if budget_remaining > 0 and combined_temporal > 0.2:
+                    node_scores[neighbor_id] = (n["similarity"], combined_temporal)
+                    frontier.append(neighbor_id)
+
+                if budget_remaining <= 0:
+                    break
+
+        results_by_ft[ft] = results
+
+    return results_by_ft
+
+
 async def retrieve_temporal(
     conn,
     query_emb_str: str,
@@ -748,3 +1110,158 @@ async def _retrieve_parallel_bfs(
             },
             temporal_constraint=None,
         )
+
+
+async def retrieve_all_fact_types_parallel(
+    pool,
+    query_text: str,
+    query_embedding_str: str,
+    bank_id: str,
+    fact_types: list[str],
+    thinking_budget: int,
+    question_date: datetime | None = None,
+    query_analyzer: Optional["QueryAnalyzer"] = None,
+    graph_retriever: GraphRetriever | None = None,
+) -> MultiFactTypeRetrievalResult:
+    """
+    Optimized retrieval for multiple fact types using batched queries.
+
+    This reduces database round-trips by:
+    1. Combining semantic + BM25 into one CTE query for ALL fact types (1 query instead of 2N)
+    2. Running graph retrieval per fact type in parallel (N parallel tasks)
+    3. Running temporal retrieval per fact type in parallel (N parallel tasks)
+
+    Args:
+        pool: Database connection pool
+        query_text: Query text
+        query_embedding_str: Query embedding as string
+        bank_id: Bank ID
+        fact_types: List of fact types to retrieve
+        thinking_budget: Budget for graph traversal and retrieval limits
+        question_date: Optional date when question was asked (for temporal filtering)
+        query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
+        graph_retriever: Graph retrieval strategy (defaults to configured retriever)
+
+    Returns:
+        MultiFactTypeRetrievalResult with results organized by fact type
+    """
+    import time
+
+    retriever = graph_retriever or get_default_graph_retriever()
+    start_time = time.time()
+    timings: dict[str, float] = {}
+
+    # Step 1: Extract temporal constraint first (CPU work, no DB)
+    # Do this before DB queries so we know if we need temporal retrieval
+    temporal_extraction_start = time.time()
+    from .temporal_extraction import extract_temporal_constraint
+
+    temporal_constraint = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
+    temporal_extraction_time = time.time() - temporal_extraction_start
+    timings["temporal_extraction"] = temporal_extraction_time
+
+    # Step 2: Run semantic + BM25 + temporal combined in ONE connection!
+    # This reduces connection usage from 2 to 1 for these operations
+    semantic_bm25_start = time.time()
+    temporal_results_by_ft: dict[str, list[RetrievalResult]] = {}
+    temporal_time = 0.0
+
+    async with acquire_with_retry(pool) as conn:
+        conn_wait = time.time() - semantic_bm25_start
+
+        # Semantic + BM25 combined
+        semantic_bm25_results = await retrieve_semantic_bm25_combined(
+            conn, query_embedding_str, query_text, bank_id, fact_types, thinking_budget
+        )
+        semantic_bm25_time = time.time() - semantic_bm25_start
+
+        # Temporal combined (if constraint detected) - same connection!
+        if temporal_constraint:
+            tc_start, tc_end = temporal_constraint
+            temporal_start = time.time()
+            temporal_results_by_ft = await retrieve_temporal_combined(
+                conn,
+                query_embedding_str,
+                bank_id,
+                fact_types,
+                tc_start,
+                tc_end,
+                budget=thinking_budget,
+                semantic_threshold=0.1,
+            )
+            temporal_time = time.time() - temporal_start
+
+    timings["semantic_bm25_combined"] = semantic_bm25_time
+    timings["temporal_combined"] = temporal_time
+
+    # Step 3: Run graph retrieval for each fact type in parallel
+    async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
+        graph_start = time.time()
+        results, mpfp_timing = await retriever.retrieve(
+            pool=pool,
+            query_embedding_str=query_embedding_str,
+            bank_id=bank_id,
+            fact_type=ft,
+            budget=thinking_budget,
+            query_text=query_text,
+            semantic_seeds=None,
+            temporal_seeds=None,
+        )
+        return ft, results, time.time() - graph_start, mpfp_timing
+
+    # Run graph for all fact types in parallel
+    graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
+    graph_results_list = await asyncio.gather(*graph_tasks)
+
+    # Organize results by fact type
+    results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
+    max_conn_wait = conn_wait  # Single connection for semantic+bm25+temporal
+    all_mpfp_timings: list[MPFPTimings] = []
+
+    for ft in fact_types:
+        # Get semantic + bm25 results for this fact type
+        semantic_results, bm25_results = semantic_bm25_results.get(ft, ([], []))
+
+        # Find graph results for this fact type
+        graph_results = []
+        graph_time = 0.0
+        mpfp_timing = None
+        for gr in graph_results_list:
+            if gr[0] == ft:
+                graph_results = gr[1]
+                graph_time = gr[2]
+                mpfp_timing = gr[3]
+                if mpfp_timing:
+                    all_mpfp_timings.append(mpfp_timing)
+                break
+
+        # Get temporal results for this fact type from combined result
+        temporal_results = temporal_results_by_ft.get(ft) if temporal_constraint else None
+        if temporal_results is not None and len(temporal_results) == 0:
+            temporal_results = None
+
+        results_by_fact_type[ft] = ParallelRetrievalResult(
+            semantic=semantic_results,
+            bm25=bm25_results,
+            graph=graph_results,
+            temporal=temporal_results,
+            timings={
+                "semantic": semantic_bm25_time / 2,  # Approximate split
+                "bm25": semantic_bm25_time / 2,
+                "graph": graph_time,
+                "temporal": temporal_time,  # Same for all fact types (single query)
+                "temporal_extraction": temporal_extraction_time,
+            },
+            temporal_constraint=temporal_constraint,
+            mpfp_timings=[mpfp_timing] if mpfp_timing else [],
+            max_conn_wait=max_conn_wait,
+        )
+
+    total_time = time.time() - start_time
+    timings["total"] = total_time
+
+    return MultiFactTypeRetrievalResult(
+        results_by_fact_type=results_by_fact_type,
+        timings=timings,
+        max_conn_wait=max_conn_wait,
+    )

hindsight-api/hindsight_api/main.py

@@ -187,12 +187,15 @@ def main():
             embeddings_provider=config.embeddings_provider,
             embeddings_local_model=config.embeddings_local_model,
             embeddings_tei_url=config.embeddings_tei_url,
+            embeddings_openai_base_url=config.embeddings_openai_base_url,
+            embeddings_cohere_base_url=config.embeddings_cohere_base_url,
             reranker_provider=config.reranker_provider,
             reranker_local_model=config.reranker_local_model,
             reranker_tei_url=config.reranker_tei_url,
             reranker_tei_batch_size=config.reranker_tei_batch_size,
             reranker_tei_max_concurrent=config.reranker_tei_max_concurrent,
             reranker_max_candidates=config.reranker_max_candidates,
+            reranker_cohere_base_url=config.reranker_cohere_base_url,
             host=args.host,
             port=args.port,
             log_level=args.log_level,
@@ -200,6 +203,7 @@ def main():
             graph_retriever=config.graph_retriever,
             mpfp_top_k_neighbors=config.mpfp_top_k_neighbors,
             recall_max_concurrent=config.recall_max_concurrent,
+            recall_connection_budget=config.recall_connection_budget,
             observation_min_facts=config.observation_min_facts,
             observation_top_entities=config.observation_top_entities,
             retain_max_completion_tokens=config.retain_max_completion_tokens,

hindsight-api/hindsight_api/metrics.py

@@ -28,6 +28,15 @@ from opentelemetry.sdk.resources import Resource
 if TYPE_CHECKING:
     import asyncpg
 
+
+def _get_tenant() -> str:
+    """Get current tenant (schema) from context for metrics labeling."""
+    # Import here to avoid circular imports
+    from hindsight_api.engine.memory_engine import get_current_schema
+
+    return get_current_schema()
+
+
 # Custom bucket boundaries for operation duration (in seconds)
 # Fine granularity in 0-30s range where most operations complete
 DURATION_BUCKETS = (0.1, 0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 5.0, 7.5, 10.0, 15.0, 20.0, 30.0, 60.0, 120.0)
@@ -323,6 +332,7 @@ class MetricsCollector(MetricsCollectorBase):
             "operation": operation,
             "bank_id": bank_id,
             "source": source,
+            "tenant": _get_tenant(),
         }
         if budget:
             attributes["budget"] = budget
@@ -373,6 +383,7 @@ class MetricsCollector(MetricsCollectorBase):
             "model": model,
             "scope": scope,
             "success": str(success).lower(),
+            "tenant": _get_tenant(),
         }
 
         # Record duration
@@ -425,10 +436,14 @@ class MetricsCollector(MetricsCollectorBase):
             status_code = status_code_getter()
             status_class = f"{status_code // 100}xx"
 
+            # Get tenant from context (may be set during request processing)
+            tenant = _get_tenant()
+
             attributes = {
                 **base_attributes,
                 "status_code": str(status_code),
                 "status_class": status_class,
+                "tenant": tenant,
             }
 
             # Record duration and count

hindsight-api/hindsight_api/migrations.py

@@ -22,6 +22,7 @@ from pathlib import Path
 
 from alembic import command
 from alembic.config import Config
+from alembic.script.revision import ResolutionError
 from sqlalchemy import create_engine, text
 
 logger = logging.getLogger(__name__)
@@ -78,7 +79,18 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
         alembic_cfg.set_main_option("target_schema", schema)
 
     # Run migrations
-    command.upgrade(alembic_cfg, "head")
+    try:
+        command.upgrade(alembic_cfg, "head")
+    except ResolutionError as e:
+        # This happens during rolling deployments when a newer version of the code
+        # has already run migrations, and this older replica doesn't have the new
+        # migration files. The database is already at a newer revision than we know.
+        # This is safe to ignore - the newer code has already applied its migrations.
+        logger.warning(
+            f"Database is at a newer migration revision than this code version knows about. "
+            f"This is expected during rolling deployments. Skipping migrations. Error: {e}"
+        )
+        return
 
     logger.info(f"Database migrations completed successfully for schema '{schema_name}'")
 

hindsight-docs/docs/developer/configuration.md

@@ -139,13 +139,18 @@ 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`) | - |
 | `HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL` | OpenAI embedding model | `text-embedding-3-small` |
+| `HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL` | Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI) | - |
 | `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
@@ -157,6 +162,12 @@ export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
 export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxxxxxxxxxx  # or reuses HINDSIGHT_API_LLM_API_KEY
 export HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small  # 1536 dimensions
 
+# Azure OpenAI - embeddings via Azure endpoint
+export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
+export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=your-azure-api-key
+export HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
+export HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://your-resource.openai.azure.com/openai/deployments/your-deployment
+
 # TEI - HuggingFace Text Embeddings Inference (recommended for production)
 export HINDSIGHT_API_EMBEDDINGS_PROVIDER=tei
 export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
@@ -165,6 +176,18 @@ export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
 export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere
 export HINDSIGHT_API_COHERE_API_KEY=your-api-key
 export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0  # 1024 dimensions
+
+# Azure-hosted Cohere - embeddings via custom endpoint
+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
@@ -187,13 +210,15 @@ 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 | - |
 | `HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE` | Batch size for TEI reranking | `128` |
 | `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
@@ -208,8 +233,27 @@ export HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
 export HINDSIGHT_API_RERANKER_PROVIDER=cohere
 export HINDSIGHT_API_COHERE_API_KEY=your-api-key  # shared with embeddings
 export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0
+
+# Azure-hosted Cohere - reranking via custom endpoint
+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:

monitoring/grafana/dashboards/hindsight-api-service.json

@@ -54,7 +54,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_http_requests_total)",
+          "expr": "sum(hindsight_http_requests_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -99,7 +99,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_http_requests_total[1m]))",
+          "expr": "sum(rate(hindsight_http_requests_total{tenant=~\"$tenant\"}[1m]))",
           "refId": "A"
         }
       ],
@@ -191,7 +191,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"5xx\"}[5m])) / sum(rate(hindsight_http_requests_total[5m]))",
+          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"5xx\", tenant=~\"$tenant\"}[5m])) / sum(rate(hindsight_http_requests_total{tenant=~\"$tenant\"}[5m]))",
           "refId": "A"
         }
       ],
@@ -236,7 +236,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_http_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_http_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "refId": "A"
         }
       ],
@@ -313,7 +313,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum by (endpoint) (rate(hindsight_http_requests_total[1m]))",
+          "expr": "sum by (endpoint) (rate(hindsight_http_requests_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{endpoint}}",
           "refId": "A"
         }
@@ -404,17 +404,17 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "histogram_quantile(0.50, sum by (le) (rate(hindsight_http_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.50, sum by (le) (rate(hindsight_http_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p50",
           "refId": "A"
         },
         {
-          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_http_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_http_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p95",
           "refId": "B"
         },
         {
-          "expr": "histogram_quantile(0.99, sum by (le) (rate(hindsight_http_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.99, sum by (le) (rate(hindsight_http_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p99",
           "refId": "C"
         }
@@ -505,12 +505,12 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"5xx\"}[1m])) / sum(rate(hindsight_http_requests_total[1m]))",
+          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"5xx\", tenant=~\"$tenant\"}[1m])) / sum(rate(hindsight_http_requests_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "5xx Error Rate",
           "refId": "A"
         },
         {
-          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"4xx\"}[1m])) / sum(rate(hindsight_http_requests_total[1m]))",
+          "expr": "sum(rate(hindsight_http_requests_total{status_class=\"4xx\", tenant=~\"$tenant\"}[1m])) / sum(rate(hindsight_http_requests_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "4xx Error Rate",
           "refId": "B"
         }
@@ -1276,7 +1276,36 @@
   "schemaVersion": 38,
   "tags": ["hindsight", "api", "service"],
   "templating": {
-    "list": []
+    "list": [
+      {
+        "allValue": ".*",
+        "current": {
+          "selected": true,
+          "text": "All",
+          "value": "$__all"
+        },
+        "datasource": {
+          "type": "prometheus",
+          "uid": "prometheus"
+        },
+        "definition": "label_values(hindsight_http_requests_total, tenant)",
+        "hide": 0,
+        "includeAll": true,
+        "label": "Tenant",
+        "multi": false,
+        "name": "tenant",
+        "options": [],
+        "query": {
+          "query": "label_values(hindsight_http_requests_total, tenant)",
+          "refId": "PrometheusVariableQueryEditor-VariableQuery"
+        },
+        "refresh": 2,
+        "regex": "",
+        "skipUrlSync": false,
+        "sort": 1,
+        "type": "query"
+      }
+    ]
   },
   "time": {
     "from": "now-30m",

monitoring/grafana/dashboards/hindsight-llm.json

@@ -46,7 +46,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_llm_calls_total)",
+          "expr": "sum(hindsight_llm_calls_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -91,7 +91,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_llm_tokens_input_tokens_total) + sum(hindsight_llm_tokens_output_tokens_total)",
+          "expr": "sum(hindsight_llm_tokens_input_tokens_total{tenant=~\"$tenant\"}) + sum(hindsight_llm_tokens_output_tokens_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -137,7 +137,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_llm_tokens_input_tokens_total)",
+          "expr": "sum(hindsight_llm_tokens_input_tokens_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -183,7 +183,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_llm_tokens_output_tokens_total)",
+          "expr": "sum(hindsight_llm_tokens_output_tokens_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -260,7 +260,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum by (scope) (rate(hindsight_llm_calls_total[1m]))",
+          "expr": "sum by (scope) (rate(hindsight_llm_calls_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{scope}}",
           "refId": "A"
         }
@@ -347,12 +347,12 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_llm_tokens_input_tokens_total[1m]))",
+          "expr": "sum(rate(hindsight_llm_tokens_input_tokens_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "Input",
           "refId": "A"
         },
         {
-          "expr": "sum(rate(hindsight_llm_tokens_output_tokens_total[1m]))",
+          "expr": "sum(rate(hindsight_llm_tokens_output_tokens_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "Output",
           "refId": "B"
         }
@@ -430,7 +430,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "histogram_quantile(0.95, sum by (scope, le) (rate(hindsight_llm_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.95, sum by (scope, le) (rate(hindsight_llm_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "{{scope}}",
           "refId": "A"
         }
@@ -508,12 +508,12 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum by (scope) (rate(hindsight_llm_tokens_input_tokens_total[1m]))",
+          "expr": "sum by (scope) (rate(hindsight_llm_tokens_input_tokens_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{scope}} (input)",
           "refId": "A"
         },
         {
-          "expr": "sum by (scope) (rate(hindsight_llm_tokens_output_tokens_total[1m]))",
+          "expr": "sum by (scope) (rate(hindsight_llm_tokens_output_tokens_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{scope}} (output)",
           "refId": "B"
         }
@@ -526,7 +526,36 @@
   "schemaVersion": 38,
   "tags": ["hindsight", "llm"],
   "templating": {
-    "list": []
+    "list": [
+      {
+        "allValue": ".*",
+        "current": {
+          "selected": true,
+          "text": "All",
+          "value": "$__all"
+        },
+        "datasource": {
+          "type": "prometheus",
+          "uid": "prometheus"
+        },
+        "definition": "label_values(hindsight_llm_calls_total, tenant)",
+        "hide": 0,
+        "includeAll": true,
+        "label": "Tenant",
+        "multi": false,
+        "name": "tenant",
+        "options": [],
+        "query": {
+          "query": "label_values(hindsight_llm_calls_total, tenant)",
+          "refId": "PrometheusVariableQueryEditor-VariableQuery"
+        },
+        "refresh": 2,
+        "regex": "",
+        "skipUrlSync": false,
+        "sort": 1,
+        "type": "query"
+      }
+    ]
   },
   "time": {
     "from": "now-30m",

monitoring/grafana/dashboards/hindsight-operations.json

@@ -46,7 +46,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(hindsight_operation_operations_total)",
+          "expr": "sum(hindsight_operation_operations_total{tenant=~\"$tenant\"})",
           "refId": "A"
         }
       ],
@@ -91,7 +91,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_operation_operations_total[1m]))",
+          "expr": "sum(rate(hindsight_operation_operations_total{tenant=~\"$tenant\"}[1m]))",
           "refId": "A"
         }
       ],
@@ -137,7 +137,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"retain\"}[1m]))",
+          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"retain\", tenant=~\"$tenant\"}[1m]))",
           "refId": "A"
         }
       ],
@@ -183,7 +183,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"recall\"}[1m]))",
+          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"recall\", tenant=~\"$tenant\"}[1m]))",
           "refId": "A"
         }
       ],
@@ -229,7 +229,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"reflect\"}[1m]))",
+          "expr": "sum(rate(hindsight_operation_operations_total{operation=\"reflect\", tenant=~\"$tenant\"}[1m]))",
           "refId": "A"
         }
       ],
@@ -319,7 +319,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum by (operation) (rate(hindsight_operation_operations_total[1m]))",
+          "expr": "sum by (operation) (rate(hindsight_operation_operations_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{operation}}",
           "refId": "A"
         }
@@ -410,17 +410,17 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "histogram_quantile(0.50, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\"}[5m])))",
+          "expr": "histogram_quantile(0.50, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\", tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p50",
           "refId": "A"
         },
         {
-          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\"}[5m])))",
+          "expr": "histogram_quantile(0.95, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\", tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p95",
           "refId": "B"
         },
         {
-          "expr": "histogram_quantile(0.99, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\"}[5m])))",
+          "expr": "histogram_quantile(0.99, sum by (le) (rate(hindsight_operation_duration_seconds_bucket{operation=\"recall\", tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "p99",
           "refId": "C"
         }
@@ -498,7 +498,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "histogram_quantile(0.95, sum by (operation, le) (rate(hindsight_operation_duration_seconds_bucket[5m])))",
+          "expr": "histogram_quantile(0.95, sum by (operation, le) (rate(hindsight_operation_duration_seconds_bucket{tenant=~\"$tenant\"}[5m])))",
           "legendFormat": "{{operation}}",
           "refId": "A"
         }
@@ -576,7 +576,7 @@
       "pluginVersion": "10.0.0",
       "targets": [
         {
-          "expr": "sum by (bank_id) (rate(hindsight_operation_operations_total[1m]))",
+          "expr": "sum by (bank_id) (rate(hindsight_operation_operations_total{tenant=~\"$tenant\"}[1m]))",
           "legendFormat": "{{bank_id}}",
           "refId": "A"
         }
@@ -589,7 +589,36 @@
   "schemaVersion": 38,
   "tags": ["hindsight"],
   "templating": {
-    "list": []
+    "list": [
+      {
+        "allValue": ".*",
+        "current": {
+          "selected": true,
+          "text": "All",
+          "value": "$__all"
+        },
+        "datasource": {
+          "type": "prometheus",
+          "uid": "prometheus"
+        },
+        "definition": "label_values(hindsight_operation_operations_total, tenant)",
+        "hide": 0,
+        "includeAll": true,
+        "label": "Tenant",
+        "multi": false,
+        "name": "tenant",
+        "options": [],
+        "query": {
+          "query": "label_values(hindsight_operation_operations_total, tenant)",
+          "refId": "PrometheusVariableQueryEditor-VariableQuery"
+        },
+        "refresh": 2,
+        "regex": "",
+        "skipUrlSync": false,
+        "sort": 1,
+        "type": "query"
+      }
+    ]
   },
   "time": {
     "from": "now-30m",