diff --git a/docs/litellm-provider.md b/docs/litellm-provider.md index e1164b3a0..86a6efa5e 100644 --- a/docs/litellm-provider.md +++ b/docs/litellm-provider.md @@ -44,6 +44,8 @@ All options can be set in config or as environment variables. | `semantic_search_enabled` | `BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED` | Auto | Set to `true` to force vector/hybrid support on. | | `semantic_embedding_provider` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_PROVIDER` | `fastembed` | Set to `litellm` for the LiteLLM provider. | | `semantic_embedding_model` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_MODEL` | `bge-small-en-v1.5` | With `litellm`, the default is remapped to `openai/text-embedding-3-small`. | +| `semantic_embedding_api_base` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_BASE` | Unset | Optional custom endpoint for the LiteLLM provider, including local or self-hosted OpenAI-compatible servers. | +| `semantic_embedding_api_key` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_KEY` | Unset | Optional API key passed directly to the LiteLLM provider. When unset, LiteLLM continues to read provider credential env vars such as `OPENAI_API_KEY`. | | `semantic_embedding_dimensions` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DIMENSIONS` | Provider default | Required for non-default LiteLLM models because vector tables are dimensioned before the first API call. | | `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. | @@ -84,6 +86,27 @@ Set this only when your deployment supports it: export BASIC_MEMORY_SEMANTIC_EMBEDDING_FORWARD_DIMENSIONS=true ``` +## Custom OpenAI-Compatible Endpoints + +Set `semantic_embedding_api_base` when an OpenAI-compatible embedding server is +available somewhere other than the provider's default endpoint. Include the API +version prefix expected by the server, commonly `/v1`: + +```bash +export BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED=true +export BASIC_MEMORY_SEMANTIC_EMBEDDING_PROVIDER=litellm +export BASIC_MEMORY_SEMANTIC_EMBEDDING_MODEL=openai/local-embedding-model +export BASIC_MEMORY_SEMANTIC_EMBEDDING_API_BASE=http://127.0.0.1:8080/v1 +export BASIC_MEMORY_SEMANTIC_EMBEDDING_API_KEY=local-key +export BASIC_MEMORY_SEMANTIC_EMBEDDING_DIMENSIONS=768 +``` + +`semantic_embedding_api_key` is useful when you want credentials in Basic +Memory's config instead of the process environment. Leave it unset to preserve +LiteLLM's normal provider environment lookup, including `OPENAI_API_KEY`. The +API key can be a placeholder when the local server does not authenticate, but +LiteLLM or the selected backend may still require it to be present. + ## Asymmetric Models Some embedding models use different request roles for indexed documents and @@ -102,8 +125,8 @@ export BASIC_MEMORY_SEMANTIC_EMBEDDING_QUERY_INPUT_TYPE=query ``` Changing provider, model, dimensions, dimension-forwarding, or document/query -roles changes the meaning of stored vectors. Rebuild embeddings after any of -those changes: +roles changes Basic Memory's stored vector identity. Rebuild embeddings after +any of those changes: ```bash bm reindex --embeddings @@ -273,6 +296,13 @@ cat > /tmp/litellm-cases.json <<'JSON' "api_key_env": "NVIDIA_NIM_API_KEY", "document_input_type": "passage", "query_input_type": "query" + }, + { + "name": "local-openai-compatible", + "model": "openai/local-embedding-model", + "dimensions": 768, + "api_key_env": "OPENAI_API_KEY", + "api_base": "http://127.0.0.1:8080/v1" } ] JSON diff --git a/docs/semantic-search.md b/docs/semantic-search.md index e5665c226..31f6c4a5d 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -101,6 +101,8 @@ All settings are fields on `BasicMemoryConfig` and can be set via environment va | `semantic_search_enabled` | `BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED` | Auto (`true` when semantic deps are available) | Enable semantic search. Required before vector/hybrid modes work. | | `semantic_embedding_provider` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_PROVIDER` | `"fastembed"` | Embedding provider: `"fastembed"` (local), `"openai"` (API), or `"litellm"` (multi-provider API, **experimental** — advanced users only). | | `semantic_embedding_model` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_MODEL` | `"bge-small-en-v1.5"` | Model identifier. Auto-adjusted per provider if left at default. | +| `semantic_embedding_api_base` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_BASE` | Unset | Optional custom endpoint for the LiteLLM provider, including local or self-hosted OpenAI-compatible servers. | +| `semantic_embedding_api_key` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_KEY` | Unset | Optional API key passed directly to the LiteLLM provider. When unset, LiteLLM continues to read provider credential env vars such as `OPENAI_API_KEY`. | | `semantic_embedding_dimensions` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_DIMENSIONS` | Provider default | Vector dimensions. 384 for FastEmbed, 1536 for OpenAI/LiteLLM OpenAI. Required when using a non-default LiteLLM model. | | `semantic_embedding_forward_dimensions` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_FORWARD_DIMENSIONS` | Auto | LiteLLM-only override for whether configured dimensions are sent as a provider-side output-size request. | | `semantic_embedding_batch_size` | `BASIC_MEMORY_SEMANTIC_EMBEDDING_BATCH_SIZE` | `2` | Number of texts to embed per batch. | @@ -162,6 +164,20 @@ dimensions. If an Azure/OpenAI deployment uses an arbitrary LiteLLM model string `azure/` and the underlying model supports reduced dimensions, set `BASIC_MEMORY_SEMANTIC_EMBEDDING_FORWARD_DIMENSIONS=true`. +For a local or self-hosted OpenAI-compatible embedding server, set the custom +endpoint explicitly: + +```bash +export BASIC_MEMORY_SEMANTIC_EMBEDDING_PROVIDER=litellm +export BASIC_MEMORY_SEMANTIC_EMBEDDING_MODEL=openai/local-embedding-model +export BASIC_MEMORY_SEMANTIC_EMBEDDING_API_BASE=http://127.0.0.1:8080/v1 +export BASIC_MEMORY_SEMANTIC_EMBEDDING_API_KEY=local-key +export BASIC_MEMORY_SEMANTIC_EMBEDDING_DIMENSIONS=768 +``` + +Leave `BASIC_MEMORY_SEMANTIC_EMBEDDING_API_KEY` unset to keep LiteLLM's normal +environment credential lookup, including `OPENAI_API_KEY`. + 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` @@ -245,7 +261,8 @@ just test-litellm-live --cases-file /tmp/litellm-nvidia-cases.json For repeatable local runs, put the same JSON array in a file and pass `just test-litellm-live --cases-file path/to/litellm-cases.json`. -When switching providers, models, dimensions, or LiteLLM document/query input types, rebuild embeddings: +When switching providers, models, dimensions, or LiteLLM document/query input types, +rebuild embeddings: ```bash bm reindex --embeddings diff --git a/src/basic_memory/config.py b/src/basic_memory/config.py index 768deed69..729893407 100644 --- a/src/basic_memory/config.py +++ b/src/basic_memory/config.py @@ -247,6 +247,21 @@ def __init__(self, **data: Any) -> None: ... default="bge-small-en-v1.5", description="Embedding model identifier used by the local provider.", ) + semantic_embedding_api_base: str | None = Field( + default=None, + description=( + "Optional custom API base URL for the LiteLLM embedding provider. " + "Use this for OpenAI-compatible local or self-hosted embedding servers." + ), + ) + semantic_embedding_api_key: str | None = Field( + default=None, + description=( + "Optional API key passed directly to the LiteLLM embedding provider. " + "When unset, LiteLLM continues to resolve credentials from provider " + "environment variables such as OPENAI_API_KEY." + ), + ) semantic_embedding_dimensions: int | None = Field( default=None, description=( diff --git a/src/basic_memory/mcp/tools/basic_memory_diagnostics.py b/src/basic_memory/mcp/tools/basic_memory_diagnostics.py index b6b658d04..d97a47e47 100644 --- a/src/basic_memory/mcp/tools/basic_memory_diagnostics.py +++ b/src/basic_memory/mcp/tools/basic_memory_diagnostics.py @@ -10,11 +10,11 @@ from basic_memory.mcp.server import mcp # Fields in BasicMemoryConfig that contain secrets and must never be surfaced. -_SECRET_FIELDS = frozenset({"cloud_api_key"}) +_SECRET_FIELDS = frozenset({"cloud_api_key", "semantic_embedding_api_key"}) # Fields whose values are URLs that may embed user:password credentials. # The userinfo component is stripped before surfacing. -_URL_FIELDS = frozenset({"database_url"}) +_URL_FIELDS = frozenset({"database_url", "semantic_embedding_api_base"}) _SECRET_QUERY_KEYS = frozenset( { diff --git a/src/basic_memory/repository/embedding_provider_factory.py b/src/basic_memory/repository/embedding_provider_factory.py index 60e028b23..9d0cf5316 100644 --- a/src/basic_memory/repository/embedding_provider_factory.py +++ b/src/basic_memory/repository/embedding_provider_factory.py @@ -1,5 +1,6 @@ """Factory for creating configured semantic embedding providers.""" +import hashlib import os from threading import Lock @@ -9,16 +10,18 @@ from basic_memory.repository.embedding_provider import EmbeddingProvider # Cache key fields are limited to values that change the *identity* of the loaded -# model (provider, model_name, dimensions, LiteLLM role/input-type/forward-dimension -# settings, batch/request knobs that affect the LiteLLM identity, 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 +# 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 # 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. type ProviderCacheKey = tuple[ str, str, + str | None, + str | None, int | None, bool | None, int, @@ -33,10 +36,17 @@ _FASTEMBED_MAX_THREADS = 8 +def _sensitive_value_digest(value: str | None) -> str | None: + """Return a stable non-secret token for process-local cache diagnostics.""" + if not value: + return None + return hashlib.sha256(value.encode("utf-8")).hexdigest() + + def _resolve_cache_dir(app_config: BasicMemoryConfig) -> str: """Resolve the effective FastEmbed cache dir for this config. - Uses an explicit ``is not None`` check — an empty string override from + Uses an explicit ``is not None`` check - an empty string override from config or ``BASIC_MEMORY_SEMANTIC_EMBEDDING_CACHE_DIR`` is an invalid path, not a request to fall back to the default, and FastEmbed's error message is clearer than silently swapping in a different directory. @@ -86,9 +96,9 @@ def _resolve_fastembed_runtime_knobs( def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey: - """Build a stable cache key from model-identity semantic embedding config. + """Build a stable cache key from process-local embedding provider config. - Uses the *resolved* cache dir — not the raw config field — so different + Uses the *resolved* cache dir - not the raw config field - so different FASTEMBED_CACHE_PATH values produce distinct cache keys even when the config field itself is unset. @@ -96,9 +106,18 @@ def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey: execution, not which model weights are loaded, and resolving them from the runtime CPU budget makes the key drift between calls in a container (#872). """ + provider_name = app_config.semantic_embedding_provider.strip().lower() + litellm_api_base_digest = None + litellm_api_key_digest = None + if provider_name == "litellm": + litellm_api_base_digest = _sensitive_value_digest(app_config.semantic_embedding_api_base) + litellm_api_key_digest = _sensitive_value_digest(app_config.semantic_embedding_api_key) + return ( - app_config.semantic_embedding_provider.strip().lower(), + provider_name, app_config.semantic_embedding_model, + litellm_api_base_digest, + litellm_api_key_digest, app_config.semantic_embedding_dimensions, app_config.semantic_embedding_forward_dimensions, app_config.semantic_embedding_batch_size, @@ -130,19 +149,19 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide # deliberately build it *outside* the lock to avoid serializing every caller # behind a single cold start. This opens a by-design TOCTOU window where both # threads may construct a provider. - # Outcome: the second check-and-set below resolves the race — the first writer + # Outcome: the second check-and-set below resolves the race - the first writer # wins and the loser's redundant provider is discarded, so the cache still # yields a single process-wide singleton per key. with _EMBEDDING_PROVIDER_CACHE_LOCK: if cached_provider := _EMBEDDING_PROVIDER_CACHE.get(cache_key): return cached_provider - provider_name = app_config.semantic_embedding_provider.strip().lower() extra_kwargs: dict = {} if app_config.semantic_embedding_dimensions is not None: extra_kwargs["dimensions"] = app_config.semantic_embedding_dimensions provider: EmbeddingProvider + provider_name = app_config.semantic_embedding_provider.strip().lower() if provider_name == "fastembed": # Deferred import: fastembed (and its onnxruntime dep) may not be installed from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider @@ -194,6 +213,8 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide ) provider = LiteLLMEmbeddingProvider( model_name=model_name, + api_key=app_config.semantic_embedding_api_key, + api_base=app_config.semantic_embedding_api_base, batch_size=app_config.semantic_embedding_batch_size, request_concurrency=app_config.semantic_embedding_request_concurrency, document_input_type=app_config.semantic_embedding_document_input_type, @@ -210,8 +231,8 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide # Trigger: a distinct cache key is being inserted while the cache already # holds entries for other keys. # Why: the provider is meant to be a process-wide singleton (#872). A second - # key means something bypassed reuse — a real config change, or a regression - # that reintroduces volatile fields into the key — and each new key reloads + # key means something bypassed reuse - a real config change, or a regression + # that reintroduces volatile fields into the key - and each new key reloads # the ~2.3GB ONNX model into a CPU arena that never releases memory. # Outcome: surface the bypass so future leaks are diagnosable from logs. if _EMBEDDING_PROVIDER_CACHE: diff --git a/src/basic_memory/repository/litellm_provider.py b/src/basic_memory/repository/litellm_provider.py index c15a66775..2e8a7d384 100644 --- a/src/basic_memory/repository/litellm_provider.py +++ b/src/basic_memory/repository/litellm_provider.py @@ -93,6 +93,7 @@ def __init__( request_concurrency: int = 4, dimensions: int = 1536, api_key: str | None = None, + api_base: str | None = None, timeout: float = 30.0, document_input_type: str | None = None, query_input_type: str | None = None, @@ -103,6 +104,7 @@ def __init__( self.batch_size = batch_size self.request_concurrency = request_concurrency self._api_key = api_key + self._api_base = api_base self._timeout = timeout default_document_input_type, default_query_input_type = _default_input_types(model_name) self.document_input_type = document_input_type or default_document_input_type @@ -130,12 +132,13 @@ def identity_key(self) -> str: forward_dimensions = str( _should_forward_dimensions(self.model_name, self.forward_dimensions) ).lower() - return ( + identity = ( 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}" ) + return identity async def _embed(self, texts: list[str], *, input_type: str | None) -> list[list[float]]: if not texts: @@ -162,6 +165,8 @@ async def embed_batch(batch_index: int, batch: list[str]) -> None: params["dimensions"] = self.dimensions if self._api_key: params["api_key"] = self._api_key + if self._api_base is not None: + params["api_base"] = self._api_base if input_type: params["input_type"] = input_type diff --git a/test-int/semantic/litellm_live_harness.py b/test-int/semantic/litellm_live_harness.py index a6b9178aa..50d013308 100644 --- a/test-int/semantic/litellm_live_harness.py +++ b/test-int/semantic/litellm_live_harness.py @@ -33,6 +33,7 @@ class LiteLLMLiveCase: model: str dimensions: int api_key_env: str | None = None + api_base: str | None = None document_input_type: str | None = None query_input_type: str | None = None forward_dimensions: bool | None = None @@ -111,6 +112,7 @@ def load_custom_cases(raw: str | None) -> list[LiteLLMLiveCase]: model=_required_string(case_data, "model"), dimensions=dimensions, api_key_env=_optional_string(case_data, "api_key_env"), + api_base=_optional_string(case_data, "api_base"), document_input_type=_optional_string(case_data, "document_input_type"), query_input_type=_optional_string(case_data, "query_input_type"), forward_dimensions=_optional_bool(case_data, "forward_dimensions"), @@ -195,6 +197,7 @@ async def evaluate_case( dimensions=case.dimensions, batch_size=2, api_key=api_key, + api_base=case.api_base, timeout=60.0, document_input_type=case.document_input_type, query_input_type=case.query_input_type, diff --git a/test-int/semantic/test_litellm_live_harness.py b/test-int/semantic/test_litellm_live_harness.py index 14ec09fdb..8b06b4a69 100644 --- a/test-int/semantic/test_litellm_live_harness.py +++ b/test-int/semantic/test_litellm_live_harness.py @@ -37,8 +37,8 @@ async def embed_documents(self, texts: list[str]) -> list[list[float]]: return [[0.0, 1.0], [1.0, 0.0]] -def test_load_custom_cases_accepts_roles_and_dimension_forwarding(): - """Custom JSON should preserve provider-specific role and dimensions options.""" +def test_load_custom_cases_accepts_endpoint_roles_and_dimension_forwarding(): + """Custom JSON should preserve endpoint, role, and dimensions options.""" raw = json.dumps( [ { @@ -46,6 +46,7 @@ def test_load_custom_cases_accepts_roles_and_dimension_forwarding(): "model": "azure/basic-memory-embeddings", "dimensions": 512, "api_key_env": "AZURE_API_KEY", + "api_base": "https://example.openai.azure.com", "document_input_type": "passage", "query_input_type": "query", "forward_dimensions": True, @@ -61,6 +62,7 @@ def test_load_custom_cases_accepts_roles_and_dimension_forwarding(): model="azure/basic-memory-embeddings", dimensions=512, api_key_env="AZURE_API_KEY", + api_base="https://example.openai.azure.com", document_input_type="passage", query_input_type="query", forward_dimensions=True, @@ -111,6 +113,28 @@ async def test_evaluate_case_reports_metrics_for_valid_provider(): assert result.total_latency_ms >= 0 +@pytest.mark.asyncio +async def test_evaluate_case_forwards_api_base_to_provider(): + """Custom live cases should exercise the configured endpoint.""" + provider_kwargs: dict[str, object] = {} + + def provider_factory(**kwargs): + provider_kwargs.update(kwargs) + return FakeProvider(**kwargs) + + await evaluate_case( + LiteLLMLiveCase( + name="local-provider", + model="openai/local-embedding-model", + dimensions=2, + api_base="http://127.0.0.1:8080/v1", + ), + provider_factory=provider_factory, + ) + + assert provider_kwargs["api_base"] == "http://127.0.0.1:8080/v1" + + @pytest.mark.asyncio async def test_evaluate_case_rejects_wrong_ranking(): """A case should fail when the query ranks the distractor document higher.""" diff --git a/tests/mcp/test_tool_basic_memory_diagnostics.py b/tests/mcp/test_tool_basic_memory_diagnostics.py index a5b7c5f94..99b0710d6 100644 --- a/tests/mcp/test_tool_basic_memory_diagnostics.py +++ b/tests/mcp/test_tool_basic_memory_diagnostics.py @@ -33,6 +33,33 @@ def test_redact_config_removes_cloud_api_key(): assert "projects" in result +def test_redact_config_removes_semantic_embedding_api_key(): + raw = { + "semantic_embedding_api_key": "provider-secret", + "semantic_embedding_api_base": "https://embeddings.example.com/v1", + } + + result = _redact_config(raw) + + assert "semantic_embedding_api_key" not in result + assert result["semantic_embedding_api_base"] == "https://embeddings.example.com/v1" + + +def test_redact_config_scrubs_semantic_embedding_api_base_credentials(): + raw = { + "semantic_embedding_api_base": ( + "https://provider-user:provider-password@embeddings.example.com/v1" + "?api_key=query-secret&timeout=30" + ), + } + + result = _redact_config(raw) + + assert result["semantic_embedding_api_base"] == ( + "https://***@embeddings.example.com/v1?api_key=%2A%2A%2A&timeout=30" + ) + + def test_redact_config_passes_through_safe_fields(): raw = {"default_project": "main", "log_level": "INFO", "env": "dev"} result = _redact_config(raw) @@ -114,6 +141,30 @@ def test_diagnostics_redacts_cloud_api_key(tmp_path): assert "cloud_api_key" not in result +def test_diagnostics_redacts_semantic_embedding_api_key(tmp_path): + """Provider credentials must never appear in diagnostic output.""" + config_data = { + "semantic_embedding_api_key": "provider-super-secret", + "semantic_embedding_api_base": ( + "https://provider-user:provider-password@embeddings.example.com/v1" + "?api_key=query-secret&timeout=30" + ), + "projects": {}, + } + config_file = tmp_path / "config.json" + config_file.write_text(json.dumps(config_data)) + + result = basic_memory_diagnostics() + + assert "provider-super-secret" not in result + assert "provider-user" not in result + assert "provider-password" not in result + assert "query-secret" not in result + assert "semantic_embedding_api_key" not in result + assert "embeddings.example.com" in result + assert "timeout=30" in result + + def test_diagnostics_config_missing(tmp_path): """When config file does not exist, output should say so.""" config_file = tmp_path / "config.json" diff --git a/tests/repository/test_litellm_provider.py b/tests/repository/test_litellm_provider.py index d21d97c56..912bdb984 100644 --- a/tests/repository/test_litellm_provider.py +++ b/tests/repository/test_litellm_provider.py @@ -9,6 +9,7 @@ import pytest from basic_memory.config import BasicMemoryConfig +import basic_memory.repository.embedding_provider_factory as embedding_provider_factory_module from basic_memory.repository.embedding_provider_factory import ( create_embedding_provider, reset_embedding_provider_cache, @@ -122,6 +123,49 @@ async def test_litellm_provider_api_key_omitted_when_none(monkeypatch): assert "api_key" not in calls[0] +@pytest.mark.asyncio +async def test_litellm_provider_api_base_forwarded(monkeypatch): + """api_base should be passed to litellm.aembedding when configured.""" + calls = _install_litellm_stub(monkeypatch) + provider = LiteLLMEmbeddingProvider( + model_name="openai/local-embedding-model", + api_base="http://127.0.0.1:8080/v1", + dimensions=3, + ) + + await provider.embed_query("test") + + assert calls[0]["api_base"] == "http://127.0.0.1:8080/v1" + + +@pytest.mark.asyncio +async def test_litellm_provider_api_base_omitted_when_none(monkeypatch): + """api_base should not override LiteLLM's endpoint resolution when unset.""" + calls = _install_litellm_stub(monkeypatch) + provider = LiteLLMEmbeddingProvider(dimensions=3) + + await provider.embed_query("test") + + assert "api_base" not in calls[0] + + +def test_litellm_provider_identity_key_ignores_api_base(): + """Endpoint routing is not part of Basic Memory's persisted vector identity.""" + default_endpoint = LiteLLMEmbeddingProvider(dimensions=3) + custom_endpoint = LiteLLMEmbeddingProvider( + dimensions=3, + api_base="http://127.0.0.1:8080/v1", + ) + + assert default_endpoint.identity_key() == ( + "openai/text-embedding-3-small:3:document_input_type=-:" + "query_input_type=-:forward_dimensions=true" + ) + assert custom_endpoint.identity_key() == default_endpoint.identity_key() + assert "api_base" not in custom_endpoint.identity_key() + assert "127.0.0.1" not in custom_endpoint.identity_key() + + @pytest.mark.asyncio async def test_litellm_provider_drop_params_always_set(monkeypatch): """drop_params=True should always be in the call kwargs.""" @@ -333,6 +377,128 @@ def test_factory_maps_default_model_for_litellm(): assert provider.model_name == "openai/text-embedding-3-small" +@pytest.mark.asyncio +async def test_factory_forwards_litellm_api_base(monkeypatch): + """Factory should pass the configured custom endpoint to LiteLLM calls.""" + calls = _install_litellm_stub(monkeypatch) + config = BasicMemoryConfig( + env="test", + projects={"test": "/tmp/basic-memory-test"}, + default_project="test", + semantic_search_enabled=True, + semantic_embedding_provider="litellm", + semantic_embedding_model="openai/local-embedding-model", + semantic_embedding_api_base="http://127.0.0.1:8080/v1", + semantic_embedding_dimensions=3, + ) + + provider = create_embedding_provider(config) + assert isinstance(provider, LiteLLMEmbeddingProvider) + await provider.embed_query("test") + + assert calls[0]["api_base"] == "http://127.0.0.1:8080/v1" + + +@pytest.mark.asyncio +async def test_factory_forwards_litellm_api_key(monkeypatch): + """Factory should pass configured LiteLLM API keys without requiring env vars.""" + calls = _install_litellm_stub(monkeypatch) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + config = BasicMemoryConfig( + env="test", + projects={"test": "/tmp/basic-memory-test"}, + default_project="test", + semantic_search_enabled=True, + semantic_embedding_provider="litellm", + semantic_embedding_model="openai/local-embedding-model", + semantic_embedding_api_key="config-key", + semantic_embedding_dimensions=3, + ) + + provider = create_embedding_provider(config) + assert isinstance(provider, LiteLLMEmbeddingProvider) + await provider.embed_query("test") + + assert calls[0]["api_key"] == "config-key" + + +@pytest.mark.asyncio +async def test_factory_omits_litellm_api_key_when_unset(monkeypatch): + """Unset config should preserve LiteLLM's environment credential resolution.""" + calls = _install_litellm_stub(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + config = BasicMemoryConfig( + env="test", + projects={"test": "/tmp/basic-memory-test"}, + default_project="test", + semantic_search_enabled=True, + semantic_embedding_provider="litellm", + semantic_embedding_model="openai/text-embedding-3-small", + semantic_embedding_dimensions=3, + ) + + provider = create_embedding_provider(config) + assert isinstance(provider, LiteLLMEmbeddingProvider) + await provider.embed_query("test") + + assert "api_key" not in calls[0] + + +def test_factory_cache_separates_litellm_api_bases_without_exposing_endpoint(): + """Endpoint routing should separate cached providers without leaking raw URLs.""" + shared_config = { + "env": "test", + "projects": {"test": "/tmp/basic-memory-test"}, + "default_project": "test", + "semantic_search_enabled": True, + "semantic_embedding_provider": "litellm", + "semantic_embedding_model": "openai/text-embedding-3-small", + } + first_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_api_base="http://token@example.test/v1", + ) + second_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_api_base="http://other-token@example.test/v1", + ) + first_provider = create_embedding_provider(first_config) + second_provider = create_embedding_provider(second_config) + + assert first_provider is not second_provider + first_cache_key = repr(embedding_provider_factory_module._provider_cache_key(first_config)) + assert "token@example.test" not in first_cache_key + assert "api_base" not in first_cache_key + + +def test_factory_cache_separates_litellm_api_keys_without_exposing_secret(): + """Configured key routing should separate cached providers without leaking keys.""" + shared_config = { + "env": "test", + "projects": {"test": "/tmp/basic-memory-test"}, + "default_project": "test", + "semantic_search_enabled": True, + "semantic_embedding_provider": "litellm", + "semantic_embedding_model": "openai/text-embedding-3-small", + } + first_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_api_key="first-secret-key", + ) + second_config = BasicMemoryConfig( + **shared_config, + semantic_embedding_api_key="second-secret-key", + ) + first_provider = create_embedding_provider(first_config) + second_provider = create_embedding_provider(second_config) + + assert first_provider is not second_provider + first_cache_key = repr(embedding_provider_factory_module._provider_cache_key(first_config)) + second_cache_key = repr(embedding_provider_factory_module._provider_cache_key(second_config)) + assert "first-secret-key" not in first_cache_key + assert "second-secret-key" not in second_cache_key + + def test_factory_forwards_litellm_document_and_query_input_types(): """Factory should pass role-specific LiteLLM input_type config to the provider.""" config = BasicMemoryConfig( diff --git a/tests/repository/test_sqlite_vector_search_repository.py b/tests/repository/test_sqlite_vector_search_repository.py index 9bcd72da7..d99b8d4f5 100644 --- a/tests/repository/test_sqlite_vector_search_repository.py +++ b/tests/repository/test_sqlite_vector_search_repository.py @@ -347,6 +347,23 @@ def test_sqlite_embedding_model_key_includes_litellm_role_settings(): assert "query_input_type=query" in passage_query_key +def test_sqlite_embedding_model_key_ignores_litellm_api_base(): + """LiteLLM endpoint routing is not part of stored vector identity.""" + repo = _make_sqlite_repo_for_unit_tests() + repo._embedding_provider = LiteLLMEmbeddingProvider(dimensions=3) + default_endpoint_key = repo._embedding_model_key() + + repo._embedding_provider = LiteLLMEmbeddingProvider( + dimensions=3, + api_base="http://token@example.test/v1", + ) + custom_endpoint_key = repo._embedding_model_key() + + assert default_endpoint_key == custom_endpoint_key + assert "api_base" not in custom_endpoint_key + assert "token@example.test" not in custom_endpoint_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."""