diff --git a/src/backend/base/langflow/api/v1/memories.py b/src/backend/base/langflow/api/v1/memories.py index 53ce559477b1..423540028936 100644 --- a/src/backend/base/langflow/api/v1/memories.py +++ b/src/backend/base/langflow/api/v1/memories.py @@ -47,7 +47,7 @@ from langflow.services.jobs import DuplicateJobError from langflow.services.memory_base.kb_path_helpers import BackendProvisioningError from langflow.services.memory_base.provider_scope import MemoryBaseFlowNotFoundError -from langflow.services.memory_base.service import PreprocessingValidationError +from langflow.services.memory_base.service import EmbeddingProviderValidationError, PreprocessingValidationError router = APIRouter(tags=["Memories"], prefix="/memories", include_in_schema=False) @@ -160,7 +160,7 @@ async def create_memory_base( except ModelProviderPolicyError as exc: # Keep a hidden provider indistinguishable from one that does not exist. raise HTTPException(status_code=404, detail="Model provider not found") from exc - except PreprocessingValidationError as exc: + except (PreprocessingValidationError, EmbeddingProviderValidationError) as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc except BackendProvisioningError as exc: # Bad remote vector-store config (unreachable / wrong credentials) — diff --git a/src/backend/base/langflow/services/database/models/memory_base/model.py b/src/backend/base/langflow/services/database/models/memory_base/model.py index da14cb016147..327aa416eb12 100644 --- a/src/backend/base/langflow/services/database/models/memory_base/model.py +++ b/src/backend/base/langflow/services/database/models/memory_base/model.py @@ -50,6 +50,15 @@ class MemoryBaseCreate(MemoryBaseBase): backend_type: str | None = None backend_config: dict = Field(default_factory=dict) + # Provider that serves ``embedding_model``, as selected by the caller. + # Declared here rather than on ``MemoryBaseBase`` for the same reason as the + # backend fields above: no column is added to the ``memory_base`` table + # because the value is persisted in the backing ``knowledge_base`` row's + # ``model_selection``, which every ingestion/retrieval path resolves against. + # ``None`` falls back to name-based inference, which cannot see + # live-discovered models (e.g. an OpenAI-Compatible endpoint's catalog). + embedding_provider: str | None = None + @model_validator(mode="after") def preprocessing_defaults(self) -> "MemoryBaseCreate": if self.preprocessing and not self.preproc_model: diff --git a/src/backend/base/langflow/services/memory_base/embedding_helpers.py b/src/backend/base/langflow/services/memory_base/embedding_helpers.py index 3a4b184c7b70..de41d4383804 100644 --- a/src/backend/base/langflow/services/memory_base/embedding_helpers.py +++ b/src/backend/base/langflow/services/memory_base/embedding_helpers.py @@ -38,6 +38,12 @@ def infer_embedding_provider(embedding_model: str) -> str: """Derive embedding provider name from a model string. + Fallback only: used when the caller did not supply the provider it + selected. Inference is a guess — it has no user context, so it cannot see + live-discovered models (an OpenAI-Compatible endpoint's catalog is + per-user) and mislabels them via the ``"OpenAI"`` default. Callers that + know the selected provider must pass it instead. + Looks up the model in the unified models catalog first so the answer matches what the UI dropdown shows; falls back to pattern-based inference for legacy/edge cases. diff --git a/src/backend/base/langflow/services/memory_base/service.py b/src/backend/base/langflow/services/memory_base/service.py index d87c9bbea2df..6017d009e3f0 100644 --- a/src/backend/base/langflow/services/memory_base/service.py +++ b/src/backend/base/langflow/services/memory_base/service.py @@ -20,8 +20,9 @@ from lfx.base.knowledge_bases.backends import is_local_chroma from lfx.base.knowledge_bases.backends.postgres import resolve_default_kb_backend from lfx.base.knowledge_bases.validation import validate_collection_name -from lfx.base.models.provider_registry import is_api_key_optional +from lfx.base.models.provider_registry import is_api_key_optional, provider_name_for_id, resolve_provider_id from lfx.base.models.unified_models import get_api_key_for_provider +from lfx.base.models.unified_models.class_registry import EMBEDDING_PROVIDER_CLASS_MAPPING from lfx.services.model_provider_policy import ( ModelProviderPolicyPurpose, aresolve_model_provider_policy, @@ -29,7 +30,7 @@ ) from sqlmodel import col, select -from langflow.api.utils.kb_helpers import local_chroma_rejection_reason +from langflow.api.utils.kb_helpers import local_chroma_rejection_reason, resolve_embedding_selection from langflow.services.base import Service from langflow.services.database.models.memory_base.model import ( MemoryBase, @@ -82,6 +83,10 @@ class PreprocessingValidationError(ValueError): """Raised when preprocessing is enabled but the provider API key is absent.""" +class EmbeddingProviderValidationError(ValueError): + """Raised when the caller-selected embedding provider cannot serve embeddings.""" + + def _require_preprocessing_model_provider(user_id: uuid.UUID, preproc_model: str | None) -> str | None: """Require CONFIGURE access for a supplied preprocessing model identity.""" provider = _infer_preprocessing_model_provider(preproc_model) @@ -105,18 +110,67 @@ def _infer_preprocessing_model_provider(preproc_model: str | None) -> str | None raise PreprocessingValidationError(str(exc)) from exc +# ``get_embedding_provider`` reports this sentinel for a knowledge_base row whose +# ``model_selection`` carries no provider; it must never be authorized or persisted. +_UNKNOWN_PROVIDER = "Unknown" + + +def _select_embedding_provider(embedding_provider: str | None, embedding_model: str) -> str: + """Return the canonical embedding provider for a Memory Base. + + The caller's explicit selection wins. It is canonicalized through the provider + registry so the persisted value is the exact key every downstream embedding + lookup uses (``EMBEDDING_PROVIDER_CLASS_MAPPING`` is matched verbatim, while the + policy layer matches case- and alias-insensitively): ``"openai"`` becomes + ``"OpenAI"`` and ``"IBM watsonx.ai"`` becomes ``"IBM WatsonX"``. Names the + registry does not know are kept as supplied so the policy layer can reject + them. Name-based inference is the fallback only when nothing usable was given. + """ + supplied = (embedding_provider or "").strip() + if not supplied or supplied == _UNKNOWN_PROVIDER: + return infer_embedding_provider(embedding_model) + return provider_name_for_id(resolve_provider_id(supplied)) or supplied + + +def _require_embedding_class(provider: str) -> None: + """Reject a caller-selected provider that cannot serve embeddings. + + Runs after the policy preflight on create only. The OSS policy allows every + provider name, so without this check a typo or a chat-only provider would be + persisted and fail at the first ingestion with a misleading credential error. + Stored providers on existing Memory Bases are not re-checked so an uninstalled + bundle never blocks deactivating or renaming a Memory Base. + + Raises: + EmbeddingProviderValidationError: ``provider`` has no registered embedding class. + """ + if provider not in EMBEDDING_PROVIDER_CLASS_MAPPING: + msg = f"Embedding provider '{provider}' is not available for embeddings." + raise EmbeddingProviderValidationError(msg) + + async def _preflight_memory_provider_configuration( *, flow, actor_user_id: uuid.UUID, actor_is_superuser: bool, embedding_model: str, + embedding_provider: str | None, preproc_model: str | None, ) -> tuple[str | None, str]: - """Authorize selected configuration providers before any owner credential read.""" + """Authorize selected configuration providers before any owner credential read. + + ``embedding_provider`` is the provider the caller actually selected and is + authoritative when supplied. Name-based inference is only the fallback: it + cannot see live-discovered models (an OpenAI-Compatible endpoint's catalog is + per-user), so guessing from the model name labels those models as OpenAI and + every later credential lookup asks for the wrong key. + """ preprocessing_provider = _infer_preprocessing_model_provider(preproc_model) - embedding_provider = infer_embedding_provider(embedding_model) - providers = list(dict.fromkeys(provider for provider in (preprocessing_provider, embedding_provider) if provider)) + selected_embedding_provider = _select_embedding_provider(embedding_provider, embedding_model) + providers = list( + dict.fromkeys(provider for provider in (preprocessing_provider, selected_embedding_provider) if provider) + ) with scoped_model_provider_policy_for_flow( flow, user_id=actor_user_id, @@ -129,7 +183,7 @@ async def _preflight_memory_provider_configuration( ) for provider in providers: provider_policy.require(provider) - return preprocessing_provider, embedding_provider + return preprocessing_provider, selected_embedding_provider def _validate_preprocessing_api_key(user_id: uuid.UUID, preproc_model: str | None) -> None: @@ -229,8 +283,12 @@ async def create( actor_user_id=user_id, actor_is_superuser=is_superuser, embedding_model=payload.embedding_model, + embedding_provider=payload.embedding_provider, preproc_model=payload.preproc_model, ) + # Policy first so a hidden provider stays indistinguishable from a missing one. + if (payload.embedding_provider or "").strip(): + _require_embedding_class(embedding_provider) if payload.preprocessing: _validate_preprocessing_provider_api_key( user_id, @@ -307,9 +365,9 @@ async def create( raise ValueError(msg) mb = MemoryBase( - # ``backend_type``/``backend_config`` live on the knowledge_base - # row created above, not on this table. - **payload.model_dump(exclude={"user_id", "backend_type", "backend_config"}), + # ``backend_type``/``backend_config``/``embedding_provider`` live + # on the knowledge_base row created above, not on this table. + **payload.model_dump(exclude={"user_id", "backend_type", "backend_config", "embedding_provider"}), user_id=user_id, kb_name=kb_name, ) @@ -414,11 +472,21 @@ async def update( return None flow = await resolve_owned_memory_flow(db, flow_id=mb.flow_id, user_id=owner_user_id) + # The embedding provider chosen at create time is persisted on the + # backing knowledge_base row (the memory_base table stores only the + # model name). Re-inferring it from that name would relabel a + # live-discovered model — e.g. one served by an OpenAI-Compatible + # endpoint — as OpenAI and authorize the wrong provider. + stored_embedding_provider, _stored_embedding_model = await resolve_embedding_selection( + user_id=owner_user_id, + kb_name=mb.kb_name, + ) preprocessing_provider, _embedding_provider = await _preflight_memory_provider_configuration( flow=flow, actor_user_id=actor_user_id, actor_is_superuser=actor_is_superuser, embedding_model=mb.embedding_model, + embedding_provider=stored_embedding_provider, preproc_model=mb.preproc_model if mb.preprocessing else None, ) diff --git a/src/backend/tests/unit/test_memory_bases.py b/src/backend/tests/unit/test_memory_bases.py index 19e773240735..17dfa5993d78 100644 --- a/src/backend/tests/unit/test_memory_bases.py +++ b/src/backend/tests/unit/test_memory_bases.py @@ -714,7 +714,10 @@ async def authorize_providers(*, user_id, providers, purpose): with ( patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), - patch("langflow.services.memory_base.service.infer_embedding_provider", return_value="OpenAI"), + patch( + "langflow.services.memory_base.service.resolve_embedding_selection", + AsyncMock(return_value=("OpenAI", mb.embedding_model)), + ), patch("langflow.services.memory_base.service.infer_llm_provider", return_value="Anthropic"), patch( "langflow.services.memory_base.service.aresolve_model_provider_policy", @@ -757,7 +760,10 @@ async def test_update_denied_embedding_provider_stops_before_mutation_or_persist with ( patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), - patch("langflow.services.memory_base.service.infer_embedding_provider", return_value="Anthropic"), + patch( + "langflow.services.memory_base.service.resolve_embedding_selection", + AsyncMock(return_value=("Anthropic", mb.embedding_model)), + ), patch("langflow.services.memory_base.service.aresolve_model_provider_policy", resolve_policy), pytest.raises(ModelProviderPolicyError), ): @@ -874,6 +880,10 @@ async def deny_current_project(**_kwargs): with ( patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), + patch( + "langflow.services.memory_base.service.resolve_embedding_selection", + AsyncMock(return_value=("OpenAI", mb.embedding_model)), + ), patch( "langflow.services.memory_base.service.aresolve_model_provider_policy", side_effect=deny_current_project, @@ -925,6 +935,198 @@ async def test_update_missing_stored_flow_fails_before_provider_or_persistence(s db.commit.assert_not_awaited() +class TestMemoryBaseEmbeddingProviderSelection: + """The provider the caller selected is authoritative over name-based inference. + + Regression for the OpenAI-Compatible case: those models are discovered + per-user, so ``infer_embedding_provider`` cannot see them and defaults them + to ``"OpenAI"``. The wrong label was persisted on the backing + ``knowledge_base`` row, and ingestion then demanded an OpenAI API key. + """ + + @pytest.fixture + def service(self): + from langflow.services.memory_base.service import MemoryBaseService + + return MemoryBaseService() + + @staticmethod + def _fake_scope(mock_db): + class _FakeCtx: + async def __aenter__(self): + return mock_db + + async def __aexit__(self, *_args): + pass + + scope = MagicMock() + scope.return_value = _FakeCtx() + return scope + + @staticmethod + def _create_db(flow): + flow_result = MagicMock() + flow_result.first.return_value = flow + missing_result = MagicMock() + missing_result.first.return_value = None + db = AsyncMock() + db.exec = AsyncMock(side_effect=[flow_result, missing_result, missing_result]) + db.add = MagicMock() + db.commit = AsyncMock() + db.refresh = AsyncMock() + return db + + @pytest.mark.asyncio + async def test_create_persists_the_selected_provider_without_inferring(self, service): + from langflow.services.database.models.flow.model import Flow + + user_id = uuid.uuid4() + flow_id = uuid.uuid4() + flow = Flow(id=flow_id, user_id=user_id, name="flow") + payload = MemoryBaseCreate( + name="mb", + flow_id=flow_id, + embedding_model="mock-embed-1", + embedding_provider="OpenAI Compatible", + ) + db = self._create_db(flow) + create_record = AsyncMock() + infer = MagicMock(return_value="OpenAI") + + with ( + patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), + patch("langflow.services.memory_base.service.resolve_kb_username", AsyncMock(return_value="testuser")), + patch("langflow.services.memory_base.service.infer_embedding_provider", infer), + patch( + "langflow.services.memory_base.service.aresolve_model_provider_policy", + AsyncMock(return_value=MagicMock()), + ), + patch("langflow.services.memory_base.service.initialize_kb", AsyncMock()), + patch("langflow.api.utils.knowledge_base_service.create_record", create_record), + # The bundle registers this at startup; unit tests run without extensions loaded. + patch.dict( + "langflow.services.memory_base.service.EMBEDDING_PROVIDER_CLASS_MAPPING", + {"OpenAI Compatible": "OpenAIEmbeddings"}, + ), + ): + await service.create(payload, user_id=user_id) + + assert create_record.await_args.kwargs["model_selection"] == { + "name": "mock-embed-1", + "provider": "OpenAI Compatible", + } + infer.assert_not_called() + + @pytest.mark.asyncio + async def test_create_rejects_a_provider_without_embedding_class_before_provisioning(self, service): + from langflow.services.database.models.flow.model import Flow + from langflow.services.memory_base.service import EmbeddingProviderValidationError + + user_id = uuid.uuid4() + flow_id = uuid.uuid4() + flow = Flow(id=flow_id, user_id=user_id, name="flow") + payload = MemoryBaseCreate( + name="mb", + flow_id=flow_id, + embedding_model="text-embedding-3-small", + embedding_provider="OpenAl", + ) + db = self._create_db(flow) + create_record = AsyncMock() + initialize = AsyncMock() + + with ( + patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), + patch("langflow.services.memory_base.service.resolve_kb_username", AsyncMock(return_value="testuser")), + patch( + "langflow.services.memory_base.service.aresolve_model_provider_policy", + AsyncMock(return_value=MagicMock()), + ), + patch("langflow.services.memory_base.service.initialize_kb", initialize), + patch("langflow.api.utils.knowledge_base_service.create_record", create_record), + pytest.raises(EmbeddingProviderValidationError, match="OpenAl"), + ): + await service.create(payload, user_id=user_id) + + initialize.assert_not_awaited() + create_record.assert_not_awaited() + + @pytest.mark.asyncio + async def test_create_without_a_provider_falls_back_to_inference(self, service): + from langflow.services.database.models.flow.model import Flow + + user_id = uuid.uuid4() + flow_id = uuid.uuid4() + flow = Flow(id=flow_id, user_id=user_id, name="flow") + payload = MemoryBaseCreate(name="mb", flow_id=flow_id, embedding_model="text-embedding-3-small") + db = self._create_db(flow) + create_record = AsyncMock() + infer = MagicMock(return_value="OpenAI") + + with ( + patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), + patch("langflow.services.memory_base.service.resolve_kb_username", AsyncMock(return_value="testuser")), + patch("langflow.services.memory_base.service.infer_embedding_provider", infer), + patch( + "langflow.services.memory_base.service.aresolve_model_provider_policy", + AsyncMock(return_value=MagicMock()), + ), + patch("langflow.services.memory_base.service.initialize_kb", AsyncMock()), + patch("langflow.api.utils.knowledge_base_service.create_record", create_record), + ): + await service.create(payload, user_id=user_id) + + infer.assert_called_once_with("text-embedding-3-small") + assert create_record.await_args.kwargs["model_selection"] == { + "name": "text-embedding-3-small", + "provider": "OpenAI", + } + + @pytest.mark.asyncio + async def test_update_authorizes_the_stored_provider_not_an_inferred_one(self, service): + from langflow.services.database.models.flow.model import Flow + from lfx.services.model_provider_policy import ModelProviderPolicyPurpose + + user_id = uuid.uuid4() + mb = _make_mb(user_id=user_id, threshold=10) + mb.embedding_model = "mock-embed-1" + flow = Flow(id=mb.flow_id, user_id=user_id, name="flow") + mb_result = MagicMock() + mb_result.first.return_value = mb + flow_result = MagicMock() + flow_result.first.return_value = flow + db = AsyncMock() + db.exec = AsyncMock(side_effect=[mb_result, flow_result]) + db.add = MagicMock() + policy = MagicMock() + resolve_policy = AsyncMock(return_value=policy) + infer = MagicMock(return_value="OpenAI") + resolve_selection = AsyncMock(return_value=("OpenAI Compatible", "mock-embed-1")) + + with ( + patch("langflow.services.memory_base.service.session_scope", self._fake_scope(db)), + patch("langflow.services.memory_base.service.resolve_embedding_selection", resolve_selection), + patch("langflow.services.memory_base.service.infer_embedding_provider", infer), + patch("langflow.services.memory_base.service.aresolve_model_provider_policy", resolve_policy), + ): + await service.update( + mb.id, + owner_user_id=user_id, + patch=MemoryBaseUpdate(threshold=99), + actor_user_id=user_id, + ) + + resolve_selection.assert_awaited_once_with(user_id=user_id, kb_name=mb.kb_name) + resolve_policy.assert_awaited_once_with( + user_id=user_id, + providers=["OpenAI Compatible"], + purpose=ModelProviderPolicyPurpose.CONFIGURE, + ) + policy.require.assert_called_once_with("OpenAI Compatible") + infer.assert_not_called() + assert mb.threshold == 99 + + class TestMemoryBaseGuardPassesRealKbIdentity: """The ID-bearing guards must pass the REAL kb identity, not actor-as-owner. @@ -2598,6 +2800,34 @@ async def test_create_missing_api_key_returns_422(self, mock_user): assert exc_info.value.status_code == 422 assert "API key" in exc_info.value.detail + @pytest.mark.asyncio + async def test_create_unusable_embedding_provider_returns_422(self, mock_user): + from fastapi import HTTPException + from langflow.api.v1.memories import create_memory_base + from langflow.services.memory_base.service import EmbeddingProviderValidationError + + payload = MemoryBaseCreate( + name="mb", + flow_id=uuid.uuid4(), + user_id=mock_user.id, + kb_name="kb", + embedding_model="text-embedding-3-small", + embedding_provider="OpenAl", + ) + + svc = MagicMock() + svc.create = AsyncMock( + side_effect=EmbeddingProviderValidationError("Embedding provider 'OpenAl' is not available for embeddings.") + ) + with ( + patch("langflow.api.v1.memories.get_memory_base_service", return_value=svc), + pytest.raises(HTTPException) as exc_info, + ): + await create_memory_base(current_user=mock_user, payload=payload) + + assert exc_info.value.status_code == 422 + assert "OpenAl" in exc_info.value.detail + @pytest.mark.asyncio async def test_update_missing_api_key_returns_403(self, mock_user): from fastapi import HTTPException @@ -3773,3 +4003,49 @@ async def test_backfill_creates_row_for_orphan_memory_base(self): assert kwargs["backend_type"] == "chroma" assert kwargs["source_types"] == ["memory"] assert kwargs["model_selection"]["name"] == "text-embedding-3-small" + + +class TestSelectEmbeddingProvider: + """``_select_embedding_provider`` persists exactly the key downstream lookups use. + + The policy layer matches providers case- and alias-insensitively, but + ``EMBEDDING_PROVIDER_CLASS_MAPPING`` is matched verbatim, so a variant that + is authorized but persisted as-is would create a Memory Base that can never + embed. The ``"Unknown"`` sentinel from ``get_embedding_provider`` must fall + back to inference instead of being authorized as a provider name. + """ + + @pytest.mark.parametrize( + ("supplied", "expected"), + [ + ("OpenAI", "OpenAI"), + ("openai", "OpenAI"), + (" OpenAI ", "OpenAI"), + ("IBM watsonx.ai", "IBM WatsonX"), + ], + ) + def test_supplied_provider_is_canonicalized(self, supplied, expected): + from langflow.services.memory_base.service import _select_embedding_provider + + assert _select_embedding_provider(supplied, "text-embedding-3-small") == expected + + def test_unregistered_provider_is_kept_for_policy_to_reject(self): + from langflow.services.memory_base.service import _select_embedding_provider + + assert _select_embedding_provider("Not A Provider", "some-model") == "Not A Provider" + + @pytest.mark.parametrize("provider", ["OpenAl", "Not A Provider", "Anthropic"]) + def test_provider_without_embedding_class_is_rejected(self, provider): + """Typos, unknown names, and chat-only providers fail fast instead of breaking ingestion later.""" + from langflow.services.memory_base.service import EmbeddingProviderValidationError, _require_embedding_class + + with pytest.raises(EmbeddingProviderValidationError, match=provider): + _require_embedding_class(provider) + + @pytest.mark.parametrize("supplied", [None, "", " ", "Unknown"]) + def test_missing_or_unknown_provider_falls_back_to_inference(self, supplied): + from langflow.services.memory_base import service as service_module + + with patch.object(service_module, "infer_embedding_provider", return_value="Inferred") as infer: + assert service_module._select_embedding_provider(supplied, "text-embedding-3-small") == "Inferred" + infer.assert_called_once_with("text-embedding-3-small") diff --git a/src/frontend/src/controllers/API/queries/memories/__tests__/memories-mutation-hooks-cache.test.ts b/src/frontend/src/controllers/API/queries/memories/__tests__/memories-mutation-hooks-cache.test.ts index 70e72b840226..7d211921d0cf 100644 --- a/src/frontend/src/controllers/API/queries/memories/__tests__/memories-mutation-hooks-cache.test.ts +++ b/src/frontend/src/controllers/API/queries/memories/__tests__/memories-mutation-hooks-cache.test.ts @@ -134,6 +134,7 @@ describe("memories mutation hooks cache wiring", () => { name: "New Memory", flow_id: "flow-1", embedding_model: "text-embedding-3-small", + embedding_provider: "OpenAI", }); }); diff --git a/src/frontend/src/controllers/API/queries/memories/types.ts b/src/frontend/src/controllers/API/queries/memories/types.ts index be198a42cd4d..fef495de88a5 100644 --- a/src/frontend/src/controllers/API/queries/memories/types.ts +++ b/src/frontend/src/controllers/API/queries/memories/types.ts @@ -77,6 +77,11 @@ export interface CreateMemoryPayload { name: string; flow_id: string; embedding_model: string; + // Provider that serves `embedding_model`. Required: the server can only guess + // it from the model name otherwise, and that guess cannot see live-discovered + // models (e.g. an OpenAI-Compatible endpoint's catalog), which it would + // mislabel as OpenAI. + embedding_provider: string; threshold?: number; auto_capture?: boolean; preprocessing?: boolean; diff --git a/src/frontend/src/modals/createMemoryModal/__tests__/useCreateMemoryModal.test.tsx b/src/frontend/src/modals/createMemoryModal/__tests__/useCreateMemoryModal.test.tsx index cb7e5c5402cb..df5717d2435a 100644 --- a/src/frontend/src/modals/createMemoryModal/__tests__/useCreateMemoryModal.test.tsx +++ b/src/frontend/src/modals/createMemoryModal/__tests__/useCreateMemoryModal.test.tsx @@ -339,6 +339,9 @@ describe("useCreateMemoryModal", () => { name: "My Memory", flow_id: "flow-1", embedding_model: "text-embedding-3-small", + // Sent explicitly so the server never has to guess the provider from + // the model name (that guess mislabels live-discovered models). + embedding_provider: "OpenAI", preproc_model: "gpt-4o-mini", preproc_instructions: "summarize", preprocessing: true, @@ -347,6 +350,60 @@ describe("useCreateMemoryModal", () => { ); }); + it("forwards a non-OpenAI embedding provider verbatim", () => { + // Regression: an OpenAI-Compatible model is discovered per-user, so the + // server's name-based inference defaults it to OpenAI and then demands an + // OpenAI API key. The selected provider must travel with the model name. + mockModelProvidersResult = { + ...mockModelProvidersResult, + data: [ + { + provider: "OpenAI Compatible", + is_enabled: true, + icon: "Bot", + models: [ + { + model_name: "mock-embed-1", + metadata: { model_type: "embeddings" }, + }, + ], + }, + ], + }; + mockEnabledModelsResult = { + ...mockEnabledModelsResult, + data: { + enabled_models: { "OpenAI Compatible": { "mock-embed-1": true } }, + }, + }; + + const { result } = renderHook(() => + useCreateMemoryModal({ flowId: "flow-1", onClose: jest.fn() }), + ); + + act(() => { + result.current.setName("My Memory"); + result.current.setSelectedEmbeddingModel([ + { + id: "mock-embed-1", + name: "mock-embed-1", + provider: "OpenAI Compatible", + } as ModelOption, + ]); + }); + + act(() => { + result.current.handleSubmit(); + }); + + expect(mockMutate).toHaveBeenCalledWith( + expect.objectContaining({ + embedding_model: "mock-embed-1", + embedding_provider: "OpenAI Compatible", + }), + ); + }); + it("omits the implicit Chroma type so the server can choose its default", () => { const { result } = renderHook(() => useCreateMemoryModal({ flowId: "flow-1", onClose: jest.fn() }), diff --git a/src/frontend/src/modals/createMemoryModal/useCreateMemoryModal.ts b/src/frontend/src/modals/createMemoryModal/useCreateMemoryModal.ts index 8204860c0219..ef1394984c2b 100644 --- a/src/frontend/src/modals/createMemoryModal/useCreateMemoryModal.ts +++ b/src/frontend/src/modals/createMemoryModal/useCreateMemoryModal.ts @@ -322,6 +322,7 @@ export function useCreateMemoryModal({ name: name.trim(), flow_id: flowId, embedding_model: embeddingSelection?.name, + embedding_provider: embeddingSelection?.provider, auto_capture: true, threshold: parsedThreshold, preprocessing: preprocessingEnabled,