From a3739fa72e5fb68e83647f4b708980e9f9c8d5c8 Mon Sep 17 00:00:00 2001
From: phernandez <paul@basicmachines.co>
Date: Thu, 28 May 2026 20:14:44 -0500
Subject: [PATCH] fix(core): scope LiteLLM dimension requests

Signed-off-by: phernandez <paul@basicmachines.co>
---
 docs/semantic-search.md                        |  4 ++++
 .../repository/litellm_provider.py             | 18 +++++++++++++++++-
 tests/repository/test_litellm_provider.py      | 15 +++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/docs/semantic-search.md b/docs/semantic-search.md
index 3d6f3df9..f7ae30a5 100644
--- a/docs/semantic-search.md
+++ b/docs/semantic-search.md
@@ -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`
diff --git a/src/basic_memory/repository/litellm_provider.py b/src/basic_memory/repository/litellm_provider.py
index c92875eb..f8ce339b 100644
--- a/src/basic_memory/repository/litellm_provider.py
+++ b/src/basic_memory/repository/litellm_provider.py
@@ -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:
diff --git a/tests/repository/test_litellm_provider.py b/tests/repository/test_litellm_provider.py
index e07ef463..6f49070c 100644
--- a/tests/repository/test_litellm_provider.py
+++ b/tests/repository/test_litellm_provider.py
@@ -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."""