| 49 | |
| 50 | |
| 51 | class OpenAIEmbedding: |
| 52 | name = "openai" |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | model: str, |
| 57 | dimensions: int, |
| 58 | api_key: str, |
| 59 | base_url: Optional[str] = None, |
| 60 | timeout: float = 60.0, |
| 61 | ) -> None: |
| 62 | from openai import OpenAI # local import keeps dependency optional |
| 63 | |
| 64 | self.model = model |
| 65 | self.dimensions = dimensions |
| 66 | self._client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout) |
| 67 | |
| 68 | def embed(self, text: str) -> List[float]: |
| 69 | return self.embed_batch([text])[0] |
| 70 | |
| 71 | def embed_batch(self, texts: Iterable[str]) -> List[List[float]]: |
| 72 | inputs = [text or " " for text in texts] |
| 73 | kwargs: dict[str, object] = {"model": self.model, "input": inputs} |
| 74 | if self.model.startswith("text-embedding-3"): |
| 75 | kwargs["dimensions"] = self.dimensions |
| 76 | response = self._client.embeddings.create(**kwargs) |
| 77 | return [_resize(list(item.embedding), self.dimensions) for item in response.data] |
| 78 | |
| 79 | |
| 80 | class SentenceTransformersEmbedding: |
no outgoing calls
no test coverage detected