feat: add Ollama embedder with batching and retries
This commit is contained in:
21
tests/test_embedder.py
Normal file
21
tests/test_embedder.py
Normal 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
|
||||
Reference in New Issue
Block a user