hindsight-api 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

hindsight_api/engine/cross_encoder.py

@@ -2,10 +2,23 @@
 Cross-encoder abstraction for reranking.
 
 Provides an interface for reranking with different backends.
+
+Configuration via environment variables - see hindsight_api.config for all env var names.
 """
 from abc import ABC, abstractmethod
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 import logging
+import os
+
+import httpx
+
+from ..config import (
+    ENV_RERANKER_PROVIDER,
+    ENV_RERANKER_LOCAL_MODEL,
+    ENV_RERANKER_TEI_URL,
+    DEFAULT_RERANKER_PROVIDER,
+    DEFAULT_RERANKER_LOCAL_MODEL,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -17,12 +30,18 @@ class CrossEncoderModel(ABC):
     Cross-encoders take query-document pairs and return relevance scores.
     """
 
+    @property
     @abstractmethod
-    def load(self) -> None:
+    def provider_name(self) -> str:
+        """Return a human-readable name for this provider (e.g., 'local', 'tei')."""
+        pass
+
+    @abstractmethod
+    async def initialize(self) -> None:
         """
-        Load the cross-encoder model.
+        Initialize the cross-encoder model asynchronously.
 
-        This should be called during initialization to load the model
+        This should be called during startup to load/connect to the model
         and avoid cold start latency on first predict() call.
         """
         pass
@@ -41,11 +60,11 @@ class CrossEncoderModel(ABC):
         pass
 
 
-class SentenceTransformersCrossEncoder(CrossEncoderModel):
+class LocalSTCrossEncoder(CrossEncoderModel):
     """
-    Cross-encoder implementation using SentenceTransformers.
+    Local cross-encoder implementation using SentenceTransformers.
 
-    Call load() during initialization to load the model and avoid cold starts.
+    Call initialize() during startup to load the model and avoid cold starts.
 
     Default model is cross-encoder/ms-marco-MiniLM-L-6-v2:
     - Fast inference (~80ms for 100 pairs on CPU)
@@ -53,18 +72,22 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
     - Trained for passage re-ranking
     """
 
-    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
+    def __init__(self, model_name: Optional[str] = None):
         """
-        Initialize SentenceTransformers cross-encoder.
+        Initialize local SentenceTransformers cross-encoder.
 
         Args:
             model_name: Name of the CrossEncoder model to use.
                        Default: cross-encoder/ms-marco-MiniLM-L-6-v2
         """
-        self.model_name = model_name
+        self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
         self._model = None
 
-    def load(self) -> None:
+    @property
+    def provider_name(self) -> str:
+        return "local"
+
+    async def initialize(self) -> None:
         """Load the cross-encoder model."""
         if self._model is not None:
             return
@@ -73,18 +96,18 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
             from sentence_transformers import CrossEncoder
         except ImportError:
             raise ImportError(
-                "sentence-transformers is required for SentenceTransformersCrossEncoder. "
+                "sentence-transformers is required for LocalSTCrossEncoder. "
                 "Install it with: pip install sentence-transformers"
             )
 
-        logger.info(f"Loading cross-encoder model: {self.model_name}...")
+        logger.info(f"Reranker: initializing local provider with model {self.model_name}")
         # Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
         # Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
         self._model = CrossEncoder(
             self.model_name,
             model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
         )
-        logger.info("Cross-encoder model loaded")
+        logger.info("Reranker: local provider initialized")
 
     def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
         """
@@ -97,6 +120,187 @@ class SentenceTransformersCrossEncoder(CrossEncoderModel):
             List of relevance scores (raw logits from the model)
         """
         if self._model is None:
-            self.load()
+            raise RuntimeError("Reranker not initialized. Call initialize() first.")
         scores = self._model.predict(pairs, show_progress_bar=False)
         return scores.tolist() if hasattr(scores, 'tolist') else list(scores)
