A wrapper around AsyncOpenAI that provides a unified interface for OpenAI embeddings. This class is specifically designed for embedding models like text-embedding-3-small, text-embedding-3-large, etc.
| 21 | |
| 22 | |
| 23 | class EmbeddingOpenAI(BaseModel): |
| 24 | """ |
| 25 | A wrapper around AsyncOpenAI that provides a unified interface for OpenAI embeddings. |
| 26 | |
| 27 | This class is specifically designed for embedding models like text-embedding-3-small, text-embedding-3-large, etc. |
| 28 | """ |
| 29 | |
| 30 | model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow") |
| 31 | |
| 32 | # Model configuration |
| 33 | model: Union[ChatModel, str] = "text-embedding-3-small" |
| 34 | |
| 35 | # Client initialization parameters |
| 36 | api_key: Optional[str] = None |
| 37 | organization: Optional[str] = None |
| 38 | project: Optional[str] = None |
| 39 | base_url: Optional[Union[str, httpx.URL]] = None |
| 40 | websocket_base_url: Optional[Union[str, httpx.URL]] = None |
| 41 | timeout: Optional[Union[float, httpx.Timeout]] = None |
| 42 | max_retries: int = 5 |
| 43 | default_headers: Optional[Mapping[str, str]] = None |
| 44 | default_query: Optional[Mapping[str, object]] = None |
| 45 | http_client: Optional[httpx.AsyncClient] = None |
| 46 | _strict_response_validation: bool = False |
| 47 | |
| 48 | # Embedding parameters |
| 49 | dimensions: Optional[int] = None # Optional dimensions for text-embedding-3 models |
| 50 | encoding_format: Optional[str] = None # "float" or "base64" |
| 51 | |
| 52 | @property |
| 53 | def provider(self) -> str: |
| 54 | return 'openai' |
| 55 | |
| 56 | def _get_client_params(self) -> dict[str, Any]: |
| 57 | """Prepare client parameters dictionary.""" |
| 58 | base_params = { |
| 59 | 'api_key': self.api_key, |
| 60 | 'organization': self.organization, |
| 61 | 'project': self.project, |
| 62 | 'base_url': self.base_url, |
| 63 | 'websocket_base_url': self.websocket_base_url, |
| 64 | 'timeout': self.timeout, |
| 65 | 'max_retries': self.max_retries, |
| 66 | 'default_headers': self.default_headers, |
| 67 | 'default_query': self.default_query, |
| 68 | '_strict_response_validation': self._strict_response_validation, |
| 69 | } |
| 70 | |
| 71 | # Create client_params dict with non-None values |
| 72 | client_params = {k: v for k, v in base_params.items() if v is not None} |
| 73 | |
| 74 | # Add http_client if provided |
| 75 | if self.http_client is not None: |
| 76 | client_params['http_client'] = self.http_client |
| 77 | |
| 78 | return client_params |
| 79 | |
| 80 | def get_client(self) -> AsyncOpenAI: |