Coverage for node / src / stigmem_node / embedding / openai_adapter.py: 53%

1"""Cloud opt-in embedding adapter via OpenAI API (design memo §2). 

2 

3Activated by setting ``STIGMEM_EMBED_MODEL_PROVIDER=openai``. 

4Requires the OpenAI Python SDK:: 

5 

6 pip install 'stigmem-node[embed-openai]' 

7 

8The API key is read from the environment variable named by 

9``STIGMEM_EMBED_OPENAI_API_KEY_ENV`` (default: ``OPENAI_API_KEY``). 

10""" 

11 

12from __future__ import annotations 

13 

14import logging 

15import os 

16 

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

18 

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

20 

21_DEFAULT_MODEL = "text-embedding-3-small" 

22_DEFAULT_DIM = 1536 

23 

24 

25class OpenAIEmbeddingModel(EmbeddingModel): 

26 """Embed via OpenAI's embeddings API.""" 

27 

28 def __init__( 

29 self, 

30 model_id: str = _DEFAULT_MODEL, 

31 api_key_env: str = "OPENAI_API_KEY", 

32 dimension: int = _DEFAULT_DIM, 

33 ) -> None: 

34 self._model_id = model_id 

35 self._api_key_env = api_key_env 

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 openai 

49 except ImportError as exc: 

50 raise EmbeddingError( 

51 "openai package is required for the OpenAI embedding adapter. " 

52 "Install it with: pip install 'stigmem-node[embed-openai]'" 

53 ) from exc 

54 

55 api_key = os.environ.get(self._api_key_env) 

56 if not api_key: 

57 logger.debug("OpenAI embedding API key environment variable is unset") 

58 raise EmbeddingError( 

59 "OpenAI embedding credentials are not configured. " 

60 "Set the configured API key environment variable or use a " 

61 "different embed_model_provider." 

62 ) 

63 

64 try: 

65 client = openai.OpenAI(api_key=api_key) 

66 response = client.embeddings.create(model=self._model_id, input=texts) 

67 vecs = [item.embedding for item in response.data] 

68 except Exception as exc: 

69 raise EmbeddingError(f"OpenAI embedding failed: {exc}") from exc 

70 

71 return [l2_normalize(v) for v in vecs]