+
+
+class RemoteTEICrossEncoder(CrossEncoderModel):
+    """
+    Remote cross-encoder implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
+
+    TEI supports reranking via the /rerank endpoint.
+    See: https://github.com/huggingface/text-embeddings-inference
+
+    Note: The TEI server must be running a cross-encoder/reranker model.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 30.0,
+        batch_size: int = 32,
+        max_retries: int = 3,
+        retry_delay: float = 0.5,
+    ):
+        """
+        Initialize remote TEI cross-encoder client.
+
+        Args:
+            base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
+            timeout: Request timeout in seconds (default: 30.0)
+            batch_size: Maximum batch size for rerank requests (default: 32)
+            max_retries: Maximum number of retries for failed requests (default: 3)
+            retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
+        """
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.batch_size = batch_size
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self._client: Optional[httpx.Client] = None
+        self._model_id: Optional[str] = None
+
+    @property
+    def provider_name(self) -> str:
+        return "tei"
+
+    def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
+        """Make an HTTP request with automatic retries on transient errors."""
+        import time
+        last_error = None
+        delay = self.retry_delay
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                if method == "GET":
+                    response = self._client.get(url, **kwargs)
+                else:
+                    response = self._client.post(url, **kwargs)
+                response.raise_for_status()
+                return response
+            except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
+                last_error = e
+                if attempt < self.max_retries:
+                    logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
+                    time.sleep(delay)
+                    delay *= 2  # Exponential backoff
+            except httpx.HTTPStatusError as e:
+                # Retry on 5xx server errors
+                if e.response.status_code >= 500 and attempt < self.max_retries:
+                    last_error = e
+                    logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
+                    time.sleep(delay)
+                    delay *= 2
+                else:
+                    raise
+
+        raise last_error
+
+    async def initialize(self) -> None:
+        """Initialize the HTTP client and verify server connectivity."""
+        if self._client is not None:
+            return
+
+        logger.info(f"Reranker: initializing TEI provider at {self.base_url}")
+        self._client = httpx.Client(timeout=self.timeout)
+
+        # Verify server is reachable and get model info
+        try:
+            response = self._request_with_retry("GET", f"{self.base_url}/info")
+            info = response.json()
+            self._model_id = info.get("model_id", "unknown")
+            logger.info(f"Reranker: TEI provider initialized (model: {self._model_id})")
+        except httpx.HTTPError as e:
+            raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
+
+    def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
+        """
+        Score query-document pairs using the remote TEI reranker.
+
+        Args:
+            pairs: List of (query, document) tuples to score
+
+        Returns:
+            List of relevance scores
+        """
+        if self._client is None:
+            raise RuntimeError("Reranker not initialized. Call initialize() first.")
+
+        if not pairs:
+            return []
+
+        all_scores = []
+
+        # Process in batches
+        for i in range(0, len(pairs), self.batch_size):
+            batch = pairs[i:i + self.batch_size]
+
+            # TEI rerank endpoint expects query and texts separately
+            # All pairs in a batch should have the same query for optimal performance
+            # but we handle mixed queries by making separate requests per unique query
+            query_groups: dict[str, list[tuple[int, str]]] = {}
+            for idx, (query, text) in enumerate(batch):
+                if query not in query_groups:
+                    query_groups[query] = []
+                query_groups[query].append((idx, text))
+
+            batch_scores = [0.0] * len(batch)
+
+            for query, indexed_texts in query_groups.items():
+                texts = [text for _, text in indexed_texts]
+                indices = [idx for idx, _ in indexed_texts]
+
+                try:
+                    response = self._request_with_retry(
+                        "POST",
+                        f"{self.base_url}/rerank",
+                        json={
+                            "query": query,
+                            "texts": texts,
+                            "return_text": False,
+                        },
+                    )
+                    results = response.json()
+
+                    # TEI returns results sorted by score descending, with original index
+                    for result in results:
+                        original_idx = result["index"]
+                        score = result["score"]
+                        # Map back to batch position
+                        batch_scores[indices[original_idx]] = score
+
+                except httpx.HTTPError as e:
+                    raise RuntimeError(f"TEI rerank request failed: {e}")
+
+            all_scores.extend(batch_scores)
+
+        return all_scores
+
+
+def create_cross_encoder_from_env() -> CrossEncoderModel:
+    """
+    Create a CrossEncoderModel instance based on environment variables.
+
+    See hindsight_api.config for environment variable names and defaults.
+
+    Returns:
+        Configured CrossEncoderModel instance
+    """
+    provider = os.environ.get(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER).lower()
+
+    if provider == "tei":
+        url = os.environ.get(ENV_RERANKER_TEI_URL)
+        if not url:
+            raise ValueError(
+                f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'"
+            )
+        return RemoteTEICrossEncoder(base_url=url)
+    elif provider == "local":
+        model = os.environ.get(ENV_RERANKER_LOCAL_MODEL)
+        model_name = model or DEFAULT_RERANKER_LOCAL_MODEL
+        return LocalSTCrossEncoder(model_name=model_name)
+    else:
+        raise ValueError(
+            f"Unknown reranker provider: {provider}. Supported: 'local', 'tei'"
+        )

hindsight_api/engine/embeddings.py

@@ -5,15 +5,26 @@ Provides an interface for generating embeddings with different backends.
 
 IMPORTANT: All embeddings must produce 384-dimensional vectors to match
 the database schema (pgvector column defined as vector(384)).
+
+Configuration via environment variables - see hindsight_api.config for all env var names.
 """
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Optional
 import logging
