hindsight-api 0.0.21__py3-none-any.whl → 0.1.1__py3-none-any.whl

hindsight_api/config.py

@@ -0,0 +1,154 @@
+"""
+Centralized configuration for Hindsight API.
+
+All environment variables and their defaults are defined here.
+"""
+import os
+from dataclasses import dataclass
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Environment variable names
+ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
+ENV_LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
+ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
+ENV_LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
+ENV_LLM_BASE_URL = "HINDSIGHT_API_LLM_BASE_URL"
+
+ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
+ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
+ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
+
+ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER"
+ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
+ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
+
+ENV_HOST = "HINDSIGHT_API_HOST"
+ENV_PORT = "HINDSIGHT_API_PORT"
+ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
+ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
+
+# Default values
+DEFAULT_DATABASE_URL = "pg0"
+DEFAULT_LLM_PROVIDER = "groq"
+DEFAULT_LLM_MODEL = "openai/gpt-oss-20b"
+
+DEFAULT_EMBEDDINGS_PROVIDER = "local"
+DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
+
+DEFAULT_RERANKER_PROVIDER = "local"
+DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+DEFAULT_HOST = "0.0.0.0"
+DEFAULT_PORT = 8888
+DEFAULT_LOG_LEVEL = "info"
+DEFAULT_MCP_ENABLED = True
+
+# Required embedding dimension for database schema
+EMBEDDING_DIMENSION = 384
+
+
+@dataclass
+class HindsightConfig:
+    """Configuration container for Hindsight API."""
+
+    # Database
+    database_url: str
+
+    # LLM
+    llm_provider: str
+    llm_api_key: Optional[str]
+    llm_model: str
+    llm_base_url: Optional[str]
+
+    # Embeddings
+    embeddings_provider: str
+    embeddings_local_model: str
+    embeddings_tei_url: Optional[str]
+
+    # Reranker
+    reranker_provider: str
+    reranker_local_model: str
+    reranker_tei_url: Optional[str]
+
+    # Server
+    host: str
+    port: int
+    log_level: str
+    mcp_enabled: bool
+
+    @classmethod
+    def from_env(cls) -> "HindsightConfig":
+        """Create configuration from environment variables."""
+        return cls(
+            # Database
+            database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
+
+            # LLM
+            llm_provider=os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER),
+            llm_api_key=os.getenv(ENV_LLM_API_KEY),
+            llm_model=os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL),
+            llm_base_url=os.getenv(ENV_LLM_BASE_URL) or None,
+
+            # Embeddings
+            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),
+
+            # Reranker
+            reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
+            reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
+            reranker_tei_url=os.getenv(ENV_RERANKER_TEI_URL),
+
+            # Server
+            host=os.getenv(ENV_HOST, DEFAULT_HOST),
+            port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
+            log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
+            mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
+        )
+
+    def get_llm_base_url(self) -> str:
+        """Get the LLM base URL, with provider-specific defaults."""
+        if self.llm_base_url:
+            return self.llm_base_url
+
+        provider = self.llm_provider.lower()
+        if provider == "groq":
+            return "https://api.groq.com/openai/v1"
+        elif provider == "ollama":
+            return "http://localhost:11434/v1"
+        else:
+            return ""
+
+    def get_python_log_level(self) -> int:
+        """Get the Python logging level from the configured log level string."""
+        log_level_map = {
+            "critical": logging.CRITICAL,
+            "error": logging.ERROR,
+            "warning": logging.WARNING,
+            "info": logging.INFO,
+            "debug": logging.DEBUG,
+            "trace": logging.DEBUG,  # Python doesn't have TRACE, use DEBUG
+        }
+        return log_level_map.get(self.log_level.lower(), logging.INFO)
+
+    def configure_logging(self) -> None:
+        """Configure Python logging based on the log level."""
+        logging.basicConfig(
+            level=self.get_python_log_level(),
+            format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+        )
+
+    def log_config(self) -> None:
+        """Log the current configuration (without sensitive values)."""
+        logger.info(f"Database: {self.database_url}")
+        logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
+        logger.info(f"Embeddings: provider={self.embeddings_provider}")
+        logger.info(f"Reranker: provider={self.reranker_provider}")
+
+
+def get_config() -> HindsightConfig:
+    """Get the current configuration from environment variables."""
+    return HindsightConfig.from_env()

hindsight_api/engine/__init__.py

@@ -9,7 +9,8 @@ This package contains all the implementation details of the memory engine:
 
 from .memory_engine import MemoryEngine
 from .db_utils import acquire_with_retry
-from .embeddings import Embeddings, SentenceTransformersEmbeddings
+from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
+from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
 from .search.trace import (
     SearchTrace,
     QueryInfo,
@@ -29,7 +30,11 @@ __all__ = [
     "MemoryEngine",
     "acquire_with_retry",
     "Embeddings",
-    "SentenceTransformersEmbeddings",
+    "LocalSTEmbeddings",
+    "RemoteTEIEmbeddings",
+    "CrossEncoderModel",
+    "LocalSTCrossEncoder",
+    "RemoteTEICrossEncoder",
     "SearchTrace",
     "SearchTracer",
     "QueryInfo",

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,13 +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}...")
-        self._model = CrossEncoder(self.model_name)
-        logger.info("Cross-encoder model loaded")
+        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("Reranker: local provider initialized")
 
     def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
         """
@@ -92,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'"
+        )