feat: add Ollama embedder with batching and retries

This commit is contained in:
2026-04-13 14:11:48 -04:00
parent 95687fad2e
commit 6e7f347bd3
2 changed files with 65 additions and 0 deletions

21
tests/test_embedder.py Normal file
View File

@@ -0,0 +1,21 @@
import pytest
import respx
from httpx import Response
from companion.rag.embedder import OllamaEmbedder
@respx.mock
def test_embed_batch():
route = respx.post("http://localhost:11434/api/embed").mock(
return_value=Response(200, json={"embeddings": [[0.1] * 1024, [0.2] * 1024]})
)
embedder = OllamaEmbedder(
base_url="http://localhost:11434", model="mxbai-embed-large", batch_size=2
)
result = embedder.embed(["hello world", "goodbye world"])
assert len(result) == 2
assert len(result[0]) == 1024
assert result[0][0] == 0.1
assert result[1][0] == 0.2
assert route.called