+import os
 
-logger = logging.getLogger(__name__)
+import httpx
 
-# Fixed embedding dimension required by database schema
-EMBEDDING_DIMENSION = 384
+from ..config import (
+    ENV_EMBEDDINGS_PROVIDER,
+    ENV_EMBEDDINGS_LOCAL_MODEL,
+    ENV_EMBEDDINGS_TEI_URL,
+    DEFAULT_EMBEDDINGS_PROVIDER,
+    DEFAULT_EMBEDDINGS_LOCAL_MODEL,
+    EMBEDDING_DIMENSION,
+)
+
+logger = logging.getLogger(__name__)
 
 
 class Embeddings(ABC):
@@ -24,12 +35,18 @@ class Embeddings(ABC):
     the database schema.
     """
 
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """Return a human-readable name for this provider (e.g., 'local', 'tei')."""
+        pass
+
     @abstractmethod
-    def load(self) -> None:
+    async def initialize(self) -> None:
         """
-        Load the embedding model.
+        Initialize the embedding model asynchronously.
 
-        This should be called during initialization to load the model
+        This should be called during startup to load/connect to the model
         and avoid cold start latency on first encode() call.
         """
         pass
@@ -48,29 +65,33 @@ class Embeddings(ABC):
         pass
 
 
-class SentenceTransformersEmbeddings(Embeddings):
+class LocalSTEmbeddings(Embeddings):
     """
-    Embeddings implementation using SentenceTransformers.
+    Local embeddings implementation using SentenceTransformers.
 
-    Call load() during initialization to load the model and avoid cold starts.
+    Call initialize() during startup to load the model and avoid cold starts.
 
     Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional
     embeddings matching the database schema.
     """
 
-    def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
+    def __init__(self, model_name: Optional[str] = None):
         """
-        Initialize SentenceTransformers embeddings.
+        Initialize local SentenceTransformers embeddings.
 
         Args:
             model_name: Name of the SentenceTransformer model to use.
                        Must produce 384-dimensional embeddings.
                        Default: BAAI/bge-small-en-v1.5
         """
-        self.model_name = model_name
+        self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
         self._model = None
 
-    def load(self) -> None:
+    @property
+    def provider_name(self) -> str:
+        return "local"
+
+    async def initialize(self) -> None:
         """Load the embedding model."""
         if self._model is not None:
             return
@@ -79,11 +100,11 @@ class SentenceTransformersEmbeddings(Embeddings):
             from sentence_transformers import SentenceTransformer
         except ImportError:
             raise ImportError(
-                "sentence-transformers is required for SentenceTransformersEmbeddings. "
+                "sentence-transformers is required for LocalSTEmbeddings. "
                 "Install it with: pip install sentence-transformers"
             )
 
-        logger.info(f"Loading embedding model: {self.model_name}...")
+        logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
         # Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
         # Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
         self._model = SentenceTransformer(
@@ -100,7 +121,7 @@ class SentenceTransformersEmbeddings(Embeddings):
                 f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings."
             )
 
-        logger.info(f"Model loaded (embedding dim: {model_dim})")
+        logger.info(f"Embeddings: local provider initialized (dim: {model_dim})")
 
     def encode(self, texts: List[str]) -> List[List[float]]:
         """
