"""
Farnsworth Embeddings - Multi-Model Embedding Support
Supports multiple embedding backends:
- Sentence Transformers (local)
- Ollama embeddings
- OpenAI-compatible APIs
"""
import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional
import hashlib
from loguru import logger
try:
import numpy as np
except ImportError:
np = None
@dataclass
class EmbeddingResult:
"""Result from embedding operation."""
embedding: list[float]
model: str
dimensions: int
cached: bool = False
class EmbeddingBackend(ABC):
"""Abstract embedding backend."""
@property
@abstractmethod
def dimensions(self) -> int:
pass
@abstractmethod
async def embed(self, text: str) -> list[float]:
pass
@abstractmethod
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
pass
class SentenceTransformerBackend(EmbeddingBackend):
"""Sentence Transformers embedding backend."""
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
self.model_name = model_name
self._model = None
self._dimensions = 384 # Default for MiniLM
@property
def dimensions(self) -> int:
return self._dimensions
def _load_model(self):
if self._model is None:
try:
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer(self.model_name)
self._dimensions = self._model.get_sentence_embedding_dimension()
except ImportError:
raise RuntimeError("sentence-transformers not installed")
async def embed(self, text: str) -> list[float]:
self._load_model()
loop = asyncio.get_event_loop()
embedding = await loop.run_in_executor(
None,
lambda: self._model.encode(text, normalize_embeddings=True)
)
return embedding.tolist()
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
self._load_model()
loop = asyncio.get_event_loop()
embeddings = await loop.run_in_executor(
None,
lambda: self._model.encode(texts, normalize_embeddings=True, batch_size=32)
)
return [e.tolist() for e in embeddings]
class OllamaEmbeddingBackend(EmbeddingBackend):
"""Ollama embedding backend."""
def __init__(
self,
model_name: str = "nomic-embed-text",
host: str = "http://localhost:11434",
):
self.model_name = model_name
self.host = host
self._dimensions = 768 # Default for nomic
@property
def dimensions(self) -> int:
return self._dimensions
async def embed(self, text: str) -> list[float]:
try:
import ollama
client = ollama.AsyncClient(host=self.host)
response = await client.embeddings(model=self.model_name, prompt=text)
embedding = response.get("embedding", [])
self._dimensions = len(embedding)
return embedding
except Exception as e:
logger.error(f"Ollama embedding error: {e}")
raise
async def embed_batch(self, texts: list[str]) -> list[list[float]]:
# Ollama doesn't have native batch support, process sequentially
embeddings = []
for text in texts:
emb = await self.embed(text)
embeddings.append(emb)
return embeddings
class EmbeddingManager:
"""
Manages embedding generation with caching and fallbacks.
Features:
- LRU cache for embeddings
- Multiple backend support
- Fallback chain
- Batch optimization
"""
def __init__(
self,
backend: str = "sentence_transformers",
model_name: Optional[str] = None,
cache_size: int = 10000,
):
self.cache: dict[str, list[float]] = {}
self.cache_size = cache_size
self.cache_order: list[str] = []
# Initialize backend
if backend == "sentence_transformers":
self.backend = SentenceTransformerBackend(
model_name or "all-MiniLM-L6-v2"
)
elif backend == "ollama":
self.backend = OllamaEmbeddingBackend(
model_name or "nomic-embed-text"
)
else:
raise ValueError(f"Unknown backend: {backend}")
self._stats = {
"total_requests": 0,
"cache_hits": 0,
"cache_misses": 0,
}
@property
def dimensions(self) -> int:
return self.backend.dimensions
def _cache_key(self, text: str) -> str:
"""Generate cache key for text."""
return hashlib.md5(text.encode()).hexdigest()
def _add_to_cache(self, key: str, embedding: list[float]):
"""Add to cache with LRU eviction."""
if key in self.cache:
# Move to end (most recent)
self.cache_order.remove(key)
self.cache_order.append(key)
return
# Evict if full
while len(self.cache) >= self.cache_size:
oldest = self.cache_order.pop(0)
del self.cache[oldest]
self.cache[key] = embedding
self.cache_order.append(key)
async def embed(self, text: str, use_cache: bool = True) -> EmbeddingResult:
"""Get embedding for text."""
self._stats["total_requests"] += 1
if use_cache:
key = self._cache_key(text)
if key in self.cache:
self._stats["cache_hits"] += 1
return EmbeddingResult(
embedding=self.cache[key],
model=self.backend.model_name if hasattr(self.backend, 'model_name') else "unknown",
dimensions=len(self.cache[key]),
cached=True,
)
self._stats["cache_misses"] += 1
embedding = await self.backend.embed(text)
if use_cache:
self._add_to_cache(key, embedding)
return EmbeddingResult(
embedding=embedding,
model=self.backend.model_name if hasattr(self.backend, 'model_name') else "unknown",
dimensions=len(embedding),
cached=False,
)
async def embed_batch(
self,
texts: list[str],
use_cache: bool = True,
) -> list[EmbeddingResult]:
"""Get embeddings for multiple texts."""
results = []
uncached_texts = []
uncached_indices = []
# Check cache first
for i, text in enumerate(texts):
if use_cache:
key = self._cache_key(text)
if key in self.cache:
self._stats["cache_hits"] += 1
results.append(EmbeddingResult(
embedding=self.cache[key],
model=self.backend.model_name if hasattr(self.backend, 'model_name') else "unknown",
dimensions=len(self.cache[key]),
cached=True,
))
else:
uncached_texts.append(text)
uncached_indices.append(i)
results.append(None) # Placeholder
else:
uncached_texts.append(text)
uncached_indices.append(i)
results.append(None)
# Batch embed uncached texts
if uncached_texts:
self._stats["cache_misses"] += len(uncached_texts)
embeddings = await self.backend.embed_batch(uncached_texts)
for idx, (text, embedding) in zip(uncached_indices, zip(uncached_texts, embeddings)):
if use_cache:
self._add_to_cache(self._cache_key(text), embedding)
results[idx] = EmbeddingResult(
embedding=embedding,
model=self.backend.model_name if hasattr(self.backend, 'model_name') else "unknown",
dimensions=len(embedding),
cached=False,
)
self._stats["total_requests"] += len(texts)
return results
def get_embedding_sync(self, text: str) -> list[float]:
"""Synchronous embedding (for compatibility)."""
loop = asyncio.get_event_loop()
if loop.is_running():
# Create new loop in thread
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as pool:
future = pool.submit(asyncio.run, self.embed(text))
result = future.result()
return result.embedding
else:
result = loop.run_until_complete(self.embed(text))
return result.embedding
def clear_cache(self):
"""Clear the embedding cache."""
self.cache.clear()
self.cache_order.clear()
def get_stats(self) -> dict:
"""Get embedding statistics."""
return {
**self._stats,
"cache_size": len(self.cache),
"cache_max": self.cache_size,
"cache_hit_rate": (
self._stats["cache_hits"] / max(1, self._stats["total_requests"])
),
"dimensions": self.dimensions,
}
