Coverage for node / src / stigmem_node / embedding / local_adapter.py: 94%

1"""Local embedding adapter via Ollama HTTP API (design memo §2). 

2 

3Default model: ``nomic-embed-text-v1.5`` (768-dim, Apache-2.0, Matryoshka). 

4Requires a running Ollama instance; install the model with:: 

5 

6 ollama pull nomic-embed-text 

7 

8Install the optional dependency:: 

9 

10 pip install 'stigmem-node[embed-local]' # includes httpx 

11""" 

12 

13from __future__ import annotations 

14 

15import logging 

16 

17from .base import EmbeddingError, EmbeddingModel, Vector, l2_normalize 

18 

19logger = logging.getLogger("stigmem.embed.local") 

20 

21_DEFAULT_MODEL = "nomic-embed-text-v1.5" 

22_DEFAULT_OLLAMA_URL = "http://localhost:11434" 

23 

24 

25class OllamaEmbeddingModel(EmbeddingModel): 

26 """Embed via Ollama's ``/api/embeddings`` endpoint.""" 

27 

28 def __init__( 

29 self, 

30 model_id: str = _DEFAULT_MODEL, 

31 ollama_url: str = _DEFAULT_OLLAMA_URL, 

32 dimension: int = 768, 

33 ) -> None: 

34 self._model_id = model_id 

35 self._ollama_url = ollama_url.rstrip("/") 

36 self._dimension = dimension 

37 

38 @property 

39 def model_id(self) -> str: 

40 return self._model_id 

41 

42 @property 

43 def dimension(self) -> int: 

44 return self._dimension 

45 

46 def embed(self, texts: list[str]) -> list[Vector]: 

47 try: 

48 import httpx 

49 except ImportError as exc: 

50 raise EmbeddingError( 

51 "httpx is required for the Ollama embedding adapter. " 

52 "It is already a core stigmem-node dependency." 

53 ) from exc 

54 

55 results: list[Vector] = [] 

56 url = f"{self._ollama_url}/api/embeddings" 

57 for text in texts: 

58 try: 

59 resp = httpx.post( 

60 url, 

61 json={"model": self._model_id, "prompt": text}, 

62 timeout=30.0, 

63 ) 

64 resp.raise_for_status() 

65 vec = resp.json()["embedding"] 

66 except Exception as exc: 

67 raise EmbeddingError(f"Ollama embedding failed: {exc}") from exc 

68 results.append(l2_normalize(vec)) 

69 return results