BUG: Limit Azure OpenAI embeddings chunk size (#13425)

Hi! 
This short PR aims at:
* Fixing `OpenAIEmbeddings`' check on `chunk_size` when used with Azure
OpenAI (thus with openai < 1.0). Azure OpenAI embeddings support at most
16 chunks per batch, I believe we are supposed to take the min between
the passed value/default value and 16, not the max - which, I suppose,
was introduced by accident while refactoring the previous version of
this check from this other PR of mine: #10707
* Porting this fix to the newest class (`AzureOpenAIEmbeddings`) for
openai >= 1.0

This fixes #13539 (closed but the issue persists).  

@baskaryan @hwchase17
This commit is contained in:
Massimiliano Pronesti 2023-11-20 03:34:51 +01:00 committed by GitHub
parent e53f59f01a
commit 6bf9b2cb51
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 18 additions and 1 deletions

View File

@ -83,6 +83,10 @@ class AzureOpenAIEmbeddings(OpenAIEmbeddings):
values["azure_ad_token"] = values["azure_ad_token"] or os.getenv(
"AZURE_OPENAI_AD_TOKEN"
)
# Azure OpenAI embedding models allow a maximum of 16 texts
# at a time in each batch
# See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#embeddings
values["chunk_size"] = min(values["chunk_size"], 16)
try:
import openai

View File

@ -291,7 +291,7 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
# Azure OpenAI embedding models allow a maximum of 16 texts
# at a time in each batch
# See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#embeddings
values["chunk_size"] = max(values["chunk_size"], 16)
values["chunk_size"] = min(values["chunk_size"], 16)
else:
default_api_version = ""
values["openai_api_version"] = get_from_dict_or_env(

View File

@ -30,12 +30,25 @@ def test_azure_openai_embedding_documents_multiple() -> None:
embedding = _get_embeddings(chunk_size=2)
embedding.embedding_ctx_length = 8191
output = embedding.embed_documents(documents)
assert embedding.chunk_size == 2
assert len(output) == 3
assert len(output[0]) == 1536
assert len(output[1]) == 1536
assert len(output[2]) == 1536
def test_azure_openai_embedding_documents_chunk_size() -> None:
"""Test openai embeddings."""
documents = ["foo bar"] * 20
embedding = _get_embeddings()
embedding.embedding_ctx_length = 8191
output = embedding.embed_documents(documents)
# Max 16 chunks per batch on Azure OpenAI embeddings
assert embedding.chunk_size == 16
assert len(output) == 20
assert all([len(out) == 1536 for out in output])
@pytest.mark.asyncio
async def test_azure_openai_embedding_documents_async_multiple() -> None:
"""Test openai embeddings."""