mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
fix(core): scope LiteLLM dimension requests
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user