feat: add Ollama embedder with batching and retries
This commit is contained in:
44
src/companion/rag/embedder.py
Normal file
44
src/companion/rag/embedder.py
Normal file
@@ -0,0 +1,44 @@
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class OllamaEmbedder:
|
||||
def __init__(self, base_url: str, model: str, batch_size: int):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.model = model
|
||||
self.batch_size = batch_size
|
||||
|
||||
def embed(
|
||||
self, texts: List[str], retries: int = 3, backoff: float = 1.0
|
||||
) -> List[List[float]]:
|
||||
all_embeddings: List[List[float]] = []
|
||||
url = f"{self.base_url}/api/embed"
|
||||
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
last_exception: Exception | None = None
|
||||
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
response = client.post(
|
||||
url,
|
||||
json={"model": self.model, "input": batch},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
embeddings = data["embeddings"]
|
||||
all_embeddings.extend(embeddings)
|
||||
break
|
||||
except Exception as exc:
|
||||
last_exception = exc
|
||||
if attempt < retries - 1:
|
||||
time.sleep(backoff * (2**attempt))
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Failed to embed batch after {retries} retries"
|
||||
) from last_exception
|
||||
|
||||
return all_embeddings
|
||||
Reference in New Issue
Block a user