tribalmemory 0.1.0__py3-none-any.whl

tribalmemory/services/embeddings.py

@@ -0,0 +1,273 @@
+"""OpenAI Embedding Service.
+
+Production embedding service using OpenAI's text-embedding-3-small model.
+"""
+
+import asyncio
+import math
+import os
+from typing import Optional
+
+import httpx
+
+from ..interfaces import IEmbeddingService
+from ..utils import normalize_embedding
+
+
+class OpenAIEmbeddingService(IEmbeddingService):
+    """OpenAI-compatible embedding service.
+    
+    Supports OpenAI, Ollama, and any OpenAI-compatible embedding API.
+    
+    Features:
+    - Async HTTP calls with retry logic
+    - Batch embedding support
+    - Rate limiting awareness
+    - Configurable model, dimensions, and API base URL
+    - Local-only mode (Ollama) — no API key needed
+    
+    Usage:
+        # OpenAI
+        service = OpenAIEmbeddingService(api_key="sk-...")
+        
+        # Ollama (local)
+        service = OpenAIEmbeddingService(
+            api_base="http://localhost:11434/v1",
+            model="nomic-embed-text",
+            dimensions=768,
+        )
+    """
+    
+    DEFAULT_MODEL = "text-embedding-3-small"
+    DEFAULT_DIMENSIONS = 1536
+    DEFAULT_API_BASE = "https://api.openai.com/v1"
+    LOCAL_API_KEY_PLACEHOLDER = "local-no-key-needed"
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = DEFAULT_MODEL,
+        dimensions: int = DEFAULT_DIMENSIONS,
+        max_retries: int = 3,
+        timeout_seconds: float = 30.0,
+        backoff_base: float = 2.0,
+        backoff_max: float = 60.0,
+        api_base: Optional[str] = None,
+    ):
+        """Initialize OpenAI-compatible embedding service.
+        
+        Supports OpenAI, Ollama, and any OpenAI-compatible embedding API.
+        
+        Args:
+            api_key: API key. Falls back to OPENAI_API_KEY env var.
+                     Not required when api_base points to a local service (e.g., Ollama).
+            model: Embedding model to use.
+            dimensions: Output embedding dimensions.
+            max_retries: Max retry attempts on transient failures.
+            timeout_seconds: Request timeout.
+            backoff_base: Base for exponential backoff (default 2.0).
+            backoff_max: Maximum backoff delay in seconds (default 60.0).
+            api_base: Base URL for the embedding API. Defaults to OpenAI.
+                      For Ollama: "http://localhost:11434/v1"
+                      For any OpenAI-compatible API: "http://host:port/v1"
+        
+        Security Note:
+            API keys are stored in memory and used in HTTP headers. Never log
+            the _client object or include it in error reports. For production,
+            consider using a secrets manager rather than environment variables.
+        """
+        # Validate dimensions
+        if dimensions < 1 or dimensions > 8192:
+            raise ValueError(
+                f"Dimensions must be between 1 and 8192, got {dimensions}"
+            )
+        
+        # Build the API URL from api_base
+        self.api_url = self._resolve_api_url(api_base)
+        
+        # Determine if this is a local (non-OpenAI) service
+        is_local = (
+            api_base is not None
+            and api_base.strip() != ""
+            and "api.openai.com" not in api_base.lower()
+        )
+        
+        # For local services (non-OpenAI), api_key is optional
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        self.backoff_base = backoff_base
+        self.backoff_max = backoff_max
+        if not self.api_key:
+            if not is_local:
+                raise ValueError(
+                    "OpenAI API key required. Pass api_key "
+                    "or set OPENAI_API_KEY env var."
+                )
+            # Use a placeholder for local services (e.g., Ollama)
+            self.api_key = self.LOCAL_API_KEY_PLACEHOLDER
+        
+        self.model = model
+        self.dimensions = dimensions
+        self.max_retries = max_retries
+        self.timeout_seconds = timeout_seconds
+        
+        self._client: Optional[httpx.AsyncClient] = None
+    
+    @staticmethod
+    def _resolve_api_url(api_base: Optional[str] = None) -> str:
+        """Resolve the full embeddings API URL from an optional base.
+        
+        Args:
+            api_base: Base URL (e.g., "http://localhost:11434/v1").
+                      If None, uses OpenAI default.
+        
+        Returns:
+            Full URL ending in /embeddings.
+        
+        Raises:
+            ValueError: If api_base is not a valid HTTP(S) URL.
+        """
+        if api_base is None or api_base.strip() == "":
+            return (
+                f"{OpenAIEmbeddingService.DEFAULT_API_BASE}"
+                "/embeddings"
+            )
+        
+        base = api_base.strip().rstrip("/")
+        
+        # Basic URL validation
+        if base and not base.startswith(("http://", "https://")):
+            raise ValueError(
+                f"api_base must be an HTTP(S) URL, got: {base}"
+            )
+        
+        # If already ends with /embeddings, use as-is
+        if base.endswith("/embeddings"):
+            return base
+        
+        return f"{base}/embeddings"
+    
+    def __repr__(self) -> str:
+        """Safe repr that masks API key to prevent accidental logging."""
+        return f"OpenAIEmbeddingService(model={self.model!r}, api_key=***)"
+    
+    async def _get_client(self) -> httpx.AsyncClient:
+        """Get or create HTTP client."""
+        if self._client is None or self._client.is_closed:
+            self._client = httpx.AsyncClient(
+                timeout=httpx.Timeout(self.timeout_seconds),
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                }
+            )
+        return self._client
+    
+    async def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        results = await self.embed_batch([text])
+        return results[0]
+    
+    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts."""
+        if not texts:
+            return []
+        
+        cleaned_texts = [self._clean_text(t) for t in texts]
+        client = await self._get_client()
+        
+        payload = {
+            "model": self.model,
+            "input": cleaned_texts,
+            "dimensions": self.dimensions,
+        }
+        
+        last_error = None
+        for attempt in range(self.max_retries):
+            try:
+                response = await client.post(self.api_url, json=payload)
+                
+                if response.status_code == 200:
+                    data = response.json()
+                    embeddings = sorted(data["data"], key=lambda x: x["index"])
+                    return [normalize_embedding(e["embedding"]) for e in embeddings]
+                
+                elif response.status_code == 429:
+                    retry_after = int(response.headers.get("Retry-After", 5))
+                    await asyncio.sleep(retry_after)
+                    continue
+                
+                elif response.status_code >= 500:
+                    backoff = min(self.backoff_base ** attempt, self.backoff_max)
+                    await asyncio.sleep(backoff)
+                    continue
+                
+                else:
+                    try:
+                        error_detail = (
+                            response.json()
+                            .get("error", {})
+                            .get("message", response.text)
+                        )
+                    except Exception:
+                        error_detail = response.text[:200]
+                    raise RuntimeError(
+                        f"Embedding API error "
+                        f"({response.status_code}): "
+                        f"{error_detail}"
+                    )
+                    
+            except httpx.TimeoutException as e:
+                last_error = e
+                backoff = min(self.backoff_base ** attempt, self.backoff_max)
+                await asyncio.sleep(backoff)
+            except httpx.RequestError as e:
+                last_error = e
+                backoff = min(self.backoff_base ** attempt, self.backoff_max)
+                await asyncio.sleep(backoff)
+        
+        raise RuntimeError(f"OpenAI API failed after {self.max_retries} retries: {last_error}")
+    
+    def similarity(self, a: list[float], b: list[float]) -> float:
+        """Calculate cosine similarity between two embeddings."""
+        if len(a) != len(b):
+            raise ValueError(f"Embedding dimensions don't match: {len(a)} vs {len(b)}")
+        
+        dot_product = sum(x * y for x, y in zip(a, b))
+        norm_a = math.sqrt(sum(x * x for x in a))
+        norm_b = math.sqrt(sum(x * x for x in b))
+        
+        if norm_a == 0 or norm_b == 0:
+            return 0.0
+        
+        return dot_product / (norm_a * norm_b)
+    
+    def _clean_text(self, text: str) -> str:
+        """Clean text for embedding.
+        
+        Normalizes whitespace and truncates to fit within token limits.
+        Uses encoding-aware truncation to avoid splitting UTF-8 characters.
+        """
+        cleaned = " ".join(text.split())
+        
+        # OpenAI has token limits; truncate very long texts
+        # text-embedding-3-small supports 8191 tokens
+        # Rough estimate: 4 bytes per token (worst case UTF-8)
+        max_bytes = 8191 * 4
+        
+        encoded = cleaned.encode('utf-8')
+        if len(encoded) > max_bytes:
+            # Truncate bytes and decode safely (ignore partial chars)
+            cleaned = encoded[:max_bytes].decode('utf-8', errors='ignore')
+        
+        return cleaned
+
+    async def close(self):
+        """Close the HTTP client."""
+        if self._client and not self._client.is_closed:
+            await self._client.aclose()
+    
+    async def __aenter__(self):
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()