mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
jope-bm
deef89a724
Signed-off-by: Joe P <joe@basicmemory.com> Signed-off-by: phernandez <paul@basicmachines.co> Signed-off-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Paul Hernandez <60959+phernandez@users.noreply.github.com> Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Paul Hernandez <phernandez@users.noreply.github.com> Co-authored-by: phernandez <paul@basicmachines.co> Co-authored-by: jope-bm <jope-bm@users.noreply.github.com>
100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
"""OpenAI-based embedding provider for cloud or API-backed semantic indexing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from typing import Any
|
|
|
|
from basic_memory.repository.embedding_provider import EmbeddingProvider
|
|
from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
|
|
|
|
|
|
class OpenAIEmbeddingProvider(EmbeddingProvider):
|
|
"""Embedding provider backed by OpenAI's embeddings API."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str = "text-embedding-3-small",
|
|
*,
|
|
batch_size: int = 64,
|
|
dimensions: int = 1536,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
timeout: float = 30.0,
|
|
) -> None:
|
|
self.model_name = model_name
|
|
self.dimensions = dimensions
|
|
self.batch_size = batch_size
|
|
self._api_key = api_key
|
|
self._base_url = base_url
|
|
self._timeout = timeout
|
|
self._client: Any | None = None
|
|
self._client_lock = asyncio.Lock()
|
|
|
|
async def _get_client(self) -> Any:
|
|
if self._client is not None:
|
|
return self._client
|
|
|
|
async with self._client_lock:
|
|
if self._client is not None:
|
|
return self._client
|
|
|
|
try:
|
|
from openai import AsyncOpenAI # type: ignore[import-not-found]
|
|
except ImportError as exc: # pragma: no cover - covered via monkeypatch tests
|
|
raise SemanticDependenciesMissingError(
|
|
"OpenAI dependency is missing. "
|
|
"Install/update basic-memory to include semantic dependencies: "
|
|
"pip install -U basic-memory"
|
|
) from exc
|
|
|
|
api_key = self._api_key or os.getenv("OPENAI_API_KEY")
|
|
if not api_key:
|
|
raise SemanticDependenciesMissingError(
|
|
"OpenAI embedding provider requires OPENAI_API_KEY."
|
|
)
|
|
|
|
self._client = AsyncOpenAI(
|
|
api_key=api_key,
|
|
base_url=self._base_url,
|
|
timeout=self._timeout,
|
|
)
|
|
return self._client
|
|
|
|
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
if not texts:
|
|
return []
|
|
|
|
client = await self._get_client()
|
|
all_vectors: list[list[float]] = []
|
|
|
|
for start in range(0, len(texts), self.batch_size):
|
|
batch = texts[start : start + self.batch_size]
|
|
response = await client.embeddings.create(
|
|
model=self.model_name,
|
|
input=batch,
|
|
)
|
|
vectors_by_index: dict[int, list[float]] = {
|
|
int(item.index): [float(value) for value in item.embedding]
|
|
for item in response.data
|
|
}
|
|
for index in range(len(batch)):
|
|
vector = vectors_by_index.get(index)
|
|
if vector is None:
|
|
raise RuntimeError(
|
|
"OpenAI embedding response is missing expected vector index."
|
|
)
|
|
all_vectors.append(vector)
|
|
|
|
if all_vectors and len(all_vectors[0]) != self.dimensions:
|
|
raise RuntimeError(
|
|
f"Embedding model returned {len(all_vectors[0])}-dimensional vectors "
|
|
f"but provider was configured for {self.dimensions} dimensions."
|
|
)
|
|
return all_vectors
|
|
|
|
async def embed_query(self, text: str) -> list[float]:
|
|
vectors = await self.embed_documents([text])
|
|
return vectors[0] if vectors else [0.0] * self.dimensions
|