diff --git a/docs/litellm-provider.md b/docs/litellm-provider.md index 86a6efa5e..4bb8ce82a 100644 --- a/docs/litellm-provider.md +++ b/docs/litellm-provider.md @@ -50,6 +50,8 @@ All options can be set in config or as environment variables. | `semantic_embedding_forward_dimensions` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_FORWARD_DIMENSIONS` | Auto | Sends `dimensions` to LiteLLM only when supported. Auto is enabled for `text-embedding-3` model strings. | | `semantic_embedding_document_input_type` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_INPUT_TYPE` | Auto | LiteLLM `input_type` for indexed notes/passages. | | `semantic_embedding_query_input_type` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_INPUT_TYPE` | Auto | LiteLLM `input_type` for search queries. | +| `semantic_embedding_document_prefix` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_PREFIX` | Unset | Literal text prefix prepended to indexed document chunks before embedding. | +| `semantic_embedding_query_prefix` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_PREFIX` | Unset | Literal text prefix prepended to search queries before embedding. | | `semantic_embedding_batch_size` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_BATCH_SIZE` | `2` | Number of text chunks per provider request. | | `semantic_embedding_request_concurrency` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_REQUEST_CONCURRENCY` | `4` | Maximum concurrent LiteLLM embedding requests. | | `semantic_embedding_sync_batch_size` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_SYNC_BATCH_SIZE` | `2` | Number of prepared vector jobs flushed through the sync pipeline together. | @@ -124,9 +126,17 @@ export BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_INPUT_TYPE=passage export BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_INPUT_TYPE=query ``` -Changing provider, model, dimensions, dimension-forwarding, or document/query -roles changes Basic Memory's stored vector identity. Rebuild embeddings after -any of those changes: +`input_type` is an API parameter. For models that require role text in the +actual input string, configure literal prefixes instead or in addition: + +```bash +export BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_PREFIX="title: none | text: " +export BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_PREFIX="task: search result | query: " +``` + +Changing provider, model, dimensions, dimension-forwarding, document/query roles, +or prefixes changes Basic Memory's stored vector identity. Rebuild embeddings +after any of those changes: ```bash bm reindex --embeddings diff --git a/docs/semantic-search.md b/docs/semantic-search.md index 31f6c4a5d..fb54f4d37 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -108,6 +108,8 @@ All settings are fields on `BasicMemoryConfig` and can be set via environment va | `semantic_embedding_batch_size` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_BATCH_SIZE` | `2` | Number of texts to embed per batch. | | `semantic_embedding_document_input_type` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_INPUT_TYPE` | Auto for known LiteLLM models | Optional LiteLLM `input_type` for indexed document/passages. | | `semantic_embedding_query_input_type` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_INPUT_TYPE` | Auto for known LiteLLM models | Optional LiteLLM `input_type` for search queries. | +| `semantic_embedding_document_prefix` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_PREFIX` | Unset | Optional literal text prefix prepended to indexed document chunks before embedding. | +| `semantic_embedding_query_prefix` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_PREFIX` | Unset | Optional literal text prefix prepended to search queries before embedding. | | `semantic_vector_k` | `BASIC_MEMORY_SEMANTIC_VECTOR_K` | `100` | Candidate count for vector nearest-neighbour retrieval. Higher values improve recall at the cost of latency. | ## Embedding Providers @@ -190,6 +192,18 @@ export BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_INPUT_TYPE=passage export BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_INPUT_TYPE=query ``` +Some asymmetric models require literal role text in the input string rather +than, or in addition to, an API `input_type` parameter: + +```bash +export BASIC_MEMORY_SEMANTIC_EMBEDDING_DOCUMENT_PREFIX="title: none | text: " +export BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_PREFIX="task: search result | query: " +``` + +The document prefix is prepended to indexed chunks during sync/reindex. The +query prefix is prepended to search text for vector and hybrid retrieval. +Prefixes work with `fastembed`, `openai`, and `litellm` providers. + #### Live LiteLLM Validation Provider APIs differ in subtle ways: some accept `dimensions`, some require separate @@ -334,6 +348,7 @@ bm reindex -p my-project - **Model change**: After changing `semantic_embedding_model` - **Dimension change**: After changing `semantic_embedding_dimensions` - **LiteLLM role change**: After changing `semantic_embedding_document_input_type` or `semantic_embedding_query_input_type` +- **Literal prefix change**: After changing `semantic_embedding_document_prefix` or `semantic_embedding_query_prefix` The reindex command shows progress with embedded/skipped/error counts: diff --git a/src/basic_memory/config.py b/src/basic_memory/config.py index 5b0740e9b..2202518f1 100644 --- a/src/basic_memory/config.py +++ b/src/basic_memory/config.py @@ -308,6 +308,20 @@ def __init__(self, **data: Any) -> None: ... "Use with asymmetric embedding models such as Cohere or NVIDIA retrieval models." ), ) + semantic_embedding_document_prefix: str | None = Field( + default=None, + description=( + "Optional literal text prefix prepended to indexed document chunks before " + "embedding. Use with prefix-sensitive asymmetric embedding models." + ), + ) + semantic_embedding_query_prefix: str | None = Field( + default=None, + description=( + "Optional literal text prefix prepended to search queries before embedding. " + "Use with prefix-sensitive asymmetric embedding models." + ), + ) semantic_embedding_sync_batch_size: int = Field( default=2, description="Batch size for vector sync orchestration flushes.", diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index fda3b6b6f..98aff4684 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -83,7 +83,9 @@ async def lifespan(app: FastMCP): f"Semantic search: provider={config.semantic_embedding_provider}, " f"model={config.semantic_embedding_model}, " f"dimensions={config.semantic_embedding_dimensions or 'auto'}, " - f"batch_size={config.semantic_embedding_batch_size}" + f"batch_size={config.semantic_embedding_batch_size}, " + f"document_prefix_set={bool(config.semantic_embedding_document_prefix)}, " + f"query_prefix_set={bool(config.semantic_embedding_query_prefix)}" ) # Log configured projects with their routing mode diff --git a/src/basic_memory/repository/embedding_provider.py b/src/basic_memory/repository/embedding_provider.py index 0d8640e7b..d79172404 100644 --- a/src/basic_memory/repository/embedding_provider.py +++ b/src/basic_memory/repository/embedding_provider.py @@ -1,6 +1,6 @@ """Embedding provider protocol for pluggable semantic backends.""" -from typing import Any, Protocol +from typing import Any, Protocol, runtime_checkable class EmbeddingProvider(Protocol): @@ -20,3 +20,19 @@ async def embed_documents(self, texts: list[str]) -> list[list[float]]: def runtime_log_attrs(self) -> dict[str, Any]: """Return provider-specific runtime settings suitable for startup logs.""" ... + + +@runtime_checkable +class EmbeddingIdentityProvider(Protocol): + """Optional capability for providers with semantics beyond model and dimensions.""" + + def identity_key(self) -> str: + """Return a stable identity for persisted-vector invalidation.""" + ... + + +def embedding_provider_identity(provider: EmbeddingProvider) -> str: + """Return a provider's explicit semantic identity or the protocol fallback.""" + if isinstance(provider, EmbeddingIdentityProvider): + return provider.identity_key() + return f"{provider.model_name}:{provider.dimensions}" diff --git a/src/basic_memory/repository/embedding_provider_factory.py b/src/basic_memory/repository/embedding_provider_factory.py index 9d0cf5316..dcee3616d 100644 --- a/src/basic_memory/repository/embedding_provider_factory.py +++ b/src/basic_memory/repository/embedding_provider_factory.py @@ -8,12 +8,17 @@ from basic_memory.config import BasicMemoryConfig, default_fastembed_cache_dir from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.prefixing_provider import ( + PrefixingEmbeddingProvider, + embedding_prefix_digest, + normalize_embedding_prefix, +) # Cache key fields are limited to values that change the *identity* of the loaded # provider instance (provider, model_name, explicit LiteLLM endpoint/key routing, -# dimensions, semantic role/input-type settings, batch/request knobs, and the -# resolved cache dir). Thread/parallel knobs are deliberately excluded - they -# change ONNX *execution* only, not the loaded weights. Including them caused #872: in a +# dimensions, semantic role/input-type/prefix settings, batch/request knobs, +# and the resolved cache dir). Thread/parallel knobs are deliberately excluded - +# they change ONNX *execution* only, not the loaded weights. Including them caused #872: in a # container/cgroup the CPU-derived thread count can drift between calls, producing # a fresh cache key and reloading the ~2.3GB model into a CPU arena that never # returns memory to the OS. @@ -28,6 +33,8 @@ int, str | None, str | None, + str | None, + str | None, str, ] @@ -124,6 +131,8 @@ def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey: app_config.semantic_embedding_request_concurrency, app_config.semantic_embedding_document_input_type, app_config.semantic_embedding_query_input_type, + embedding_prefix_digest(app_config.semantic_embedding_document_prefix), + embedding_prefix_digest(app_config.semantic_embedding_query_prefix), _resolve_cache_dir(app_config), ) @@ -225,6 +234,15 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide else: raise ValueError(f"Unsupported semantic embedding provider: {provider_name}") + document_prefix = normalize_embedding_prefix(app_config.semantic_embedding_document_prefix) + query_prefix = normalize_embedding_prefix(app_config.semantic_embedding_query_prefix) + if document_prefix is not None or query_prefix is not None: + provider = PrefixingEmbeddingProvider( + provider, + document_prefix=document_prefix, + query_prefix=query_prefix, + ) + with _EMBEDDING_PROVIDER_CACHE_LOCK: if cached_provider := _EMBEDDING_PROVIDER_CACHE.get(cache_key): return cached_provider diff --git a/src/basic_memory/repository/prefixing_provider.py b/src/basic_memory/repository/prefixing_provider.py new file mode 100644 index 000000000..6234f9d62 --- /dev/null +++ b/src/basic_memory/repository/prefixing_provider.py @@ -0,0 +1,82 @@ +"""Embedding provider wrapper for role-specific literal text prefixes.""" + +from __future__ import annotations + +import hashlib +from typing import Any + +from basic_memory.repository.embedding_provider import ( + EmbeddingProvider, + embedding_provider_identity, +) + + +def normalize_embedding_prefix(value: str | None) -> str | None: + """Treat unset and empty prefixes as disabled while preserving meaningful spaces.""" + if value == "": + return None + return value + + +def embedding_prefix_digest(value: str | None) -> str: + """Return a stable non-secret prefix identity, reserving ``-`` for unset.""" + normalized = normalize_embedding_prefix(value) + if normalized is None: + return "-" + return hashlib.sha256(normalized.encode("utf-8")).hexdigest() + + +class PrefixingEmbeddingProvider(EmbeddingProvider): + """Apply document/query text prefixes before delegating to an embedding provider.""" + + def __init__( + self, + provider: EmbeddingProvider, + *, + document_prefix: str | None = None, + query_prefix: str | None = None, + ) -> None: + self.provider = provider + self.document_prefix = normalize_embedding_prefix(document_prefix) + self.query_prefix = normalize_embedding_prefix(query_prefix) + + @property + def model_name(self) -> str: + return self.provider.model_name + + @property + def dimensions(self) -> int: + return self.provider.dimensions + + async def embed_query(self, text: str) -> list[float]: + if self.query_prefix is not None: + text = f"{self.query_prefix}{text}" + return await self.provider.embed_query(text) + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + if self.document_prefix is not None: + texts = [f"{self.document_prefix}{text}" for text in texts] + return await self.provider.embed_documents(texts) + + def runtime_log_attrs(self) -> dict[str, Any]: + attrs = self.provider.runtime_log_attrs() + attrs.update( + { + "document_prefix_set": self.document_prefix is not None, + "query_prefix_set": self.query_prefix is not None, + } + ) + if self.document_prefix is not None: + attrs["document_prefix_length"] = len(self.document_prefix) + if self.query_prefix is not None: + attrs["query_prefix_length"] = len(self.query_prefix) + return attrs + + def identity_key(self) -> str: + """Return embedding semantics without exposing literal prefix content.""" + provider_identity = embedding_provider_identity(self.provider) + return ( + f"{type(self.provider).__name__}:{provider_identity}:" + f"document_prefix_sha256={embedding_prefix_digest(self.document_prefix)}:" + f"query_prefix_sha256={embedding_prefix_digest(self.query_prefix)}" + ) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 43ede2b76..aa5f094b1 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -18,7 +18,10 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory import db -from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.embedding_provider import ( + EmbeddingProvider, + embedding_provider_identity, +) from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.semantic_errors import ( SemanticDependenciesMissingError, @@ -692,17 +695,12 @@ def _embedding_model_key(self) -> str: """Build a stable model identity for vector invalidation checks.""" assert self._embedding_provider is not None 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() - + # Trigger: providers can change request/input semantics without changing + # model/dimensions. + # Why: asymmetric providers may use role-specific API params or literal + # text-prefix transforms that change stored vector meaning. + # Outcome: reindex treats those semantic config changes as stale vectors. + provider_identity = embedding_provider_identity(provider) return f"{type(provider).__name__}:{provider_identity}" def _plan_entity_vector_shard( diff --git a/src/basic_memory/schemas/project_info.py b/src/basic_memory/schemas/project_info.py index b8c4f4f73..3db4e9ab4 100644 --- a/src/basic_memory/schemas/project_info.py +++ b/src/basic_memory/schemas/project_info.py @@ -87,6 +87,8 @@ class EmbeddingStatus(BaseModel): embedding_provider: Optional[str] = None embedding_model: Optional[str] = None embedding_dimensions: Optional[int] = None + embedding_document_prefix_set: bool = False + embedding_query_prefix_set: bool = False # Counts total_indexed_entities: int = 0 diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index ecfd623da..f112c92e5 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -1036,6 +1036,8 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: provider = config.semantic_embedding_provider model = config.semantic_embedding_model dimensions = config.semantic_embedding_dimensions + document_prefix_set = bool(config.semantic_embedding_document_prefix) + query_prefix_set = bool(config.semantic_embedding_query_prefix) is_postgres = config.database_backend == DatabaseBackend.POSTGRES @@ -1074,6 +1076,8 @@ async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: embedding_provider=provider, embedding_model=model, embedding_dimensions=dimensions, + embedding_document_prefix_set=document_prefix_set, + embedding_query_prefix_set=query_prefix_set, total_indexed_entities=total_indexed_entities, vector_tables_exist=False, reindex_recommended=True, @@ -1190,6 +1194,8 @@ async def _vec_scalar(vec_sql) -> int: embedding_provider=provider, embedding_model=model, embedding_dimensions=dimensions, + embedding_document_prefix_set=document_prefix_set, + embedding_query_prefix_set=query_prefix_set, total_indexed_entities=total_indexed_entities, vector_tables_exist=False, reindex_recommended=True, @@ -1224,6 +1230,8 @@ async def _vec_scalar(vec_sql) -> int: embedding_provider=provider, embedding_model=model, embedding_dimensions=dimensions, + embedding_document_prefix_set=document_prefix_set, + embedding_query_prefix_set=query_prefix_set, total_indexed_entities=total_indexed_entities, total_entities_with_chunks=total_entities_with_chunks, total_chunks=total_chunks, diff --git a/tests/repository/test_openai_provider.py b/tests/repository/test_openai_provider.py index 7c8841a73..e5a04db3f 100644 --- a/tests/repository/test_openai_provider.py +++ b/tests/repository/test_openai_provider.py @@ -15,6 +15,7 @@ ) from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider from basic_memory.repository.openai_provider import OpenAIEmbeddingProvider +from basic_memory.repository.prefixing_provider import PrefixingEmbeddingProvider from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError @@ -547,6 +548,69 @@ def test_embedding_provider_factory_creates_new_provider_for_different_cache_key assert provider_a is not provider_b +def test_embedding_provider_factory_wraps_provider_when_prefixes_are_configured(): + """Factory should apply literal prefixes independently of provider backend.""" + config = BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/basic-memory-test"}, + default_project="test-project", + semantic_search_enabled=True, + semantic_embedding_provider="fastembed", + semantic_embedding_document_prefix="title: none | text: ", + semantic_embedding_query_prefix="task: search result | query: ", + ) + + provider = create_embedding_provider(config) + + assert isinstance(provider, PrefixingEmbeddingProvider) + assert isinstance(provider.provider, FastEmbedEmbeddingProvider) + assert provider.document_prefix == "title: none | text: " + assert provider.query_prefix == "task: search result | query: " + + +def test_embedding_provider_factory_reuses_provider_for_same_prefixes(): + """Prefix fields participate in the process-local provider cache key.""" + config = BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/basic-memory-test"}, + default_project="test-project", + semantic_search_enabled=True, + semantic_embedding_provider="fastembed", + semantic_embedding_document_prefix="doc: ", + semantic_embedding_query_prefix="query: ", + ) + + provider_a = create_embedding_provider(config) + provider_b = create_embedding_provider(config) + + assert provider_a is provider_b + + +def test_embedding_provider_factory_separates_cache_for_different_prefixes(): + """Changing literal prefixes should not reuse a stale cached provider.""" + shared_config = { + "env": "test", + "projects": {"test-project": "/tmp/basic-memory-test"}, + "default_project": "test-project", + "semantic_search_enabled": True, + "semantic_embedding_provider": "fastembed", + "semantic_embedding_query_prefix": "query: ", + } + first_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_document_prefix="doc: ", + ) + second_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_document_prefix="document: ", + ) + + first_provider = create_embedding_provider(first_config) + second_provider = create_embedding_provider(second_config) + + assert first_provider is not second_provider + + def test_embedding_provider_factory_reuses_provider_when_only_thread_knobs_differ(): """Thread/parallel knobs tune ONNX execution, not model identity (#872). diff --git a/tests/repository/test_prefixing_provider.py b/tests/repository/test_prefixing_provider.py new file mode 100644 index 000000000..70d7fda57 --- /dev/null +++ b/tests/repository/test_prefixing_provider.py @@ -0,0 +1,147 @@ +"""Tests for role-specific literal embedding text prefixes.""" + +import hashlib +from typing import Any + +import pytest + +from basic_memory.config import BasicMemoryConfig +from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.embedding_provider_factory import _provider_cache_key +from basic_memory.repository.prefixing_provider import ( + PrefixingEmbeddingProvider, + normalize_embedding_prefix, +) + + +class _RecordingEmbeddingProvider(EmbeddingProvider): + model_name = "stub-model" + dimensions = 3 + + def __init__(self) -> None: + self.document_calls: list[list[str]] = [] + self.query_calls: list[str] = [] + + async def embed_query(self, text: str) -> list[float]: + self.query_calls.append(text) + return [1.0, 0.0, 0.0] + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + self.document_calls.append(texts) + return [[0.0, 1.0, 0.0] for _ in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {"provider_batch_size": 7} + + +@pytest.mark.asyncio +async def test_prefixing_provider_applies_role_specific_prefixes(): + """Documents and queries should receive their own literal text prefixes.""" + inner = _RecordingEmbeddingProvider() + provider = PrefixingEmbeddingProvider( + inner, + document_prefix="title: none | text: ", + query_prefix="task: search result | query: ", + ) + + await provider.embed_documents(["indexed chunk"]) + await provider.embed_query("retrieval text") + + assert inner.document_calls == [["title: none | text: indexed chunk"]] + assert inner.query_calls == ["task: search result | query: retrieval text"] + + +@pytest.mark.asyncio +async def test_prefixing_provider_preserves_unset_and_empty_prefix_behavior(): + """Unset or empty prefixes should not change provider inputs.""" + inner = _RecordingEmbeddingProvider() + provider = PrefixingEmbeddingProvider(inner, document_prefix="", query_prefix=None) + + await provider.embed_documents(["indexed chunk"]) + await provider.embed_query("retrieval text") + + assert inner.document_calls == [["indexed chunk"]] + assert inner.query_calls == ["retrieval text"] + assert normalize_embedding_prefix("") is None + + +def test_prefixing_provider_identity_key_uses_prefix_digests(): + """Prefix changes should alter identity without exposing literal values.""" + first = PrefixingEmbeddingProvider( + _RecordingEmbeddingProvider(), + document_prefix="doc: ", + query_prefix="query: ", + ) + second = PrefixingEmbeddingProvider( + _RecordingEmbeddingProvider(), + document_prefix="document: ", + query_prefix="query: ", + ) + + first_key = first.identity_key() + second_key = second.identity_key() + + assert first_key != second_key + assert f"document_prefix_sha256={hashlib.sha256(b'doc: ').hexdigest()}" in first_key + assert f"query_prefix_sha256={hashlib.sha256(b'query: ').hexdigest()}" in first_key + assert "doc: " not in first_key + assert "query: " not in first_key + + +def test_prefixing_provider_identity_key_distinguishes_unset_from_literal_dash(): + """Unset prefixes must not collide with a real dash prefix.""" + unset_document = PrefixingEmbeddingProvider( + _RecordingEmbeddingProvider(), + document_prefix=None, + query_prefix="query: ", + ) + dash_document = PrefixingEmbeddingProvider( + _RecordingEmbeddingProvider(), + document_prefix="-", + query_prefix="query: ", + ) + + unset_key = unset_document.identity_key() + dash_key = dash_document.identity_key() + + assert unset_key != dash_key + assert "document_prefix_sha256=-" in unset_key + assert f"document_prefix_sha256={hashlib.sha256(b'-').hexdigest()}" in dash_key + + +def test_provider_cache_key_does_not_expose_prefix_values(): + """Process-local cache diagnostics should contain digests, not configured text.""" + first = BasicMemoryConfig( + semantic_embedding_document_prefix="private document role: ", + semantic_embedding_query_prefix="private query role: ", + ) + second = BasicMemoryConfig( + semantic_embedding_document_prefix="different document role: ", + semantic_embedding_query_prefix="private query role: ", + ) + + first_key = _provider_cache_key(first) + second_key = _provider_cache_key(second) + cache_key_text = repr(first_key) + + assert first_key != second_key + assert "private document role: " not in cache_key_text + assert "private query role: " not in cache_key_text + assert hashlib.sha256(b"private document role: ").hexdigest() in cache_key_text + assert hashlib.sha256(b"private query role: ").hexdigest() in cache_key_text + + +def test_prefixing_provider_reports_runtime_prefix_status(): + """Runtime logs should expose whether prefixes are enabled without raw text.""" + provider = PrefixingEmbeddingProvider( + _RecordingEmbeddingProvider(), + document_prefix="doc: ", + query_prefix=None, + ) + + assert provider.runtime_log_attrs() == { + "provider_batch_size": 7, + "document_prefix_set": True, + "query_prefix_set": False, + "document_prefix_length": 5, + } diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index d99b8d4f5..751b78462 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -1,6 +1,7 @@ """SQLite sqlite-vec search repository tests.""" import asyncio +import hashlib from contextlib import asynccontextmanager from datetime import datetime, timezone from typing import Any, cast @@ -11,7 +12,9 @@ from basic_memory import db from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.litellm_provider import LiteLLMEmbeddingProvider +from basic_memory.repository.prefixing_provider import PrefixingEmbeddingProvider 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 @@ -104,7 +107,7 @@ def _relation_row( def _enable_semantic( search_repository: SQLiteSearchRepository, - embedding_provider: StubEmbeddingProvider | None = None, + embedding_provider: EmbeddingProvider | None = None, ) -> None: try: import sqlite_vec # noqa: F401 @@ -322,6 +325,20 @@ async def test_sqlite_vector_sync_skips_unchanged_and_reembeds_changed_content(s assert model_changed_result.chunks_skipped == 0 assert model_changed_result.embedding_jobs_total == model_changed_result.chunks_total + _enable_semantic( + search_repository, + PrefixingEmbeddingProvider( + StubEmbeddingProviderV2(), + document_prefix="doc: ", + query_prefix="query: ", + ), + ) + prefix_changed_result = await search_repository.sync_entity_vectors_batch([111]) + assert prefix_changed_result.entities_synced == 1 + assert prefix_changed_result.entities_skipped == 0 + assert prefix_changed_result.chunks_skipped == 0 + assert prefix_changed_result.embedding_jobs_total == prefix_changed_result.chunks_total + def test_sqlite_embedding_model_key_includes_litellm_role_settings(): """LiteLLM role changes should invalidate previously embedded document chunks.""" @@ -364,6 +381,26 @@ def test_sqlite_embedding_model_key_ignores_litellm_api_base(): assert "token@example.test" not in custom_endpoint_key +def test_sqlite_embedding_model_key_includes_literal_prefixes(): + """Literal prefixes change vector semantics and must invalidate stored chunks.""" + repo = _make_sqlite_repo_for_unit_tests() + repo._embedding_provider = PrefixingEmbeddingProvider( + StubEmbeddingProvider(), + document_prefix="title: none | text: ", + query_prefix="task: search result | query: ", + ) + + key = repo._embedding_model_key() + + assert key.startswith("PrefixingEmbeddingProvider:StubEmbeddingProvider:stub:4:") + assert f"document_prefix_sha256={hashlib.sha256(b'title: none | text: ').hexdigest()}" in key + assert ( + f"query_prefix_sha256={hashlib.sha256(b'task: search result | query: ').hexdigest()}" in key + ) + assert "title: none | text: " not in key + assert "task: search result | query: " not in 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.""" diff --git a/tests/test_config.py b/tests/test_config.py index 8409a0640..c2f1588ca 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1192,6 +1192,21 @@ def test_semantic_embedding_forward_dimensions_can_be_set(self): config = BasicMemoryConfig(semantic_embedding_forward_dimensions=True) assert config.semantic_embedding_forward_dimensions is True + def test_semantic_embedding_prefixes_default_to_none(self): + """Literal embedding text prefixes should be disabled by default.""" + config = BasicMemoryConfig() + assert config.semantic_embedding_document_prefix is None + assert config.semantic_embedding_query_prefix is None + + def test_semantic_embedding_prefixes_can_be_set(self): + """Document and query embedding prefixes should be stored independently.""" + config = BasicMemoryConfig( + semantic_embedding_document_prefix="title: none | text: ", + semantic_embedding_query_prefix="task: search result | query: ", + ) + assert config.semantic_embedding_document_prefix == "title: none | text: " + assert config.semantic_embedding_query_prefix == "task: search result | query: " + def test_semantic_postgres_prepare_concurrency_defaults_to_4(self): """Postgres prepare concurrency should default to a conservative window of 4.""" config = BasicMemoryConfig()