@@ -113,6 +134,159 @@ class SentenceTransformersEmbeddings(Embeddings):
             List of 384-dimensional embedding vectors
         """
         if self._model is None:
-            self.load()
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
         embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
         return [emb.tolist() for emb in embeddings]
+
+
+class RemoteTEIEmbeddings(Embeddings):
+    """
+    Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
+
+    TEI provides a high-performance inference server for embedding models.
+    See: https://github.com/huggingface/text-embeddings-inference
+
+    The server should be running a model that produces 384-dimensional embeddings.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        timeout: float = 30.0,
+        batch_size: int = 32,
+        max_retries: int = 3,
+        retry_delay: float = 0.5,
+    ):
+        """
+        Initialize remote TEI embeddings client.
+
+        Args:
+            base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
+            timeout: Request timeout in seconds (default: 30.0)
+            batch_size: Maximum batch size for embedding requests (default: 32)
+            max_retries: Maximum number of retries for failed requests (default: 3)
+            retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
+        """
+        self.base_url = base_url.rstrip("/")
+        self.timeout = timeout
+        self.batch_size = batch_size
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self._client: Optional[httpx.Client] = None
+        self._model_id: Optional[str] = None
+
+    @property
+    def provider_name(self) -> str:
+        return "tei"
+
+    def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
+        """Make an HTTP request with automatic retries on transient errors."""
+        import time
+        last_error = None
+        delay = self.retry_delay
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                if method == "GET":
+                    response = self._client.get(url, **kwargs)
+                else:
+                    response = self._client.post(url, **kwargs)
+                response.raise_for_status()
+                return response
+            except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
+                last_error = e
+                if attempt < self.max_retries:
+                    logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
+                    time.sleep(delay)
+                    delay *= 2  # Exponential backoff
+            except httpx.HTTPStatusError as e:
+                # Retry on 5xx server errors
+                if e.response.status_code >= 500 and attempt < self.max_retries:
+                    last_error = e
+                    logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
+                    time.sleep(delay)
+                    delay *= 2
+                else:
+                    raise
+
+        raise last_error
+
+    async def initialize(self) -> None:
+        """Initialize the HTTP client and verify server connectivity."""
+        if self._client is not None:
+            return
+
+        logger.info(f"Embeddings: initializing TEI provider at {self.base_url}")
+        self._client = httpx.Client(timeout=self.timeout)
+
+        # Verify server is reachable and get model info
+        try:
+            response = self._request_with_retry("GET", f"{self.base_url}/info")
+            info = response.json()
+            self._model_id = info.get("model_id", "unknown")
+            logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id})")
+        except httpx.HTTPError as e:
+            raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
+
+    def encode(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings using the remote TEI server.
+
+        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]
+
+            try:
+                response = self._request_with_retry(
+                    "POST",
+                    f"{self.base_url}/embed",
+                    json={"inputs": batch},
+                )
+                batch_embeddings = response.json()
+                all_embeddings.extend(batch_embeddings)
+            except httpx.HTTPError as e:
+                raise RuntimeError(f"TEI embedding request failed: {e}")
+
+        return all_embeddings
+
+
+def create_embeddings_from_env() -> Embeddings:
+    """
+    Create an Embeddings instance based on environment variables.
+
+    See hindsight_api.config for environment variable names and defaults.
+
+    Returns:
+        Configured Embeddings instance
+    """
+    provider = os.environ.get(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER).lower()
+
+    if provider == "tei":
+        url = os.environ.get(ENV_EMBEDDINGS_TEI_URL)
+        if not url:
+            raise ValueError(
+                f"{ENV_EMBEDDINGS_TEI_URL} is required when {ENV_EMBEDDINGS_PROVIDER} is 'tei'"
+            )
+        return RemoteTEIEmbeddings(base_url=url)
+    elif provider == "local":
+        model = os.environ.get(ENV_EMBEDDINGS_LOCAL_MODEL)
+        model_name = model or DEFAULT_EMBEDDINGS_LOCAL_MODEL
+        return LocalSTEmbeddings(model_name=model_name)
+    else:
+        raise ValueError(
+            f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei'"
+        )