Skip to content

Commit

Permalink
BUG: Limit Azure OpenAI embeddings chunk size (langchain-ai#13425)
Browse files Browse the repository at this point in the history
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: langchain-ai#10707
* Porting this fix to the newest class (`AzureOpenAIEmbeddings`) for
openai >= 1.0

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

@baskaryan @hwchase17
  • Loading branch information
mspronesti authored and amiaxys committed Nov 23, 2023
1 parent 8975dfb commit 7368249
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 1 deletion.
4 changes: 4 additions & 0 deletions libs/langchain/langchain/embeddings/azure_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,10 @@ def validate_environment(cls, values: Dict) -> Dict:
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

Expand Down
2 changes: 1 addition & 1 deletion libs/langchain/langchain/embeddings/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,7 @@ def validate_environment(cls, values: Dict) -> Dict:
# 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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down

0 comments on commit 7368249

Please sign in to comment.