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