Skip to content

Honor dimensions param #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jnu merged 2 commits into from
Apr 28, 2026
Merged

Honor dimensions param #131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4ef63bd
honor dimensions param
jnu Apr 28, 2026
c1fe98a
add test for embedding dimensions
jnu Apr 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion bc2/lib/embedding/openai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,13 @@ def __init__(

def embed(self, text: str) -> Embedding:
text = self._trim_input(text)
result = self.client.embeddings.create(input=text, model=self.config.model)
params = {
"input": text,
"model": self.config.model,
}
if self.config.dimensions:
params["dimensions"] = self.config.dimensions
result = self.client.embeddings.create(**params)
Comment on lines +111 to +116

@sentry sentry Bot Apr 28, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: The embed_async method in OpenAIEmbeddingGenerator does not pass the dimensions parameter to the OpenAI API, unlike its synchronous counterpart embed.
Severity: MEDIUM

Suggested Fix

Update the embed_async method to construct a params dictionary and pass the dimensions parameter if it's set in self.config, mirroring the logic in the synchronous embed method. Also, add a test case to verify the behavior of embed_async with the dimensions parameter.

Prompt for AI Agent 
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.

Location: bc2/lib/embedding/openai.py#L111-L116

Potential issue: The pull request updated the synchronous `embed` method to correctly
pass the `dimensions` parameter from `self.config` to the OpenAI API. However, the
asynchronous `embed_async` method was not updated. It still calls
`self.aclient.embeddings.create` without including the `dimensions` parameter. This will
cause users who configure a specific embedding dimension and use the async path to
receive embeddings of the model's default size (e.g., 3072) instead of the requested
size (e.g., 256), leading to a silent inconsistency and potential downstream errors.

Did we get this right? 👍 / 👎 to inform future reviews.

return self._format_result(result)

async def embed_async(self, text: str) -> Embedding:
Expand Down
63 changes: 62 additions & 1 deletion bc2/lib/embedding/test_openai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,16 @@
from unittest.mock import MagicMock

import pytest
from openai import AsyncOpenAI, OpenAI

from bc2.core.common.openai import OpenAIClientConfig
from bc2.core.common.openai_metadata import get_encoding_for_model

from .openai import OpenAIEmbeddingConfig, OpenAIEmbeddingGeneratorConfig
from .openai import (
OpenAIEmbeddingConfig,
OpenAIEmbeddingDriver,
OpenAIEmbeddingGeneratorConfig,
)


@pytest.mark.parametrize(
Expand Down Expand Up @@ -33,3 +40,57 @@ def test_trim_input(model: str, text: str, trimmed: str):

assert embed.driver._trim_input(text) == trimmed
assert len(encoding.encode(trimmed)) <= embed.generator.max_input_tokens


def _mock_embedding_response(model: str = "text-embedding-3-large") -> MagicMock:
response = MagicMock()
response.data = [MagicMock(embedding=[0.0, 0.1, 0.2])]
response.model = model
return response


@pytest.mark.parametrize(
"dimensions",
[None, 256, 1024],
ids=["unset", "256", "1024"],
)
def test_embed_passes_dimensions(dimensions: int | None):
client = MagicMock(spec=OpenAI)
client.embeddings.create.return_value = _mock_embedding_response()
aclient = MagicMock(spec=AsyncOpenAI)

config = OpenAIEmbeddingGeneratorConfig(
model="text-embedding-3-large",
dimensions=dimensions,
)
driver = OpenAIEmbeddingDriver(client, aclient, config)

driver.embed("hello")

client.embeddings.create.assert_called_once()
call_kwargs = client.embeddings.create.call_args.kwargs
assert call_kwargs["input"] == "hello"
assert call_kwargs["model"] == "text-embedding-3-large"
if dimensions is None:
assert "dimensions" not in call_kwargs
else:
assert call_kwargs["dimensions"] == dimensions


def test_embed_config_dimensions_default():
config = OpenAIEmbeddingConfig(
client=OpenAIClientConfig(api_key="test"),
generator=OpenAIEmbeddingGeneratorConfig(model="text-embedding-3-large"),
)
assert config.generator.dimensions is None


def test_embed_config_dimensions_override():
config = OpenAIEmbeddingConfig(
client=OpenAIClientConfig(api_key="test"),
generator=OpenAIEmbeddingGeneratorConfig(
model="text-embedding-3-large", dimensions=512
),
)
assert config.generator.dimensions == 512
assert config.generator.model_dimensions == 512
Loading