fix(core): scope LiteLLM dimension requests

Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2026-05-28 20:14:44 -05:00
parent 8270405a1c
commit a3739fa72e
3 changed files with 36 additions and 1 deletions
+4
View File
@@ -151,6 +151,10 @@ export COHERE_API_KEY=...
Basic Memory creates vector tables before the first embedding call, so non-default LiteLLM models must set `BASIC_MEMORY_SEMANTIC_EMBEDDING_DIMENSIONS`. The LiteLLM OpenAI default (`openai/text-embedding-3-small`) uses 1536 dimensions automatically.
For fixed-size LiteLLM models, dimensions are used as Basic Memory's local vector schema and
validation size. Basic Memory only sends dimensions as a provider-side output-size request for
`text-embedding-3` model strings, where LiteLLM/OpenAI support reduced output dimensions.
Some retrieval models are asymmetric: indexed passages and search queries must be embedded with different provider parameters. Basic Memory automatically sets LiteLLM `input_type` for known asymmetric model families:
- Cohere v3: documents use `search_document`, queries use `search_query`
@@ -46,6 +46,21 @@ def _default_input_types(model_name: str) -> tuple[str | None, str | None]:
return None, None
def _supports_dimension_parameter(model_name: str) -> bool:
"""Return whether configured dimensions should be sent to LiteLLM."""
normalized = model_name.strip().lower()
# Trigger: `dimensions` is both the Basic Memory vector schema size and a
# provider-side output-size request parameter in LiteLLM.
# Why: fixed-size models such as Cohere v3 still need the schema value for
# validation, but LiteLLM maps the request parameter to provider fields they
# reject. Only send it for model families that clearly support output-size
# control.
# Outcome: OpenAI/Azure text-embedding-3 reductions work, while fixed-size
# providers keep dimensions local to Basic Memory.
return "text-embedding-3" in normalized
def _import_litellm() -> Any:
"""Import LiteLLM without letting its import-time dotenv hook read cwd secrets."""
# Constraint: LiteLLM 1.85.0 loads .env files at import time when
@@ -119,10 +134,11 @@ class LiteLLMEmbeddingProvider(EmbeddingProvider):
params: dict[str, Any] = {
"model": self.model_name,
"input": batch,
"dimensions": self.dimensions,
"drop_params": True,
"timeout": self._timeout,
}
if _supports_dimension_parameter(self.model_name):
params["dimensions"] = self.dimensions
if self._api_key:
params["api_key"] = self._api_key
if input_type:
+15
View File
@@ -162,6 +162,21 @@ async def test_litellm_provider_uses_cohere_document_and_query_input_types(monke
assert calls[1]["input_type"] == "search_query"
@pytest.mark.asyncio
async def test_litellm_provider_does_not_forward_dimensions_to_cohere_v3(monkeypatch):
"""Cohere v3 uses configured dimensions only for schema validation."""
calls = _install_litellm_stub(monkeypatch, dim=1024)
provider = LiteLLMEmbeddingProvider(
model_name="cohere/embed-english-v3.0",
dimensions=1024,
)
await provider.embed_documents(["indexed passage"])
assert "dimensions" not in calls[0]
assert calls[0]["input_type"] == "search_document"
@pytest.mark.asyncio
async def test_litellm_provider_uses_explicit_document_and_query_input_types(monkeypatch):
"""Explicit input_type overrides should support asymmetric providers beyond Cohere."""