fix(core): include LiteLLM roles in embedding identity

Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2026-06-06 14:19:59 -05:00
parent bfb63724d4
commit a3dfd7010f
3 changed files with 52 additions and 5 deletions
@@ -123,6 +123,20 @@ class LiteLLMEmbeddingProvider(EmbeddingProvider):
attrs["forward_dimensions"] = self.forward_dimensions
return attrs
def identity_key(self) -> str:
"""Return the embedding semantics that should invalidate stored vectors."""
document_input_type = self.document_input_type or "-"
query_input_type = self.query_input_type or "-"
forward_dimensions = str(
_should_forward_dimensions(self.model_name, self.forward_dimensions)
).lower()
return (
f"{self.model_name}:{self.dimensions}:"
f"document_input_type={document_input_type}:"
f"query_input_type={query_input_type}:"
f"forward_dimensions={forward_dimensions}"
)
async def _embed(self, texts: list[str], *, input_type: str | None) -> list[list[float]]:
if not texts:
return []
@@ -586,11 +586,19 @@ class SearchRepositoryBase(ABC):
def _embedding_model_key(self) -> str:
"""Build a stable model identity for vector invalidation checks."""
assert self._embedding_provider is not None
return (
f"{type(self._embedding_provider).__name__}:"
f"{self._embedding_provider.model_name}:"
f"{self._embedding_provider.dimensions}"
)
provider = self._embedding_provider
provider_identity = f"{provider.model_name}:{provider.dimensions}"
from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider
if isinstance(provider, LiteLLMEmbeddingProvider):
# Trigger: LiteLLM can change request semantics without changing model/dimensions.
# Why: asymmetric providers use role-specific document/query params, and
# dimension forwarding changes provider-side output-size behavior.
# Outcome: reindex treats those semantic config changes as stale vectors.
provider_identity = provider.identity_key()
return f"{type(provider).__name__}:{provider_identity}"
def _plan_entity_vector_shard(
self,
@@ -11,6 +11,7 @@ from sqlalchemy import text
from basic_memory import db
from basic_memory.config import BasicMemoryConfig, DatabaseBackend
from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider
from basic_memory.repository.search_index_row import SearchIndexRow
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode
@@ -296,6 +297,30 @@ async def test_sqlite_vector_sync_skips_unchanged_and_reembeds_changed_content(s
assert model_changed_result.embedding_jobs_total == model_changed_result.chunks_total
def test_sqlite_embedding_model_key_includes_litellm_role_settings():
"""LiteLLM role changes should invalidate previously embedded document chunks."""
repo = _make_sqlite_repo_for_unit_tests()
repo._embedding_provider = LiteLLMEmbeddingProvider(
model_name="nvidia_nim/nvidia/embed-qa-4",
dimensions=1024,
document_input_type="passage",
query_input_type="query",
)
passage_query_key = repo._embedding_model_key()
repo._embedding_provider = LiteLLMEmbeddingProvider(
model_name="nvidia_nim/nvidia/embed-qa-4",
dimensions=1024,
document_input_type="document",
query_input_type="query",
)
document_query_key = repo._embedding_model_key()
assert passage_query_key != document_query_key
assert "document_input_type=passage" in passage_query_key
assert "query_input_type=query" in passage_query_key
@pytest.mark.asyncio
async def test_sqlite_prepare_window_uses_shared_reads_and_serialized_write_scope(monkeypatch):
"""SQLite should batch read-side prepare work but serialize write-side mutations."""