feat(python): add support for Azure OpenAPI SDK (#1906)

Closes #1699
This commit is contained in:
Lei Xu
2024-12-04 13:09:38 -08:00
committed by GitHub
parent ba34c3bee1
commit bd82e1f66d
2 changed files with 9 additions and 1 deletions

View File

@@ -6,6 +6,7 @@ LanceDB registers the OpenAI embeddings function in the registry by default, as
|---|---|---|---|
| `name` | `str` | `"text-embedding-ada-002"` | The name of the model. |
| `dim` | `int` | Model default | For OpenAI's newer text-embedding-3 model, we can specify a dimensionality that is smaller than the 1536 size. This feature supports it |
| `use_azure` | bool | `False` | Set true to use Azure OpenAPI SDK |
```python

View File

@@ -48,6 +48,9 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
organization: Optional[str] = None
api_key: Optional[str] = None
# Set true to use Azure OpenAI API
use_azure: bool = False
def ndims(self):
return self._ndims
@@ -123,4 +126,8 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
kwargs["organization"] = self.organization
if self.api_key:
kwargs["api_key"] = self.api_key
return openai.OpenAI(**kwargs)
if self.use_azure:
return openai.AzureOpenAI(**kwargs)
else:
return openai.OpenAI(**kwargs)