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

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,12 +100,17 @@ 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}...")
-        self._model = SentenceTransformer(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(
+            self.model_name,
+            model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
+        )
 
         # Validate dimension matches database schema
         model_dim = self._model.get_sentence_embedding_dimension()
@@ -95,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]]:
         """
@@ -108,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'"
+        )

hindsight_api/engine/entity_resolver.py

@@ -126,18 +126,20 @@ class EntityResolver:
 
         # Resolve each entity using pre-fetched candidates
         entity_ids = [None] * len(entities_data)
-        entities_to_update = []  # (entity_id, unit_event_date)
-        entities_to_create = []  # (idx, entity_data)
+        entities_to_update = []  # (entity_id, event_date)
+        entities_to_create = []  # (idx, entity_data, event_date)
 
         for idx, entity_data in enumerate(entities_data):
             entity_text = entity_data['text']
             nearby_entities = entity_data.get('nearby_entities', [])
+            # Use per-entity date if available, otherwise fall back to batch-level date
+            entity_event_date = entity_data.get('event_date', unit_event_date)
 
             candidates = all_candidates.get(entity_text, [])
 
             if not candidates:
                 # Will create new entity
-                entities_to_create.append((idx, entity_data))
+                entities_to_create.append((idx, entity_data, entity_event_date))
                 continue
 
             # Score candidates
@@ -165,9 +167,9 @@ class EntityResolver:
                     score += co_entity_score * 0.3
 
                 # 3. Temporal proximity (0-0.2)
-                if last_seen:
+                if last_seen and entity_event_date:
                     # Normalize timezone awareness for comparison
-                    event_date_utc = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=timezone.utc)
+                    event_date_utc = entity_event_date if entity_event_date.tzinfo else entity_event_date.replace(tzinfo=timezone.utc)
                     last_seen_utc = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=timezone.utc)
                     days_diff = abs((event_date_utc - last_seen_utc).total_seconds() / 86400)
                     if days_diff < 7:
@@ -183,9 +185,9 @@ class EntityResolver:
 
             if best_score > threshold:
                 entity_ids[idx] = best_candidate
-                entities_to_update.append((best_candidate, unit_event_date))
+                entities_to_update.append((best_candidate, entity_event_date))
             else:
-                entities_to_create.append((idx, entity_data))
+                entities_to_create.append((idx, entity_data, entity_event_date))
 
         # Batch update existing entities
         if entities_to_update:
@@ -199,29 +201,54 @@ class EntityResolver:
                 entities_to_update
             )
 
-        # Create new entities using INSERT ... ON CONFLICT to handle race conditions
-        # This ensures that if two concurrent transactions try to create the same entity,
-        # only one succeeds and the other gets the existing ID
+        # Batch create new entities using COPY + INSERT for maximum speed
+        # This handles duplicates via ON CONFLICT and returns all IDs
         if entities_to_create:
-            for idx, entity_data in entities_to_create:
-                # Use INSERT ... ON CONFLICT to atomically get-or-create
-                # The unique index is on (bank_id, LOWER(canonical_name))
-                row = await conn.fetchrow(
-                    """
-                    INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
-                    VALUES ($1, $2, $3, $4, 1)
-                    ON CONFLICT (bank_id, LOWER(canonical_name))
-                    DO UPDATE SET
-                        mention_count = entities.mention_count + 1,
-                        last_seen = EXCLUDED.last_seen
-                    RETURNING id
-                    """,
-                    bank_id,
-                    entity_data['text'],
-                    unit_event_date,
-                    unit_event_date
-                )
-                entity_ids[idx] = row['id']
+            # Group entities by canonical name (lowercase) to handle duplicates within batch
+            # For duplicates, we only insert once and reuse the ID
+            unique_entities = {}  # lowercase_name -> (entity_data, event_date, [indices])
+            for idx, entity_data, event_date in entities_to_create:
+                name_lower = entity_data['text'].lower()
+                if name_lower not in unique_entities:
+                    unique_entities[name_lower] = (entity_data, event_date, [idx])
+                else:
+                    # Same entity appears multiple times - add index to list
+                    unique_entities[name_lower][2].append(idx)
+
+            # Batch insert unique entities and get their IDs
+            # Use a single query with unnest for speed
+            entity_names = []
+            entity_dates = []
+            indices_map = []  # Maps result index -> list of original indices
+
+            for name_lower, (entity_data, event_date, indices) in unique_entities.items():
+                entity_names.append(entity_data['text'])
+                entity_dates.append(event_date)
+                indices_map.append(indices)
+
+            # Batch INSERT ... ON CONFLICT with RETURNING
+            # This is much faster than individual inserts
+            rows = await conn.fetch(
+                """
+                INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
+                SELECT $1, name, event_date, event_date, 1
+                FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
+                ON CONFLICT (bank_id, LOWER(canonical_name))
+                DO UPDATE SET
+                    mention_count = entities.mention_count + 1,
+                    last_seen = EXCLUDED.last_seen
+                RETURNING id
+                """,
+                bank_id,
+                entity_names,
+                entity_dates
+            )
+
+            # Map returned IDs back to original indices
+            for result_idx, row in enumerate(rows):
+                entity_id = row['id']
+                for original_idx in indices_map[result_idx]:
+                    entity_ids[original_idx] = entity_id
 
         return entity_ids