From cd6425bf9c4870afab94282e4e13fca33db8b676 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Fri, 22 May 2026 09:27:04 -0700 Subject: [PATCH] fix: disable chroma server-side embedding functions --- .../base/langflow/api/utils/kb_helpers.py | 5 ++-- .../base/langflow/api/v1/knowledge_bases.py | 3 +- .../services/memory_base/ingestion.py | 8 +++++- .../services/memory_base/kb_path_helpers.py | 3 +- .../langflow/services/memory_base/task.py | 8 +++++- .../knowledge_bases/test_chroma_backend.py | 7 +++++ .../test_chroma_vector_store_component.py | 24 ++++++++++++++++ .../tests/unit/test_knowledge_bases_api.py | 5 ++++ .../base/knowledge_bases/backends/chroma.py | 3 ++ .../lfx/base/vectorstores/chroma_security.py | 28 +++++++++++++++++++ src/lfx/src/lfx/components/chroma/chroma.py | 2 ++ .../files_and_knowledge/knowledge.py | 2 ++ .../files_and_knowledge/memory_retrieval.py | 2 ++ .../lfx/components/vectorstores/local_db.py | 2 ++ 14 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 src/lfx/src/lfx/base/vectorstores/chroma_security.py diff --git a/src/backend/base/langflow/api/utils/kb_helpers.py b/src/backend/base/langflow/api/utils/kb_helpers.py index e5d4e611f4..fc58999ddf 100644 --- a/src/backend/base/langflow/api/utils/kb_helpers.py +++ b/src/backend/base/langflow/api/utils/kb_helpers.py @@ -37,6 +37,7 @@ from lfx.base.knowledge_bases.ingestion_sources import ( KBIngestionSource, ) from lfx.base.knowledge_bases.ingestion_sources.base import IngestionItemStatus, IngestionRunStatus +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.components.models_and_agents.embedding_model import EmbeddingModelComponent from lfx.log import logger @@ -192,7 +193,7 @@ class KBStorageHelper: has_data = any((kb_path / m).exists() for m in ["chroma", "chroma.sqlite3", "index"]) if has_data: client = KBStorageHelper.get_fresh_chroma_client(kb_path) - chroma = Chroma(client=client, collection_name=kb_name) + chroma = Chroma(client=client, collection_name=kb_name, **chroma_langchain_collection_kwargs()) with contextlib.suppress(Exception): chroma.delete_collection() chroma = None @@ -419,7 +420,7 @@ class KBAnalysisHelper: try: if created_locally: client = KBStorageHelper.get_fresh_chroma_client(kb_path) - chroma = Chroma(client=client, collection_name=kb_path.name) + chroma = Chroma(client=client, collection_name=kb_path.name, **chroma_langchain_collection_kwargs()) if chroma is None: return diff --git a/src/backend/base/langflow/api/v1/knowledge_bases.py b/src/backend/base/langflow/api/v1/knowledge_bases.py index 3666c85316..c2c9006c27 100644 --- a/src/backend/base/langflow/api/v1/knowledge_bases.py +++ b/src/backend/base/langflow/api/v1/knowledge_bases.py @@ -21,6 +21,7 @@ from lfx.base.knowledge_bases.ingestion_sources import ( get_source_class, registered_sources, ) +from lfx.base.vectorstores.chroma_security import chroma_client_create_collection_kwargs from lfx.log import logger from pydantic import BaseModel, Field @@ -588,7 +589,7 @@ async def create_knowledge_base( # This ensures files exist for read operations and avoids 'readonly' errors later try: client = KBStorageHelper.get_fresh_chroma_client(kb_path) - client.create_collection(name=kb_name) + client.create_collection(name=kb_name, **chroma_client_create_collection_kwargs()) except (OSError, ValueError, chromadb.errors.ChromaError) as e: logger.warning("Initial Chroma setup for %s failed: %s", kb_name, e) finally: diff --git a/src/backend/base/langflow/services/memory_base/ingestion.py b/src/backend/base/langflow/services/memory_base/ingestion.py index b3359ec1b9..b693335659 100644 --- a/src/backend/base/langflow/services/memory_base/ingestion.py +++ b/src/backend/base/langflow/services/memory_base/ingestion.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING import chromadb.errors from langchain_chroma import Chroma +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.log.logger import logger from sqlmodel import col, func, select @@ -431,7 +432,12 @@ async def _delete_chunks_for_session( client = KBStorageHelper.get_fresh_chroma_client(kb_path) try: - chroma = Chroma(client=client, embedding_function=embeddings, collection_name=kb_name) + chroma = Chroma( + client=client, + embedding_function=embeddings, + collection_name=kb_name, + **chroma_langchain_collection_kwargs(), + ) await chroma.adelete(where={"session_id": {"$eq": session_id}}) # Refresh on-disk metrics so the UI reflects the post-purge state. try: diff --git a/src/backend/base/langflow/services/memory_base/kb_path_helpers.py b/src/backend/base/langflow/services/memory_base/kb_path_helpers.py index 3887a22425..221273920d 100644 --- a/src/backend/base/langflow/services/memory_base/kb_path_helpers.py +++ b/src/backend/base/langflow/services/memory_base/kb_path_helpers.py @@ -13,6 +13,7 @@ import uuid from datetime import datetime, timezone from typing import TYPE_CHECKING +from lfx.base.vectorstores.chroma_security import chroma_client_create_collection_kwargs from lfx.log.logger import logger from sqlmodel import select @@ -108,7 +109,7 @@ async def initialize_kb( # Initialize Chroma collection so the directory is non-empty and readable try: client = KBStorageHelper.get_fresh_chroma_client(kb_path) - client.create_collection(name=kb_name) + client.create_collection(name=kb_name, **chroma_client_create_collection_kwargs()) except (OSError, ValueError, chromadb.errors.ChromaError) as exc: await logger.awarning("Initial Chroma setup for %s failed: %s", kb_name, exc) finally: diff --git a/src/backend/base/langflow/services/memory_base/task.py b/src/backend/base/langflow/services/memory_base/task.py index 2d0183cbf2..862ce26b9e 100644 --- a/src/backend/base/langflow/services/memory_base/task.py +++ b/src/backend/base/langflow/services/memory_base/task.py @@ -29,6 +29,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING from langchain_chroma import Chroma +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.log.logger import logger from sqlalchemy import text from sqlmodel import Session, col, select @@ -386,7 +387,12 @@ async def ingest_memory_task(*, request: IngestionRequest) -> dict: client = KBStorageHelper.get_fresh_chroma_client(kb_path) written = 0 try: - chroma = Chroma(client=client, embedding_function=embeddings, collection_name=kb_name) + chroma = Chroma( + client=client, + embedding_function=embeddings, + collection_name=kb_name, + **chroma_langchain_collection_kwargs(), + ) written = await KBIngestionHelper.write_documents_to_chroma( documents=documents, diff --git a/src/backend/tests/unit/base/knowledge_bases/test_chroma_backend.py b/src/backend/tests/unit/base/knowledge_bases/test_chroma_backend.py index aaed9bc220..70c585060d 100644 --- a/src/backend/tests/unit/base/knowledge_bases/test_chroma_backend.py +++ b/src/backend/tests/unit/base/knowledge_bases/test_chroma_backend.py @@ -318,11 +318,18 @@ class TestChromaCloudMode: with ( patch("chromadb.CloudClient", return_value=mock_client) as mock_cloud, patch("chromadb.PersistentClient") as mock_local, + patch("lfx.base.knowledge_bases.backends.chroma.Chroma", return_value=MagicMock()) as mock_chroma, ): bk._build_vector_store() mock_cloud.assert_called_once() mock_local.assert_not_called() + mock_chroma.assert_called_once_with( + client=mock_client, + collection_name="cloud_test_kb", + embedding_function=bk.embedding_function, + collection_configuration={"embedding_function": None}, + ) def test_get_cloud_client_passes_optional_host_port(self, tmp_path: Path): from unittest.mock import patch diff --git a/src/backend/tests/unit/components/vectorstores/test_chroma_vector_store_component.py b/src/backend/tests/unit/components/vectorstores/test_chroma_vector_store_component.py index 8ee0a20441..667edddd3d 100644 --- a/src/backend/tests/unit/components/vectorstores/test_chroma_vector_store_component.py +++ b/src/backend/tests/unit/components/vectorstores/test_chroma_vector_store_component.py @@ -3,6 +3,10 @@ from typing import Any from unittest.mock import MagicMock, patch import pytest +from lfx.base.vectorstores.chroma_security import ( + chroma_client_create_collection_kwargs, + chroma_langchain_collection_kwargs, +) from lfx.components.chroma import ChromaVectorStoreComponent from lfx.schema.data import Data @@ -37,9 +41,29 @@ def test_remote_chroma_server_uses_http_client() -> None: client=mock_client, embedding_function=None, collection_name="remote_collection", + collection_configuration={"embedding_function": None}, ) +def test_chroma_collection_security_kwargs_disable_server_side_embedding_functions() -> None: + assert chroma_langchain_collection_kwargs() == { + "collection_configuration": {"embedding_function": None}, + } + assert chroma_client_create_collection_kwargs() == { + "configuration": {"embedding_function": None}, + "embedding_function": None, + } + + +def test_chroma_collection_security_kwargs_are_fresh_dicts() -> None: + first = chroma_langchain_collection_kwargs() + first["collection_configuration"]["embedding_function"] = "unsafe" + + assert chroma_langchain_collection_kwargs() == { + "collection_configuration": {"embedding_function": None}, + } + + @pytest.mark.api_key_required class TestChromaVectorStoreComponent(ComponentTestBaseWithoutClient): @pytest.fixture diff --git a/src/backend/tests/unit/test_knowledge_bases_api.py b/src/backend/tests/unit/test_knowledge_bases_api.py index 831b6eed53..2bf0fe4777 100644 --- a/src/backend/tests/unit/test_knowledge_bases_api.py +++ b/src/backend/tests/unit/test_knowledge_bases_api.py @@ -223,6 +223,11 @@ class TestKnowledgeBaseAPI: assert data["name"] == "New KB" assert data["backend_type"] == "opensearch" assert data["backend_config"] == {"index_name": "new_kb_index"} + mock_fresh_client.return_value.create_collection.assert_called_once_with( + name=kb_name, + configuration={"embedding_function": None}, + embedding_function=None, + ) record = await knowledge_base_service.get_by_user_and_name(active_user.id, kb_name) assert record is not None assert record.model_selection == model_selection diff --git a/src/lfx/src/lfx/base/knowledge_bases/backends/chroma.py b/src/lfx/src/lfx/base/knowledge_bases/backends/chroma.py index 75eb2009c2..b08c137c52 100644 --- a/src/lfx/src/lfx/base/knowledge_bases/backends/chroma.py +++ b/src/lfx/src/lfx/base/knowledge_bases/backends/chroma.py @@ -36,6 +36,7 @@ from lfx.base.knowledge_bases.backends.base import ( IngestedDocument, TestConnectionResult, ) +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.log.logger import logger if TYPE_CHECKING: @@ -110,6 +111,7 @@ class ChromaLocalBackend(BaseVectorStoreBackend): client=self._client, collection_name=self.kb_name, embedding_function=self.embedding_function, + **chroma_langchain_collection_kwargs(), ) # ---- overrides -------------------------------------------------------- @@ -304,6 +306,7 @@ class ChromaCloudBackend(BaseVectorStoreBackend): client=self._client, collection_name=self.kb_name, embedding_function=self.embedding_function, + **chroma_langchain_collection_kwargs(), ) # ---- overrides -------------------------------------------------------- diff --git a/src/lfx/src/lfx/base/vectorstores/chroma_security.py b/src/lfx/src/lfx/base/vectorstores/chroma_security.py new file mode 100644 index 0000000000..a1bba796f6 --- /dev/null +++ b/src/lfx/src/lfx/base/vectorstores/chroma_security.py @@ -0,0 +1,28 @@ +"""Safe defaults for Langflow-created Chroma collections. + +Chroma collection configuration can persist server-side embedding functions. +Some embedding functions accept model-loading kwargs such as trust_remote_code. +Langflow embeds client-side, so built-in collection creation should not register +server-side embedding code. +""" + +from __future__ import annotations + +from typing import Any + + +def chroma_langchain_collection_kwargs() -> dict[str, Any]: + """Return LangChain Chroma kwargs that disable Chroma server-side embedding code.""" + return {"collection_configuration": _chroma_collection_configuration_without_embedding_function()} + + +def chroma_client_create_collection_kwargs() -> dict[str, Any]: + """Return chromadb client kwargs that create collections without server-side embedding code.""" + return { + "configuration": _chroma_collection_configuration_without_embedding_function(), + "embedding_function": None, + } + + +def _chroma_collection_configuration_without_embedding_function() -> dict[str, Any]: + return {"embedding_function": None} diff --git a/src/lfx/src/lfx/components/chroma/chroma.py b/src/lfx/src/lfx/components/chroma/chroma.py index 3f2633f39b..2a02fb3158 100644 --- a/src/lfx/src/lfx/components/chroma/chroma.py +++ b/src/lfx/src/lfx/components/chroma/chroma.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING from langchain_chroma import Chroma from typing_extensions import override +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store from lfx.base.vectorstores.utils import chroma_collection_to_data from lfx.inputs.inputs import BoolInput, DropdownInput, HandleInput, IntInput, StrInput @@ -119,6 +120,7 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent): client=client, embedding_function=self.embedding, collection_name=self.collection_name, + **chroma_langchain_collection_kwargs(), ) except Exception as e: if isinstance(e, ChromaError): diff --git a/src/lfx/src/lfx/components/files_and_knowledge/knowledge.py b/src/lfx/src/lfx/components/files_and_knowledge/knowledge.py index c252293814..e6500fecf5 100644 --- a/src/lfx/src/lfx/components/files_and_knowledge/knowledge.py +++ b/src/lfx/src/lfx/components/files_and_knowledge/knowledge.py @@ -39,6 +39,7 @@ from lfx.base.knowledge_bases.ingestion_sources.base import ( from lfx.base.knowledge_bases.ingestion_sources.flow_component import FlowComponentSource from lfx.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases from lfx.base.models.unified_models import get_embedding_model_options, get_embeddings +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.components.files_and_knowledge._kb_paths import ( get_knowledge_bases_root_path as _get_knowledge_bases_root_path, ) @@ -1030,6 +1031,7 @@ class KnowledgeComponent(Component): chroma = Chroma( persist_directory=str(kb_path), collection_name=self.knowledge_base, + **chroma_langchain_collection_kwargs(), ) all_docs = chroma.get() diff --git a/src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py b/src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py index 344eba9a37..4b549f9f71 100644 --- a/src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py +++ b/src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py @@ -20,6 +20,7 @@ from langflow.services.database.models.user.crud import get_user_by_id from langflow.services.memory_base.kb_path_helpers import hash_session_id, validate_kb_path from sqlmodel import select +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.components.files_and_knowledge._kb_paths import ( get_knowledge_bases_root_path, load_kb_metadata, @@ -234,6 +235,7 @@ class MemoryBaseComponent(Component): persist_directory=str(kb_path), embedding_function=embedding_function, collection_name=kb_name, + **chroma_langchain_collection_kwargs(), ) def _format_results(self, results: list[tuple]) -> DataFrame: diff --git a/src/lfx/src/lfx/components/vectorstores/local_db.py b/src/lfx/src/lfx/components/vectorstores/local_db.py index 0fa8dc9f80..7255a54cb6 100644 --- a/src/lfx/src/lfx/components/vectorstores/local_db.py +++ b/src/lfx/src/lfx/components/vectorstores/local_db.py @@ -4,6 +4,7 @@ from pathlib import Path from langchain_chroma import Chroma from typing_extensions import override +from lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs from lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store from lfx.base.vectorstores.utils import chroma_collection_to_data from lfx.inputs.inputs import MultilineInput @@ -226,6 +227,7 @@ class LocalDBComponent(LCVectorStoreComponent): client=None, embedding_function=self.embedding, collection_name=self.collection_name, + **chroma_langchain_collection_kwargs(), ) self._add_documents_to_vector_store(chroma)