From bd82e1f66d184bc13869030f03783cb4cb71da18 Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Wed, 4 Dec 2024 13:09:38 -0800 Subject: [PATCH] feat(python): add support for Azure OpenAPI SDK (#1906) Closes #1699 --- .../text_embedding_functions/openai_embedding.md | 1 + python/python/lancedb/embeddings/openai.py | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/docs/src/embeddings/available_embedding_models/text_embedding_functions/openai_embedding.md b/docs/src/embeddings/available_embedding_models/text_embedding_functions/openai_embedding.md index 87fd28f1..f08253c0 100644 --- a/docs/src/embeddings/available_embedding_models/text_embedding_functions/openai_embedding.md +++ b/docs/src/embeddings/available_embedding_models/text_embedding_functions/openai_embedding.md @@ -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 diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py index d126e2a2..a061d00a 100644 --- a/python/python/lancedb/embeddings/openai.py +++ b/python/python/lancedb/embeddings/openai.py @@ -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)