feat(embedding): add dimensions_mode for OpenAI dimensions= forwarding
Add EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE (auto|always|never) controlling
whether the dimensions= parameter is forwarded on OpenAI embeddings.create
calls. auto (default) sends it when the operator explicitly set
EMBEDDING_VECTOR_DIMENSIONS and the configured model is not on the
known-rejecting allowlist (currently text-embedding-ada-002).
The provenance check (was VECTOR_DIMENSIONS explicitly set?) lives as
EmbeddingSettings.resolve_send_dimensions() because it needs access to
model_fields_set, which the standalone resolver does not have. The
resolved boolean is passed into _EmbeddingClient at construction time;
the client never inspects mode or provenance.
Also pins cloudevents <2.0 — 2.0.0 reorganized the package and dropped
cloudevents.conversion and cloudevents.http, which src/telemetry/emitter.py
imports. The original `>=1.12.0` constraint allowed the broken 2.0 resolve.
With the pin, the imports resolve cleanly and the basedpyright warning
cascade (37+ warnings about unknown types) disappears.
Drive-by cleanups (all unnecessary cast/ignore comments flagged by
basedpyright after the cloudevents downgrade):
- vector_store/lancedb.py, tests/conftest.py, and
tests/deriver/test_vector_reconciliation.py — drop dead pyright ignores
- sdks/python/src/honcho/http/{async_,}client.py — drop unnecessary
cast(datetime, ...) (parsedate_to_datetime already returns datetime)
- vector_store/turbopuffer.py — cast(Any, rows) for the upsert_rows
TypedDict that the SDK exposes but our row builder doesn't satisfy
- tests/test_datetime_parsing.py — ignore reportArgumentType on the
test that deliberately passes wrong types to assert raises
This commit is contained in:
parent
5de8a3b81a
commit
29dc1e138c
|
|
@ -38,7 +38,7 @@ dependencies = [
|
|||
"cashews[redis]==7.4.4",
|
||||
"scikit-learn>=1.6.0",
|
||||
"prometheus_client>=0.21.0",
|
||||
"cloudevents>=1.12.0",
|
||||
"cloudevents>=1.12.0,<2.0",
|
||||
]
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
|
|
|
|||
|
|
@ -370,7 +370,7 @@ class AsyncHonchoHTTPClient:
|
|||
from datetime import datetime
|
||||
from email.utils import parsedate_to_datetime
|
||||
|
||||
dt: datetime = cast(datetime, parsedate_to_datetime(header))
|
||||
dt: datetime = parsedate_to_datetime(header)
|
||||
timestamp: float = dt.timestamp()
|
||||
return max(0.0, timestamp - time.time())
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -367,7 +367,7 @@ class HonchoHTTPClient:
|
|||
from datetime import datetime
|
||||
from email.utils import parsedate_to_datetime
|
||||
|
||||
dt: datetime = cast(datetime, parsedate_to_datetime(header))
|
||||
dt: datetime = parsedate_to_datetime(header)
|
||||
timestamp: float = dt.timestamp()
|
||||
return max(0.0, timestamp - time.time())
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,12 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
ModelTransport = Literal["anthropic", "openai", "gemini"]
|
||||
EmbeddingTransport = Literal["openai", "gemini"]
|
||||
EmbeddingDimensionsMode = Literal["auto", "always", "never"]
|
||||
|
||||
# OpenAI-compatible models that reject the `dimensions=` request parameter.
|
||||
_EMBEDDING_KNOWN_REJECTING_MODELS: frozenset[str] = frozenset(
|
||||
{"text-embedding-ada-002"}
|
||||
)
|
||||
|
||||
|
||||
def _default_embedding_model_for_transport(transport: EmbeddingTransport) -> str:
|
||||
|
|
@ -294,6 +300,7 @@ class ConfiguredEmbeddingModelSettings(BaseModel):
|
|||
model: str = "text-embedding-3-small"
|
||||
transport: EmbeddingTransport = "openai"
|
||||
overrides: ModelOverrideSettings = Field(default_factory=ModelOverrideSettings)
|
||||
dimensions_mode: EmbeddingDimensionsMode = "auto"
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
@ -695,6 +702,23 @@ class EmbeddingSettings(HonchoSettings):
|
|||
)
|
||||
return data # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
def resolve_send_dimensions(self) -> bool:
|
||||
"""Decide whether OpenAI embedding calls should forward ``dimensions=``.
|
||||
|
||||
Lives on the settings instance because ``auto`` mode needs access to
|
||||
``self.model_fields_set`` to tell whether the operator explicitly set
|
||||
``VECTOR_DIMENSIONS`` — a standalone resolver over
|
||||
``ConfiguredEmbeddingModelSettings`` cannot see that.
|
||||
"""
|
||||
mode = self.MODEL_CONFIG.dimensions_mode
|
||||
if mode == "always":
|
||||
return True
|
||||
if mode == "never":
|
||||
return False
|
||||
if self.MODEL_CONFIG.model in _EMBEDDING_KNOWN_REJECTING_MODELS:
|
||||
return False
|
||||
return "VECTOR_DIMENSIONS" in self.model_fields_set
|
||||
|
||||
|
||||
class DeriverSettings(HonchoSettings):
|
||||
model_config = SettingsConfigDict( # pyright: ignore
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
import logging
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from typing import NamedTuple
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import tiktoken
|
||||
from google import genai
|
||||
|
|
@ -34,10 +34,12 @@ class _EmbeddingClient:
|
|||
vector_dimensions: int,
|
||||
max_input_tokens: int,
|
||||
max_tokens_per_request: int,
|
||||
send_dimensions: bool,
|
||||
):
|
||||
self.transport: str = config.transport
|
||||
self.model: str = config.model
|
||||
self.vector_dimensions: int = vector_dimensions
|
||||
self.send_dimensions: bool = send_dimensions
|
||||
|
||||
if self.transport == "gemini":
|
||||
if not config.api_key:
|
||||
|
|
@ -101,9 +103,10 @@ class _EmbeddingClient:
|
|||
raise ValueError("No embedding returned from Gemini API")
|
||||
return self._validate_embedding_dimensions(response.embeddings[0].values)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
model=self.model, input=[query]
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {"model": self.model, "input": [query]}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
return self._validate_embedding_dimensions(response.data[0].embedding)
|
||||
|
||||
async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
|
||||
|
|
@ -138,10 +141,13 @@ class _EmbeddingClient:
|
|||
self._validate_embedding_dimensions(emb.values)
|
||||
)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
input=batch,
|
||||
model=self.model,
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {
|
||||
"input": batch,
|
||||
"model": self.model,
|
||||
}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
embeddings.extend(
|
||||
[
|
||||
self._validate_embedding_dimensions(data.embedding)
|
||||
|
|
@ -287,9 +293,13 @@ class _EmbeddingClient:
|
|||
)
|
||||
)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
model=self.model, input=[item.text for item in batch]
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"input": [item.text for item in batch],
|
||||
}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
for item, embedding_data in zip(batch, response.data, strict=True):
|
||||
result[item.text_id][item.chunk_index] = (
|
||||
self._validate_embedding_dimensions(
|
||||
|
|
@ -411,6 +421,7 @@ class EmbeddingClient:
|
|||
vector_dimensions=settings.EMBEDDING.VECTOR_DIMENSIONS,
|
||||
max_input_tokens=settings.EMBEDDING.MAX_INPUT_TOKENS,
|
||||
max_tokens_per_request=settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
|
||||
send_dimensions=settings.EMBEDDING.resolve_send_dimensions(),
|
||||
)
|
||||
self._instance_signature = signature
|
||||
logger.debug(
|
||||
|
|
@ -434,6 +445,7 @@ class EmbeddingClient:
|
|||
settings.EMBEDDING.VECTOR_DIMENSIONS,
|
||||
settings.EMBEDDING.MAX_INPUT_TOKENS,
|
||||
settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
|
||||
settings.EMBEDDING.resolve_send_dimensions(),
|
||||
)
|
||||
|
||||
async def embed(self, query: str) -> list[float]:
|
||||
|
|
|
|||
|
|
@ -99,7 +99,7 @@ class LanceDBVectorStore(VectorStore):
|
|||
fields.extend(self._metadata_fields_for_namespace(namespace))
|
||||
schema = pa.schema(fields)
|
||||
try:
|
||||
table = await db.create_table(namespace, schema=schema) # pyright: ignore[reportUnknownArgumentType]
|
||||
table = await db.create_table(namespace, schema=schema)
|
||||
return table
|
||||
except Exception:
|
||||
# Table may have been created by another worker, try to open it
|
||||
|
|
|
|||
|
|
@ -87,7 +87,7 @@ class TurbopufferVectorStore(VectorStore):
|
|||
|
||||
try:
|
||||
await ns.write(
|
||||
upsert_rows=rows,
|
||||
upsert_rows=cast(Any, rows),
|
||||
distance_metric=DISTANCE_METRIC,
|
||||
)
|
||||
return
|
||||
|
|
|
|||
|
|
@ -509,7 +509,7 @@ def mock_vector_store(request: pytest.FixtureRequest):
|
|||
from src.vector_store import (
|
||||
VectorQueryResult,
|
||||
VectorRecord,
|
||||
_hash_namespace_components, # pyright: ignore[reportPrivateUsage]
|
||||
_hash_namespace_components,
|
||||
)
|
||||
|
||||
# Create a mock vector store that stores vectors in memory
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ from src.reconciler.sync_vectors import (
|
|||
from src.vector_store import (
|
||||
VectorRecord,
|
||||
VectorStore,
|
||||
_hash_namespace_components, # pyright: ignore[reportPrivateUsage]
|
||||
_hash_namespace_components,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,8 +12,16 @@ class FakeOpenAIEmbeddingsAPI:
|
|||
self.embedding: list[float] = embedding
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def create(self, *, model: str, input: str | list[str]) -> SimpleNamespace:
|
||||
self.calls.append({"model": model, "input": input})
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[str],
|
||||
**kwargs: Any,
|
||||
) -> SimpleNamespace:
|
||||
call: dict[str, Any] = {"model": model, "input": input}
|
||||
call.update(kwargs)
|
||||
self.calls.append(call)
|
||||
if isinstance(input, list):
|
||||
data = [SimpleNamespace(embedding=self.embedding) for _ in input]
|
||||
else:
|
||||
|
|
@ -45,6 +53,7 @@ async def test_openai_embedding_client_uses_configured_model_and_dimensions(
|
|||
vector_dimensions=8,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
embedding = await client.embed("hello world")
|
||||
|
|
@ -76,6 +85,7 @@ async def test_openai_embedding_client_rejects_dimension_mismatch(
|
|||
vector_dimensions=8,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Embedding dimension mismatch"):
|
||||
|
|
@ -125,6 +135,7 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
vector_dimensions=12,
|
||||
max_input_tokens=4096,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
embedding = await client.embed("hello world")
|
||||
|
|
@ -137,3 +148,194 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
"config": {"output_dimensionality": 12},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _build_openai_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
embedding: list[float],
|
||||
model: str,
|
||||
send_dimensions: bool,
|
||||
vector_dimensions: int,
|
||||
) -> tuple[_EmbeddingClient, FakeOpenAIEmbeddingsAPI]:
|
||||
fake_embeddings = FakeOpenAIEmbeddingsAPI(embedding)
|
||||
|
||||
class FakeOpenAIClient:
|
||||
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
|
||||
self.api_key: str | None = api_key
|
||||
self.base_url: str | None = base_url
|
||||
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
|
||||
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="openai",
|
||||
model=model,
|
||||
api_key="test-key",
|
||||
),
|
||||
vector_dimensions=vector_dimensions,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=send_dimensions,
|
||||
)
|
||||
return client, fake_embeddings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_embed_forwards_dimensions_when_send_dimensions_true(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.embed("hello")
|
||||
|
||||
assert fake.calls == [
|
||||
{
|
||||
"model": "text-embedding-3-small",
|
||||
"input": ["hello"],
|
||||
"dimensions": 768,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_embed_omits_dimensions_when_send_dimensions_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 1536,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=False,
|
||||
vector_dimensions=1536,
|
||||
)
|
||||
|
||||
await client.embed("hello")
|
||||
|
||||
assert fake.calls == [{"model": "text-embedding-3-small", "input": ["hello"]}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_simple_batch_embed_forwards_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.simple_batch_embed(["a", "b"])
|
||||
|
||||
assert len(fake.calls) == 1
|
||||
assert fake.calls[0]["dimensions"] == 768
|
||||
assert fake.calls[0]["input"] == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_batch_embed_forwards_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.batch_embed({"a": "hello", "b": "world"})
|
||||
|
||||
assert len(fake.calls) == 1
|
||||
assert fake.calls[0]["dimensions"] == 768
|
||||
|
||||
|
||||
def _build_embedding_settings(
|
||||
env: dict[str, str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> Any:
|
||||
"""Construct a fresh EmbeddingSettings from the given env, isolated from os.environ."""
|
||||
from src.config import EmbeddingSettings
|
||||
|
||||
for key in (
|
||||
"EMBEDDING_VECTOR_DIMENSIONS",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL",
|
||||
"EMBEDDING_MODEL_CONFIG__TRANSPORT",
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
for key, value in env.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
return EmbeddingSettings()
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_default_dim_returns_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings({}, monkeypatch)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_explicit_dim_returns_true(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings({"EMBEDDING_VECTOR_DIMENSIONS": "768"}, monkeypatch)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_ada_002_returns_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "1536",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL": "text-embedding-ada-002",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_always_returns_true_regardless(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "always"},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_always_overrides_ada_rejecting_allowlist(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "always",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL": "text-embedding-ada-002",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_never_returns_false_regardless(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "never",
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "768",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
|
|
|||
|
|
@ -497,7 +497,7 @@ class TestErrorHandlingAndRecovery:
|
|||
# parse_datetime_iso should raise appropriate errors
|
||||
if error_case is None or not isinstance(error_case, str):
|
||||
with pytest.raises(ValueError):
|
||||
parse_datetime_iso(error_case)
|
||||
parse_datetime_iso(error_case) # pyright: ignore[reportArgumentType]
|
||||
|
||||
# Filter validation should return None for invalid types
|
||||
if isinstance(error_case, str) or error_case is None:
|
||||
|
|
|
|||
Loading…
Reference in New Issue