wehub-resource-sync
65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from opensquilla.memory.embedding import OllamaEmbeddingProvider
|
|
|
|
|
|
def _patch(monkeypatch: Any, captured: dict[str, Any], embeddings: list[list[float]]) -> None:
|
|
captured["calls"] = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
captured["calls"] += 1
|
|
captured["url"] = str(request.url)
|
|
captured["payload"] = json.loads(request.content.decode("utf-8"))
|
|
return httpx.Response(200, json={"model": "nomic-embed-text", "embeddings": embeddings})
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
real_async_client = httpx.AsyncClient
|
|
|
|
def patched_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient:
|
|
kwargs["transport"] = transport
|
|
return real_async_client(*args, **kwargs)
|
|
|
|
monkeypatch.setattr("opensquilla.memory.embedding.httpx.AsyncClient", patched_async_client)
|
|
|
|
|
|
def test_embed_query_uses_api_embed_endpoint(monkeypatch: Any) -> None:
|
|
captured: dict[str, Any] = {}
|
|
_patch(monkeypatch, captured, [[0.1, 0.2, 0.3]])
|
|
provider = OllamaEmbeddingProvider(model="nomic-embed-text")
|
|
|
|
vec = asyncio.run(provider.embed_query("hello"))
|
|
|
|
assert captured["url"] == "http://localhost:11434/api/embed"
|
|
assert captured["payload"] == {"model": "nomic-embed-text", "input": ["hello"]}
|
|
assert vec == [0.1, 0.2, 0.3]
|
|
|
|
|
|
def test_embed_batch_sends_single_request_with_all_inputs(monkeypatch: Any) -> None:
|
|
captured: dict[str, Any] = {}
|
|
_patch(monkeypatch, captured, [[1.0], [2.0], [3.0]])
|
|
provider = OllamaEmbeddingProvider(model="nomic-embed-text")
|
|
|
|
vectors = asyncio.run(provider.embed_batch(["a", "b", "c"]))
|
|
|
|
# One round trip for the whole batch, not one per text.
|
|
assert captured["calls"] == 1
|
|
assert captured["payload"]["input"] == ["a", "b", "c"]
|
|
assert vectors == [[1.0], [2.0], [3.0]]
|
|
|
|
|
|
def test_embed_batch_empty_makes_no_request(monkeypatch: Any) -> None:
|
|
captured: dict[str, Any] = {}
|
|
_patch(monkeypatch, captured, [])
|
|
provider = OllamaEmbeddingProvider(model="nomic-embed-text")
|
|
|
|
vectors = asyncio.run(provider.embed_batch([]))
|
|
|
|
assert vectors == []
|
|
assert captured["calls"] == 0
|