feat: Unified KBs and MBs to use common vector db backends - #14239
Open
dkaushik94 wants to merge 4 commits into
Open
feat: Unified KBs and MBs to use common vector db backends#14239dkaushik94 wants to merge 4 commits into
dkaushik94 wants to merge 4 commits into
Conversation
Contributor
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughMemory Base and Knowledge Base flows now persist backend configuration in database records, resolve vector stores through shared backend abstractions, reconcile legacy rows at startup, and support backend selection in the Memory creation UI. Chroma-specific ingestion, retrieval, filtering, and metadata paths are replaced with backend-neutral operations. ChangesBackend selection and persistence
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CreateMemoryModal
participant MemoryBaseService
participant KnowledgeBaseRecord
participant VectorBackend
CreateMemoryModal->>MemoryBaseService: submit backend type and config
MemoryBaseService->>VectorBackend: initialize configured backend
MemoryBaseService->>KnowledgeBaseRecord: create backing record
MemoryBaseService-->>CreateMemoryModal: create result
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 7 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (7 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
✅ Test Coverage AdvisorNo source changes detected without accompanying tests. Thanks for keeping coverage up! 🎉
|
dkaushik94
force-pushed
the
mb-refactor-with-kb
branch
from
July 23, 2026 23:36
e39d92c to
5560886
Compare
Contributor
|
✅ Migration Validation Passed All migrations follow the Expand-Contract pattern correctly. |
dkaushik94
force-pushed
the
mb-refactor-with-kb
branch
from
July 23, 2026 23:39
c6321b5 to
5d45b78
Compare
Contributor
|
Build successful! ✅ |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## release-1.12.0 #14239 +/- ##
==================================================
- Coverage 61.52% 61.49% -0.04%
==================================================
Files 2349 2349
Lines 238463 238649 +186
Branches 33541 33567 +26
==================================================
+ Hits 146719 146750 +31
- Misses 89929 90084 +155
Partials 1815 1815
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/backend/base/langflow/services/memory_base/service.py (1)
176-207: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftDuplicate Memory Base name leaves an orphaned
knowledge_baserow and provisioned collection.
initialize_kband_create_kb_record_for_memory_baserun at step 4, but the name-uniqueness check happens afterward at Lines 201–207 (and theIntegrityErrorpath at 217–221). Becausekb_namecarries a random hex suffix, the KB row never collides, so a duplicate-name creation provisions a collection and inserts an authoritativeknowledge_baserow, then raisesValueErroron theMemoryBaseinsert — leaving both behind with noMemoryBasereferencing them and no cleanup.Move the existence pre-check ahead of provisioning so the common duplicate case fails fast; for the race that only surfaces as
IntegrityError, also tear down the provisioned KB row/collection before re-raising.🔧 Sketch: pre-check before provisioning
+ # Fail fast on duplicate name BEFORE provisioning a collection / KB row, + # otherwise the row (random kb_name, never collides) is orphaned on raise. + async with session_scope() as db: + existing = await db.exec( + select(MemoryBase).where(MemoryBase.user_id == user_id).where(MemoryBase.name == payload.name) + ) + if existing.first() is not None: + msg = f"A Memory Base named '{payload.name}' already exists for this user" + raise ValueError(msg) + # 4. Provision the backing KB: the vector-store collection plus its # ``knowledge_base`` row. ... - embedding_provider = infer_embedding_provider(payload.embedding_model) await initialize_kb(...)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/backend/base/langflow/services/memory_base/service.py` around lines 176 - 207, Move the existing MemoryBase name-uniqueness query in the creation flow before `initialize_kb` and `_create_kb_record_for_memory_base`, raising the same `ValueError` before provisioning on duplicates. For the concurrent-creation path caught by `IntegrityError`, invoke the established KB cleanup mechanism to remove the provisioned `knowledge_base` row and collection, then re-raise the original error.
🧹 Nitpick comments (3)
src/backend/tests/unit/test_memory_base_task.py (1)
232-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winIngestion-task tests never verify the DB-driven backend selection actually reaches
create_backend. Both files hardcoderesolve_backend_selectionto return("chroma", {})and mockcreate_backendwith a barereturn_value, but no test assertscreate_backendis called with that resolved value (or exercises a non-Chroma backend type) — the exact pass-through this PR is meant to unify.
src/backend/tests/unit/test_memory_base_task.py#L232-L243: capture thecreate_backendmock and assert it's called with the tuple returned byresolve_backend_selection(or parametrize with a non-"chroma" backend type) in at least one representative test (e.g.test_mark_messages_ingested_called_on_success).src/backend/tests/unit/test_memory_bases.py#L1586-L1598: apply the same assertion intest_cursor_advanced_on_success, mirroring the fix above.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/backend/tests/unit/test_memory_base_task.py` around lines 232 - 243, Update the representative ingestion-task tests so they verify backend selection is passed through to create_backend: in src/backend/tests/unit/test_memory_base_task.py lines 232-243, capture the create_backend mock and assert it receives the tuple returned by resolve_backend_selection, optionally using a non-chroma backend; apply the same assertion in src/backend/tests/unit/test_memory_bases.py lines 1586-1598 within test_cursor_advanced_on_success.Source: Path instructions
src/backend/base/langflow/api/v1/knowledge_bases.py (1)
344-378: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRedundant DB round-trip:
_assert_kb_not_memory_basere-fetches a record_guard_kb_actionalready resolved.Every one of the 8 call sites first calls
_guard_kb_action(...), which already resolves and stores the KB record on_KbGuardResult.record(owner-scoped, or the cross-user-matched candidate)._assert_kb_not_memory_base(kb_name, owner_user)then unconditionally re-queriesknowledge_base_service.get_by_user_and_name(owner_user.id, kb_name)for the same row (sinceowner_useris derived from the resolved record'suser_id), adding an avoidable DB round-trip to every guarded KB endpoint (upload, folder-ingest, get, chunks, metadata-keys, connector-ingest, delete, cancel).Accept an optional pre-fetched record and only query when the caller doesn't already have one (e.g.
_kb_guard.recordisNonefor legacy disk-only KBs).♻️ Proposed fix to avoid the extra query when the caller already has the record
-async def _assert_kb_not_memory_base(kb_name: str, owner_user) -> None: +async def _assert_kb_not_memory_base( + kb_name: str, owner_user, *, record: KnowledgeBaseRecord | None = None +) -> None: """Post-resolution memory-base check. ... """ - record = await knowledge_base_service.get_by_user_and_name(owner_user.id, kb_name) + if record is None: + record = await knowledge_base_service.get_by_user_and_name(owner_user.id, kb_name) if _record_is_memory_base_associated(record):Then at each call site (e.g. line 1075):
await _assert_kb_not_memory_base(kb_name, _kb_guard.owner_user, record=_kb_guard.record).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/backend/base/langflow/api/v1/knowledge_bases.py` around lines 344 - 378, Update _assert_kb_not_memory_base to accept an optional pre-fetched record and reuse it when provided, querying knowledge_base_service only when the record is absent. Modify all guarded endpoint call sites to pass _kb_guard.record along with _kb_guard.owner_user, preserving the existing memory-base rejection behavior and fallback lookup for legacy disk-only KBs.src/backend/tests/unit/components/files_and_knowledge/test_knowledge.py (1)
355-448: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftNew test class bypasses the mandated
ComponentTestBase*fixtures and leans heavily on mocking the DB layer.
TestBackendResolutioninstantiatesKnowledgeComponentdirectly and doesn't inherit fromComponentTestBaseWithClient/ComponentTestBaseWithoutClient, and doesn't definecomponent_class/default_kwargs/file_names_mapping. Separately,get_by_user_and_nameisAsyncMock-patched in every async test (including the "wins over sidecar" and "falls back to sidecar" happy-path tests), where a real DB-backedknowledge_baserow (via existing test fixtures) would exercise the actual persistence contract instead of a stand-inSimpleNamespace. The DB-failure test (test_database_lookup_failure_propagates) is a reasonable exception since a real transient DB error is hard to simulate with a real integration.As per coding guidelines: "Component tests must use either
ComponentTestBaseWithClientfor components needing API access orComponentTestBaseWithoutClientfor pure logic components, and must includecomponent_class,default_kwargs, andfile_names_mappingfixtures" and "Backend unit tests should avoid mocking when possible and prefer real integrations for more reliable tests."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/backend/tests/unit/components/files_and_knowledge/test_knowledge.py` around lines 355 - 448, Refactor TestBackendResolution to use the appropriate ComponentTestBase fixture, defining component_class, default_kwargs, and file_names_mapping instead of directly constructing KnowledgeComponent. Replace the AsyncMock database patches and SimpleNamespace records in the happy-path tests with real persisted knowledge_base rows through the existing fixtures; retain mocking only for test_database_lookup_failure_propagates, which specifically simulates a transient DB error.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/lfx/src/lfx/_assets/component_index.json`:
- Line 9620: Normalize filter_by_session through a shared boolean conversion
that recognizes string values such as "false" as false, then use the normalized
result in both _build_where_clause() and retrieve_memory(). Ensure session_id
validation and where-clause construction follow the actual user-selected boolean
state rather than Python truthiness.
- Line 8118: The retrieval flow currently limits client-side metadata filtering
to a fixed top_k multiplier, which can omit valid lower-ranked matches. Update
the retrieval method that handles metadata_filter to pass the filter to backends
supporting native filtering; for unsupported backends, continue paginating or
expanding unfiltered fetches until top_k matches are collected or the store is
exhausted. Preserve existing result ordering and behavior when no filter is
provided, using _parse_metadata_filter and _chunk_matches_filter for client-side
filtering.
- Line 8118: Update _convert_df_to_data_objects so Chroma Cloud ingestion does
not perform local Chroma lookups for duplicate detection. Preserve the
distinction between local Chroma and Chroma Cloud after backend normalization,
and obtain existing IDs through the resolved backend when using Chroma Cloud.
Ensure existing_ids is populated before duplicate filtering without constructing
Chroma with a local persist_directory for cloud writes.
- Line 8118: Update KnowledgeComponent.retrieve_data() to resolve the embedding
model and configuration from the database-backed knowledge-base record before
attempting to load embedding_metadata.json or requiring the local sidecar.
Preserve sidecar loading only as a legacy fallback or when stored credentials
are required, so remote-backed knowledge bases retrieve successfully on replicas
without local sidecar files.
In `@src/lfx/src/lfx/components/files_and_knowledge/knowledge.py`:
- Around line 1565-1607: Update retrieve_data’s embedding-model metadata
resolution to query the database record via
knowledge_base_service.get_by_user_and_name first, reusing read_metadata where
appropriate for model_selection, chunk_size, chunk_overlap, and separator. If no
record exists, preserve the existing embedding_metadata.json sidecar fallback
for legacy knowledge bases, but do not require local metadata when a database
row is available.
In `@src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py`:
- Around line 216-245: Update _build_backend to clean up the constructed backend
when await backend.ensure_ready() raises. Wrap the readiness call and return
path with exception handling that invokes backend.teardown() on failure, then
re-raises the original exception so retrieve_memory’s existing finally behavior
remains unchanged for successfully returned backends.
- Around line 216-245: Backend setup failures can leave partially initialized
vector stores undisposed. In
src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py#L216-L245,
update _build_backend to tear down the created backend if ensure_ready() fails,
then re-raise; regenerate or re-sync the identical embedded _build_backend
source in src/backend/base/langflow/initial_setup/starter_projects/Memory
Chatbot.json#L1628-L1628 and Price Deal Finder.json#L2288-L2288. In
src/backend/base/langflow/initial_setup/starter_projects/Vector Store
RAG.json#L1314-L1314, update KnowledgeComponent._create_vector_store to tear
down the backend on failure across create_backend, ensure_ready, iter_documents,
or add_documents before re-raising.
---
Outside diff comments:
In `@src/backend/base/langflow/services/memory_base/service.py`:
- Around line 176-207: Move the existing MemoryBase name-uniqueness query in the
creation flow before `initialize_kb` and `_create_kb_record_for_memory_base`,
raising the same `ValueError` before provisioning on duplicates. For the
concurrent-creation path caught by `IntegrityError`, invoke the established KB
cleanup mechanism to remove the provisioned `knowledge_base` row and collection,
then re-raise the original error.
---
Nitpick comments:
In `@src/backend/base/langflow/api/v1/knowledge_bases.py`:
- Around line 344-378: Update _assert_kb_not_memory_base to accept an optional
pre-fetched record and reuse it when provided, querying knowledge_base_service
only when the record is absent. Modify all guarded endpoint call sites to pass
_kb_guard.record along with _kb_guard.owner_user, preserving the existing
memory-base rejection behavior and fallback lookup for legacy disk-only KBs.
In `@src/backend/tests/unit/components/files_and_knowledge/test_knowledge.py`:
- Around line 355-448: Refactor TestBackendResolution to use the appropriate
ComponentTestBase fixture, defining component_class, default_kwargs, and
file_names_mapping instead of directly constructing KnowledgeComponent. Replace
the AsyncMock database patches and SimpleNamespace records in the happy-path
tests with real persisted knowledge_base rows through the existing fixtures;
retain mocking only for test_database_lookup_failure_propagates, which
specifically simulates a transient DB error.
In `@src/backend/tests/unit/test_memory_base_task.py`:
- Around line 232-243: Update the representative ingestion-task tests so they
verify backend selection is passed through to create_backend: in
src/backend/tests/unit/test_memory_base_task.py lines 232-243, capture the
create_backend mock and assert it receives the tuple returned by
resolve_backend_selection, optionally using a non-chroma backend; apply the same
assertion in src/backend/tests/unit/test_memory_bases.py lines 1586-1598 within
test_cursor_advanced_on_success.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 0e38af72-fbe2-4f27-81ce-75ce166faa4a
📒 Files selected for processing (32)
.gitignoresrc/backend/base/langflow/api/utils/kb_helpers.pysrc/backend/base/langflow/api/utils/knowledge_base_service.pysrc/backend/base/langflow/api/v1/knowledge_bases.pysrc/backend/base/langflow/initial_setup/starter_projects/Custom Component Generator.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Document Q&A.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Knowledge Retrieval.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Meeting Summary.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Memory Chatbot.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Price Deal Finder.jsonsrc/backend/base/langflow/initial_setup/starter_projects/Vector Store RAG.jsonsrc/backend/base/langflow/main.pysrc/backend/base/langflow/services/database/models/memory_base/model.pysrc/backend/base/langflow/services/memory_base/document_builders.pysrc/backend/base/langflow/services/memory_base/ingestion.pysrc/backend/base/langflow/services/memory_base/kb_path_helpers.pysrc/backend/base/langflow/services/memory_base/service.pysrc/backend/base/langflow/services/memory_base/task.pysrc/backend/tests/unit/base/knowledge_bases/test_connector_endpoints.pysrc/backend/tests/unit/components/files_and_knowledge/test_knowledge.pysrc/backend/tests/unit/components/files_and_knowledge/test_memory_retrieval.pysrc/backend/tests/unit/test_memory_base_task.pysrc/backend/tests/unit/test_memory_bases.pysrc/frontend/src/controllers/API/queries/memories/types.tssrc/frontend/src/locales/en.jsonsrc/frontend/src/modals/createMemoryModal/__tests__/useCreateMemoryModal.test.tsxsrc/frontend/src/modals/createMemoryModal/index.tsxsrc/frontend/src/modals/createMemoryModal/useCreateMemoryModal.tssrc/lfx/src/lfx/_assets/component_index.jsonsrc/lfx/src/lfx/base/knowledge_bases/backends/opensearch.pysrc/lfx/src/lfx/components/files_and_knowledge/knowledge.pysrc/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py
| @@ -8115,7 +8115,7 @@ | |||
| "show": true, | |||
| "title_case": false, | |||
| "type": "code", | |||
| "value": "\"\"\"Unified Knowledge component — ingest into or retrieve from a knowledge base.\n\nThis component merges what used to live in ``ingestion.py`` and\n``retrieval.py`` into a single, mode-driven component. A ``TabInput``\n(\"📥 Ingest\" / \"🔍 Retrieve\") drives which inputs and which output are\nvisible. The merged shape gives users one node per KB instead of two\nparallel ones that always shared the same ``knowledge_base`` picker,\nembedding-model metadata, and backend registry.\n\nBoth legacy classes (``KnowledgeIngestionComponent``,\n``KnowledgeBaseComponent``) remain importable as thin subclasses of\n``KnowledgeComponent`` — see the sibling ``ingestion.py`` / ``retrieval.py``\nmodules — so saved flows continue to load unchanged.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport asyncio\nimport hashlib\nimport json\nimport re\nimport uuid\nfrom dataclasses import asdict, dataclass, field\nfrom datetime import datetime, timezone\nfrom typing import TYPE_CHECKING, Any\n\nimport pandas as pd\nfrom cryptography.fernet import InvalidToken\nfrom langchain_chroma import Chroma\nfrom langflow.services.auth.utils import decrypt_api_key, encrypt_api_key\n\nfrom lfx.base.knowledge_bases.backends import BackendType, BaseVectorStoreBackend, create_backend\nfrom lfx.base.knowledge_bases.ingestion_sources.base import (\n IngestionItemResult,\n IngestionItemStatus,\n IngestionRunStatus,\n IngestionSummary,\n)\nfrom lfx.base.knowledge_bases.ingestion_sources.flow_component import FlowComponentSource\nfrom lfx.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases\nfrom lfx.base.models.unified_models import get_embedding_model_options, get_embeddings\nfrom lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs\nfrom lfx.components.files_and_knowledge._kb_paths import (\n KBKeyDecryptError,\n load_kb_metadata,\n)\nfrom lfx.components.files_and_knowledge._kb_paths import (\n get_knowledge_bases_root_path as _get_knowledge_bases_root_path,\n)\nfrom lfx.components.processing.converter import convert_to_dataframe\nfrom lfx.custom import Component\nfrom lfx.io import (\n BoolInput,\n DBProviderInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MessageTextInput,\n ModelInput,\n Output,\n SecretStrInput,\n StrInput,\n TabInput,\n TableInput,\n)\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.schema.dotdict import dotdict\nfrom lfx.schema.table import EditMode\nfrom lfx.services.deps import (\n get_settings_service,\n session_scope,\n)\nfrom lfx.utils.component_utils import set_current_fields, set_field_display\nfrom lfx.utils.validate_cloud import raise_error_if_astra_cloud_disable_component\n\n\ndef _inputs_for_mode(default_mode: str) -> list:\n \"\"\"Return a fresh copy of the canonical inputs list with show flags set for the given mode.\n\n Used by the legacy subclasses so a saved flow keyed on\n ``KnowledgeIngestion`` / ``KnowledgeBase`` lands on a node template\n whose default visibility already matches the pinned mode — no\n \"wait for the user to click the mode tab\" UX gap on load.\n \"\"\"\n always_visible = set(KnowledgeComponent.default_keys) | set(KnowledgeComponent.mode_config[default_mode])\n inputs_copy = []\n for inp in KnowledgeComponent.inputs:\n clone = inp.model_copy(deep=True) if hasattr(inp, \"model_copy\") else inp\n if clone.name == \"mode\":\n clone.value = default_mode\n else:\n clone.show = clone.name in always_visible\n inputs_copy.append(clone)\n return inputs_copy\n\n\nif TYPE_CHECKING:\n from pathlib import Path\n\n# Mode constants. Plain-text labels (no emoji) for consistency with the rest of\n# Langflow's TabInput palette — see ``MemoryComponent`` for the same convention.\nMODE_INGEST = \"Ingest\"\nMODE_RETRIEVE = \"Retrieve\"\n\n\ndef _is_retrieve_mode(value: Any) -> bool:\n \"\"\"Lenient mode check: treats any label containing 'Retrieve' as retrieve mode.\n\n Older saved flows may carry the emoji-prefixed labels (\"📥 Ingest\" /\n \"🔍 Retrieve\") this component used to ship with; substring matching\n keeps those loading without forcing a flow rewrite.\n \"\"\"\n return isinstance(value, str) and \"Retrieve\" in value\n\n\n# Error message used by both the ingest and retrieve paths when the user is\n# running against an Astra cloud environment that disables these flows.\nastra_error_msg = \"Knowledge ingestion and retrieval are not supported in Astra cloud environment.\"\n\n_DEFAULT_OPENSEARCH_CONFIG = {\n \"url_variable\": \"OPENSEARCH_URL\",\n \"username_variable\": \"OPENSEARCH_USERNAME\",\n \"password_variable\": \"OPENSEARCH_PASSWORD\", # pragma: allowlist secret\n \"index_name\": \"\",\n \"vector_field\": \"vector_field\",\n \"text_field\": \"text\",\n}\n\n_DEFAULT_CHROMA_CLOUD_CONFIG = {\n \"mode\": \"cloud\",\n \"tenant_variable\": \"CHROMA_TENANT\",\n \"database_variable\": \"CHROMA_DATABASE\",\n \"api_key_variable\": \"CHROMA_API_KEY\", # pragma: allowlist secret\n}\n\n\nclass KnowledgeComponent(Component):\n \"\"\"One component for both writing into and reading from a Langflow knowledge base.\n\n A ``TabInput`` switches between ingestion and retrieval. The\n ``update_build_config`` / ``update_outputs`` hooks hide the inputs\n and the output that don't apply to the current mode so the canvas\n node stays focused.\n \"\"\"\n\n display_name = \"Knowledge\"\n description = \"Ingest into or retrieve from a Langflow knowledge base.\"\n icon = \"database\"\n name = \"Knowledge\"\n\n # ------ Mode → visible-fields wiring ---------------------------------\n # ``default_keys`` are inputs always visible regardless of mode.\n # ``mode_config`` lists the inputs unique to each mode; everything outside\n # this set is hidden when its mode is not selected.\n default_keys: list[str] = [\"mode\", \"knowledge_base\"]\n mode_config: dict[str, list[str]] = {\n MODE_INGEST: [\n \"input_df\",\n \"column_config\",\n \"chunk_size\",\n \"api_key\",\n \"allow_duplicates\",\n \"metadata_json\",\n ],\n MODE_RETRIEVE: [\n \"search_query\",\n \"top_k\",\n \"include_metadata\",\n \"include_embeddings\",\n \"metadata_filter\",\n ],\n }\n\n def __init__(self, *args, **kwargs) -> None:\n super().__init__(*args, **kwargs)\n self._cached_kb_path: Path | None = None\n\n @dataclass\n class NewKnowledgeBaseInput:\n functionality: str = \"create\"\n fields: dict[str, dict] = field(\n default_factory=lambda: {\n \"data\": {\n \"node\": {\n \"name\": \"create_knowledge_base\",\n \"description\": \"Create new knowledge in Langflow.\",\n \"display_name\": \"Create new Knowledge Base\",\n \"field_order\": [\n \"01_new_kb_name\",\n \"02_embedding_model\",\n \"03_knowledge_backend\",\n ],\n \"template\": {\n \"01_new_kb_name\": StrInput(\n name=\"new_kb_name\",\n display_name=\"Knowledge Name\",\n info=\"Name of the new knowledge to create.\",\n required=True,\n ),\n \"02_embedding_model\": ModelInput(\n name=\"embedding_model\",\n display_name=\"Choose Embedding Model\",\n info=(\n \"Select the embedding model to use for this knowledge base. \"\n \"Langflow uses the configured credentials for that model provider.\"\n ),\n required=True,\n model_type=\"embedding\",\n ),\n \"03_knowledge_backend\": DBProviderInput(\n name=\"knowledge_backend\",\n display_name=\"DB Provider\",\n info=(\n \"Select where this knowledge base stores vectors. \"\n \"OpenSearch uses the global DB Providers settings.\"\n ),\n required=True,\n ),\n },\n },\n }\n }\n )\n\n # ------ Inputs --------------------------------------------------------\n # ``knowledge_base`` is kept at position 0 so legacy tests that check\n # ``component.inputs[0].dialog_inputs`` continue to work.\n inputs = [\n DropdownInput(\n name=\"knowledge_base\",\n display_name=\"Knowledge\",\n info=\"Select the knowledge to load data from.\",\n required=True,\n options=[],\n refresh_button=True,\n real_time_refresh=True,\n dialog_inputs=asdict(NewKnowledgeBaseInput()),\n ),\n TabInput(\n name=\"mode\",\n display_name=\"Mode\",\n options=[MODE_INGEST, MODE_RETRIEVE],\n value=MODE_INGEST,\n info=\"Switch between writing new data into the knowledge base and querying it.\",\n real_time_refresh=True,\n tool_mode=True,\n ),\n # --- Ingest-only inputs (default-shown; hidden when mode == Retrieve) -\n HandleInput(\n name=\"input_df\",\n display_name=\"Input\",\n info=(\n \"Table with all original columns (already chunked / processed). \"\n \"Accepts Message, Data, or DataFrame. If Message or Data is provided, \"\n \"it is converted to a DataFrame automatically.\"\n ),\n input_types=[\"Message\", \"Data\", \"JSON\", \"DataFrame\", \"Table\"],\n required=True,\n dynamic=True,\n show=True,\n ),\n TableInput(\n name=\"column_config\",\n display_name=\"Column Configuration\",\n info=\"Configure column behavior for the knowledge base.\",\n required=True,\n table_schema=[\n {\n \"name\": \"column_name\",\n \"display_name\": \"Column Name\",\n \"type\": \"str\",\n \"description\": \"Name of the column in the source DataFrame\",\n \"edit_mode\": EditMode.INLINE,\n },\n {\n \"name\": \"vectorize\",\n \"display_name\": \"Vectorize\",\n \"type\": \"boolean\",\n \"description\": \"Create embeddings for this column\",\n \"default\": False,\n \"edit_mode\": EditMode.INLINE,\n },\n {\n \"name\": \"identifier\",\n \"display_name\": \"Identifier\",\n \"type\": \"boolean\",\n \"description\": \"Use this column as unique identifier\",\n \"default\": False,\n \"edit_mode\": EditMode.INLINE,\n },\n ],\n value=[\n {\n \"column_name\": \"text\",\n \"vectorize\": True,\n \"identifier\": True,\n },\n ],\n dynamic=True,\n show=True,\n ),\n IntInput(\n name=\"chunk_size\",\n display_name=\"Chunk Size\",\n info=\"Batch size for processing embeddings\",\n advanced=True,\n value=1000,\n dynamic=True,\n show=True,\n ),\n SecretStrInput(\n name=\"api_key\",\n display_name=\"Embedding Provider API Key\",\n info=\"Overrides global provider settings. Leave blank to use your pre-configured API Key.\",\n advanced=True,\n required=False,\n dynamic=True,\n show=True,\n ),\n BoolInput(\n name=\"allow_duplicates\",\n display_name=\"Allow Duplicates\",\n info=\"Allow duplicate rows in the knowledge base\",\n advanced=True,\n value=False,\n dynamic=True,\n show=True,\n ),\n StrInput(\n name=\"metadata_json\",\n display_name=\"Metadata\",\n info=(\n \"Optional JSON object of user metadata applied to every chunk produced by this \"\n 'run (e.g. {\"tag\": \"invoice\", \"year\": \"2026\"}). Same shape as the upload modal '\n \"Metadata section so chunks browser filters + Knowledge retrieval metadata_filter \"\n \"work uniformly across upload, folder, and flow-driven ingestion. Malformed JSON is \"\n \"ignored with a warning rather than failing the run.\"\n ),\n advanced=True,\n required=False,\n dynamic=True,\n show=True,\n ),\n # --- Retrieve-only inputs (default-hidden; shown when mode == Retrieve) -\n MessageTextInput(\n name=\"search_query\",\n display_name=\"Search Query\",\n info=\"Optional search query to filter knowledge base data.\",\n tool_mode=True,\n dynamic=True,\n show=False,\n ),\n IntInput(\n name=\"top_k\",\n display_name=\"Top K Results\",\n info=\"Number of top results to return from the knowledge base.\",\n value=5,\n advanced=True,\n required=False,\n dynamic=True,\n show=False,\n ),\n BoolInput(\n name=\"include_metadata\",\n display_name=\"Include Metadata\",\n info=\"Whether to include all metadata in the output. If false, only content is returned.\",\n value=True,\n advanced=False,\n dynamic=True,\n show=False,\n ),\n BoolInput(\n name=\"include_embeddings\",\n display_name=\"Include Embeddings\",\n info=\"Whether to include embeddings in the output. Only applicable if 'Include Metadata' is enabled.\",\n value=False,\n advanced=True,\n dynamic=True,\n show=False,\n ),\n MessageTextInput(\n name=\"metadata_filter\",\n display_name=\"Metadata Filter\",\n info=(\n \"Optional JSON object of user-metadata key/value pairs. Only chunks \"\n 'whose source_metadata matches every key are returned (e.g. {\"tag\": \"invoice\"} '\n 'or {\"tag\": [\"invoice\", \"audit\"]} for OR-of-values). Backends without '\n \"native filtering apply the match client-side after retrieval.\"\n ),\n advanced=True,\n dynamic=True,\n show=False,\n ),\n ]\n\n # ------ Outputs -------------------------------------------------------\n # Both outputs are declared at the class level so the runtime can\n # dispatch to either method depending on which one the saved flow has\n # wired up. ``update_outputs`` filters the canvas-visible output per\n # selected mode; see ``TypeConverterComponent`` and ``MemoryComponent``\n # for the same pattern.\n #\n # Output names match the legacy ``KnowledgeIngestionComponent``\n # (``dataframe_output``) and ``KnowledgeBaseComponent``\n # (``retrieve_data``) so saved flow edges keyed on those names resolve\n # cleanly against the merged component.\n outputs = [\n Output(\n display_name=\"Results\",\n name=\"dataframe_output\",\n method=\"build_kb_info\",\n types=[\"JSON\"],\n selected=\"JSON\",\n ),\n Output(\n display_name=\"Results\",\n name=\"retrieve_data\",\n method=\"retrieve_data\",\n info=\"Returns the data from the selected knowledge base.\",\n types=[\"Table\"],\n selected=\"Table\",\n ),\n ]\n\n # ------ Mode-driven UI updates ---------------------------------------\n async def update_frontend_node(self, new_frontend_node: dict, current_frontend_node: dict):\n \"\"\"Sync the visible output with the current ``mode`` value on canvas load.\n\n Saved flows hit this path: ``update_outputs`` is normally only triggered\n by ``real_time_refresh`` field edits, so without this re-sync the\n canvas could land with both outputs visible.\n \"\"\"\n await super().update_frontend_node(new_frontend_node, current_frontend_node)\n mode_value = new_frontend_node.get(\"template\", {}).get(\"mode\", {}).get(\"value\", MODE_INGEST)\n self.update_outputs(new_frontend_node, \"mode\", mode_value)\n return new_frontend_node\n\n def update_outputs(self, frontend_node: dict, field_name: str, field_value: Any) -> dict:\n \"\"\"Filter visible outputs to match the selected mode.\n\n Triggered by the ``mode`` ``TabInput`` ``real_time_refresh`` flag.\n The class-level ``outputs`` list always carries both entries so the\n runtime can resolve either ``build_kb_info`` or ``retrieve_data``\n regardless of canvas visibility — we just hide the unused one here.\n \"\"\"\n if field_name != \"mode\":\n return frontend_node\n if _is_retrieve_mode(field_value):\n frontend_node[\"outputs\"] = [\n Output(\n display_name=\"Results\",\n name=\"retrieve_data\",\n method=\"retrieve_data\",\n info=\"Returns the data from the selected knowledge base.\",\n types=[\"Table\"],\n selected=\"Table\",\n )\n ]\n else:\n frontend_node[\"outputs\"] = [\n Output(\n display_name=\"Results\",\n name=\"dataframe_output\",\n method=\"build_kb_info\",\n types=[\"JSON\"],\n selected=\"JSON\",\n )\n ]\n return frontend_node\n\n async def update_build_config(\n self,\n build_config,\n field_value: Any,\n field_name: str | None = None,\n ):\n \"\"\"Refresh KB options, drive the create-KB dialog, and hide off-mode fields.\"\"\"\n # Astra-cloud gate covers both ingest and retrieve paths.\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n # Always populate the create-KB dialog's embedding-model options so the\n # ModelInput renders correctly regardless of which input triggered the\n # refresh.\n try:\n dialog_template = (\n build_config[\"knowledge_base\"]\n .get(\"dialog_inputs\", {})\n .get(\"fields\", {})\n .get(\"data\", {})\n .get(\"node\", {})\n .get(\"template\", {})\n )\n if \"02_embedding_model\" in dialog_template:\n embedding_options = get_embedding_model_options(user_id=self.user_id)\n dialog_template[\"02_embedding_model\"][\"options\"] = embedding_options\n except Exception: # noqa: BLE001\n self.log(\"Failed to populate embedding model options in dialog\")\n\n # KB-picker refresh + create-new-KB flow (lifted from the legacy ingestion\n # component verbatim; relied on by both the canvas refresh button and the\n # dialog-submit path).\n if field_name == \"knowledge_base\":\n # Lazy import keeps lfx importable without langflow installed.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching knowledge base list.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n if isinstance(field_value, dict) and \"01_new_kb_name\" in field_value:\n if not self.is_valid_collection_name(field_value[\"01_new_kb_name\"]):\n msg = f\"Invalid knowledge base name: {field_value['01_new_kb_name']}\"\n raise ValueError(msg)\n\n model_selection = field_value[\"02_embedding_model\"]\n if isinstance(model_selection, dict):\n model_selection = [model_selection]\n\n backend_type, backend_config = self._normalize_backend_selection(\n field_value.get(\"03_knowledge_backend\")\n )\n\n embed_model = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n )\n\n try:\n await asyncio.wait_for(\n asyncio.to_thread(embed_model.embed_query, \"test\"),\n timeout=10,\n )\n except TimeoutError as e:\n msg = \"Embedding validation timed out. Please verify network connectivity and key.\"\n raise ValueError(msg) from e\n except Exception as e:\n msg = f\"Embedding validation failed: {e!s}\"\n raise ValueError(msg) from e\n\n kb_path = self._resolve_kb_path(\n _get_knowledge_bases_root_path(), kb_user, field_value[\"01_new_kb_name\"]\n )\n kb_path.mkdir(parents=True, exist_ok=True)\n\n build_config[\"knowledge_base\"][\"value\"] = field_value[\"01_new_kb_name\"]\n self._save_embedding_metadata(\n kb_path=kb_path,\n model_selection=model_selection,\n backend_type=backend_type,\n backend_config=backend_config,\n )\n await self._create_knowledge_base_record(\n user_id=self.user_id,\n name=field_value[\"01_new_kb_name\"],\n model_selection=model_selection,\n backend_type=backend_type,\n backend_config=backend_config,\n )\n\n build_config[\"knowledge_base\"][\"options\"] = await get_knowledge_bases(\n _get_knowledge_bases_root_path(),\n user_id=self.user_id,\n )\n if build_config[\"knowledge_base\"][\"value\"] not in build_config[\"knowledge_base\"][\"options\"]:\n build_config[\"knowledge_base\"][\"value\"] = None\n\n # Honor the current mode regardless of which field triggered the refresh.\n # Falls back to MODE_INGEST when ``mode`` is missing (legacy nodes).\n current_mode = build_config.get(\"mode\", {}).get(\"value\") if isinstance(build_config, dict) else None\n if field_name == \"mode\":\n current_mode = field_value\n # Map legacy/emoji-prefixed labels onto the current canonical values so\n # flows saved before the label change still toggle visibility correctly.\n if _is_retrieve_mode(current_mode):\n current_mode = MODE_RETRIEVE\n elif current_mode not in self.mode_config:\n current_mode = MODE_INGEST\n return set_current_fields(\n build_config=build_config if isinstance(build_config, dotdict) else dotdict(build_config),\n action_fields=self.mode_config,\n selected_action=current_mode,\n default_fields=self.default_keys,\n func=set_field_display,\n )\n\n # =====================================================================\n # INGESTION CODE PATH\n # =====================================================================\n def _get_kb_root(self) -> Path:\n \"\"\"Return the root directory for knowledge bases.\"\"\"\n return _get_knowledge_bases_root_path()\n\n @staticmethod\n def _resolve_kb_path(kb_root: Path, kb_user: str, kb_name: str) -> Path:\n \"\"\"Resolve the selected KB inside the authenticated user's directory.\"\"\"\n # Lazy import keeps lfx importable without langflow's DB services.\n from langflow.services.memory_base.kb_path_helpers import validate_kb_path\n\n user_root = kb_root / kb_user\n validate_kb_path(kb_root, user_root)\n\n kb_path = user_root / kb_name\n validate_kb_path(user_root, kb_path)\n return kb_path\n\n @staticmethod\n def _scalar_notna(value) -> bool:\n \"\"\"Check if a value is not NA, safely handling arrays and sequences.\n\n ``pd.notna`` returns an array when given an array-like input, which\n cannot be used directly in a boolean context. This helper collapses\n the result to a single scalar ``bool``.\n \"\"\"\n result = pd.notna(value)\n if hasattr(result, \"__iter__\") and not isinstance(result, str):\n import numpy as np\n\n arr = np.asarray(result)\n return arr.size > 0 and arr.all()\n return bool(result)\n\n def _validate_column_config(self, df_source: pd.DataFrame) -> list[dict[str, Any]]:\n \"\"\"Validate column configuration using Structured Output patterns.\"\"\"\n if not self.column_config:\n msg = \"Column configuration cannot be empty\"\n raise ValueError(msg)\n\n config_list = self.column_config if isinstance(self.column_config, list) else []\n\n df_columns = set(df_source.columns)\n for config in config_list:\n col_name = config.get(\"column_name\")\n if col_name not in df_columns:\n msg = f\"Column '{col_name}' not found in DataFrame. Available columns: {sorted(df_columns)}\"\n raise ValueError(msg)\n\n return config_list\n\n def _build_embedding_metadata(\n self,\n model_selection: list[dict[str, Any]],\n api_key: str | None = None,\n backend_type: str = BackendType.CHROMA.value,\n backend_config: dict[str, Any] | None = None,\n ) -> dict[str, Any]:\n \"\"\"Build embedding model metadata from a model selection dict.\"\"\"\n model_dict = model_selection[0] if isinstance(model_selection, list) else model_selection\n embedding_model = model_dict.get(\"name\", \"\")\n embedding_provider = model_dict.get(\"provider\", \"Unknown\")\n\n api_key_to_save = None\n if api_key and hasattr(api_key, \"get_secret_value\"):\n api_key_to_save = api_key.get_secret_value()\n elif isinstance(api_key, str):\n api_key_to_save = api_key\n\n encrypted_api_key = None\n if api_key_to_save:\n settings_service = get_settings_service()\n try:\n encrypted_api_key = encrypt_api_key(api_key_to_save, settings_service=settings_service)\n except (TypeError, ValueError) as e:\n self.log(f\"Could not encrypt API key: {e}\")\n\n return {\n \"embedding_provider\": embedding_provider,\n \"embedding_model\": embedding_model,\n \"model_selection\": model_dict,\n \"api_key\": encrypted_api_key,\n \"api_key_used\": bool(api_key),\n \"chunk_size\": self.chunk_size,\n \"backend_type\": backend_type,\n \"backend_config\": backend_config or {},\n \"created_at\": datetime.now(timezone.utc).isoformat(),\n }\n\n def _save_embedding_metadata(\n self,\n kb_path: Path,\n model_selection: list[dict[str, Any]],\n api_key: str | None = None,\n backend_type: str | None = None,\n backend_config: dict[str, Any] | None = None,\n ) -> None:\n \"\"\"Save embedding model metadata.\"\"\"\n metadata_path = kb_path / \"embedding_metadata.json\"\n existing_metadata: dict[str, Any] = {}\n if metadata_path.exists():\n try:\n existing_metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError):\n existing_metadata = {}\n\n embedding_metadata = self._build_embedding_metadata(\n model_selection,\n api_key,\n backend_type=backend_type or existing_metadata.get(\"backend_type\") or BackendType.CHROMA.value,\n backend_config=backend_config\n if backend_config is not None\n else existing_metadata.get(\"backend_config\") or {},\n )\n metadata_path.write_text(json.dumps(embedding_metadata, indent=2))\n\n def _update_metadata_metrics(self, kb_path: Path, chroma: Chroma) -> None:\n \"\"\"Update embedding_metadata.json with accurate chunk/word/character counts.\"\"\"\n import chromadb.errors\n from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper\n\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n\n try:\n metadata = json.loads(metadata_path.read_text())\n KBAnalysisHelper.update_text_metrics(kb_path, metadata, chroma)\n metadata[\"size\"] = KBStorageHelper.get_directory_size(kb_path)\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError, chromadb.errors.ChromaError) as e:\n self.log(f\"Warning: Could not update metadata metrics: {e}\")\n\n @staticmethod\n def _extract_source_types_from_df(df_source: pd.DataFrame) -> set[str]:\n \"\"\"Pull file extensions out of common path/name columns on the source DataFrame.\n\n The direct-upload ingestion path stores extensions in\n ``embedding_metadata.json[source_types]`` so the KB list can render the\n correct file-type icon. When ingestion happens via a connected\n ``input_df`` (e.g. File → Knowledge) the same field stayed empty and\n the icon defaulted to a blank tile. We look at the well-known columns\n the File / S3 / cloud-storage components produce and collect any\n plausible extension so the icon is consistent across both flows.\n \"\"\"\n candidate_columns = (\"file_path\", \"file_name\", \"filename\", \"source\", \"path\", \"mimetype\")\n extensions: set[str] = set()\n for col in candidate_columns:\n if col not in df_source.columns:\n continue\n for value in df_source[col].dropna():\n ext = KnowledgeComponent._extension_from_value(value)\n # Drop anything that doesn't look like an extension (URL\n # query strings, version segments, etc.) — the icon palette\n # keys off short alphanumeric tokens like \"pdf\"/\"docx\".\n if ext:\n extensions.add(ext)\n return extensions\n\n @staticmethod\n def _extension_from_value(value: Any) -> str | None:\n \"\"\"Return a normalized file-extension token from a path / filename / MIME string.\n\n Accepts values like ``\"report.PDF\"``, ``\"/docs/notes.txt\"``, or\n ``\"application/pdf\"`` and returns ``\"pdf\"`` / ``\"txt\"``. Returns\n ``None`` if no plausible extension can be derived.\n \"\"\"\n extension_length_limit = 10\n if value is None:\n return None\n text = str(value).strip()\n if not text:\n return None\n # MIME types like ``application/pdf`` carry the canonical extension\n # in the subtype slot — preserve them so File / S3 messages keyed\n # only on ``mimetype`` still resolve to an icon.\n if \"/\" in text and \".\" not in text.rsplit(\"/\", 1)[-1]:\n subtype = text.rsplit(\"/\", 1)[-1].strip().lower()\n return subtype if subtype and len(subtype) <= extension_length_limit and subtype.isalnum() else None\n if \".\" not in text:\n return None\n ext = text.rsplit(\".\", 1)[-1].strip().lower()\n if ext and len(ext) <= extension_length_limit and ext.isalnum():\n return ext\n return None\n\n @classmethod\n def _extract_source_types_from_mapping(cls, mapping: Any) -> set[str]:\n \"\"\"Pull extensions from a Message/Data-style ``data`` dict or a plain dict.\"\"\"\n if not isinstance(mapping, dict):\n return set()\n candidate_keys = (\"file_path\", \"file_name\", \"filename\", \"source\", \"path\", \"mimetype\")\n extensions: set[str] = set()\n for key in candidate_keys:\n ext = cls._extension_from_value(mapping.get(key))\n if ext:\n extensions.add(ext)\n return extensions\n\n @classmethod\n def _extract_source_types_from_input(cls, input_value: Any) -> set[str]:\n \"\"\"Pull file extensions out of the raw component input.\n\n ``convert_to_dataframe`` strips Message/Data fields down to ``text``\n when projecting onto a DataFrame, so file metadata attached to a\n File-component \"Raw Content\" output never reaches\n ``_extract_source_types_from_df``. Looking at the raw input first\n keeps the KB icon consistent with direct upload.\n \"\"\"\n if input_value is None:\n return set()\n if isinstance(input_value, list):\n extensions: set[str] = set()\n for item in input_value:\n extensions |= cls._extract_source_types_from_input(item)\n return extensions\n if isinstance(input_value, pd.DataFrame):\n return cls._extract_source_types_from_df(input_value)\n # Message / Data / JSON all expose a ``data`` dict via the lfx schema.\n mapping = getattr(input_value, \"data\", None)\n if mapping is None and isinstance(input_value, dict):\n mapping = input_value\n return cls._extract_source_types_from_mapping(mapping)\n\n def _merge_source_types(self, kb_path: Path, extensions: set[str]) -> None:\n \"\"\"Merge newly observed extensions into the KB's ``source_types`` metadata.\n\n Mirrors the direct-upload path in ``KBIngestionHelper`` so the icon\n rendering on the Knowledge Bases list works regardless of which\n ingestion route was used.\n \"\"\"\n if not extensions:\n return\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n try:\n metadata = json.loads(metadata_path.read_text())\n existing = set(metadata.get(\"source_types\") or [])\n metadata[\"source_types\"] = sorted(existing | extensions)\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError) as e:\n self.log(f\"Warning: Could not update source_types metadata: {e}\")\n\n async def _update_backend_metadata_metrics(self, kb_path: Path, backend: BaseVectorStoreBackend) -> None:\n \"\"\"Update metadata metrics for non-Chroma backends.\"\"\"\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n\n try:\n metadata = json.loads(metadata_path.read_text())\n chunks = await backend.count()\n characters = 0\n words = 0\n async for batch in backend.iter_documents():\n for document in batch:\n characters += len(document.content)\n words += len(document.content.split())\n\n metadata[\"chunks\"] = chunks\n metadata[\"characters\"] = characters\n metadata[\"words\"] = words\n metadata[\"avg_chunk_size\"] = characters / chunks if chunks else 0.0\n metadata[\"size\"] = await backend.storage_size_bytes()\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError) as e:\n self.log(f\"Warning: Could not update backend metadata metrics: {e}\")\n\n @staticmethod\n def _normalize_backend_selection(value: Any) -> tuple[str, dict[str, Any]]:\n \"\"\"Normalize a DBProviderInput value into backend type/config.\"\"\"\n if not value:\n return BackendType.CHROMA.value, {}\n\n if isinstance(value, str):\n backend_type = value if value == BackendType.OPENSEARCH.value else BackendType.CHROMA.value\n return (\n backend_type,\n _DEFAULT_OPENSEARCH_CONFIG.copy() if backend_type == BackendType.OPENSEARCH.value else {},\n )\n\n if not isinstance(value, dict):\n return BackendType.CHROMA.value, {}\n\n backend_type = str(value.get(\"backend_type\") or value.get(\"id\") or BackendType.CHROMA.value)\n\n if backend_type == BackendType.OPENSEARCH.value:\n backend_config = value.get(\"backend_config\") or value.get(\"config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return BackendType.OPENSEARCH.value, {**_DEFAULT_OPENSEARCH_CONFIG, **backend_config}\n\n if backend_type == \"chroma_cloud\":\n backend_config = value.get(\"backend_config\") or value.get(\"config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return BackendType.CHROMA.value, {**_DEFAULT_CHROMA_CLOUD_CONFIG, **backend_config}\n\n return BackendType.CHROMA.value, {}\n\n @staticmethod\n def _get_backend_from_metadata(kb_path: Path) -> tuple[str, dict[str, Any]]:\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return BackendType.CHROMA.value, {}\n try:\n metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError):\n return BackendType.CHROMA.value, {}\n\n backend_type = str(metadata.get(\"backend_type\") or BackendType.CHROMA.value)\n backend_config = metadata.get(\"backend_config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return backend_type, backend_config\n\n async def _create_knowledge_base_record(\n self,\n *,\n user_id: Any,\n name: str,\n model_selection: list[dict[str, Any]],\n backend_type: str,\n backend_config: dict[str, Any],\n ) -> None:\n \"\"\"Persist the component-created KB in the DB when Langflow is available.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n except ImportError:\n return\n\n try:\n await knowledge_base_service.create_record(\n user_id=user_id,\n name=name,\n model_selection=model_selection,\n column_config=self.column_config if isinstance(self.column_config, list) else [],\n backend_type=backend_type,\n backend_config=backend_config,\n )\n except Exception as exc: # noqa: BLE001\n self.log(f\"Warning: could not persist knowledge base record: {exc}\")\n\n def _save_kb_files(\n self,\n kb_path: Path,\n config_list: list[dict[str, Any]],\n ) -> None:\n \"\"\"Save KB files using File Component storage patterns.\"\"\"\n try:\n kb_path.mkdir(parents=True, exist_ok=True)\n\n cfg_path = kb_path / \"schema.json\"\n if not cfg_path.exists():\n cfg_path.write_text(json.dumps(config_list, indent=2))\n\n except (OSError, TypeError, ValueError) as e:\n self.log(f\"Error saving KB files: {e}\")\n\n def _build_column_metadata(self, config_list: list[dict[str, Any]], df_source: pd.DataFrame) -> dict[str, Any]:\n \"\"\"Build detailed column metadata.\"\"\"\n metadata: dict[str, Any] = {\n \"total_columns\": len(df_source.columns),\n \"mapped_columns\": len(config_list),\n \"unmapped_columns\": len(df_source.columns) - len(config_list),\n \"columns\": [],\n \"summary\": {\"vectorized_columns\": [], \"identifier_columns\": []},\n }\n\n for config in config_list:\n col_name = config.get(\"column_name\")\n vectorize = config.get(\"vectorize\") == \"True\" or config.get(\"vectorize\") is True\n identifier = config.get(\"identifier\") == \"True\" or config.get(\"identifier\") is True\n\n metadata[\"columns\"].append(\n {\n \"name\": col_name,\n \"vectorize\": vectorize,\n \"identifier\": identifier,\n }\n )\n\n if vectorize:\n metadata[\"summary\"][\"vectorized_columns\"].append(col_name)\n if identifier:\n metadata[\"summary\"][\"identifier_columns\"].append(col_name)\n\n return metadata\n\n async def _create_vector_store(\n self,\n df_source: pd.DataFrame,\n config_list: list[dict[str, Any]],\n embedding_function,\n ) -> BaseVectorStoreBackend:\n \"\"\"Create vector store using the configured DB provider.\"\"\"\n vector_store_dir = await self._kb_path()\n if not vector_store_dir:\n msg = \"Knowledge base path is not set. Please create a new knowledge base first.\"\n raise ValueError(msg)\n vector_store_dir.mkdir(parents=True, exist_ok=True)\n\n backend_type, backend_config = self._get_backend_from_metadata(vector_store_dir)\n backend = create_backend(\n backend_type,\n kb_name=self.knowledge_base,\n kb_path=vector_store_dir,\n backend_config=backend_config,\n embedding_function=embedding_function,\n user_id=self.user_id,\n )\n await backend.ensure_ready()\n\n existing_ids = None\n if backend_type != BackendType.CHROMA.value and not self.allow_duplicates:\n existing_ids = set()\n async for batch in backend.iter_documents():\n for document in batch:\n doc_id = document.metadata.get(\"_id\")\n if doc_id:\n existing_ids.add(doc_id)\n\n data_objects = await self._convert_df_to_data_objects(df_source, config_list, existing_ids=existing_ids)\n\n user_metadata_tag = self._resolve_user_metadata_tag()\n\n documents = []\n for data_obj in data_objects:\n doc = data_obj.to_lc_document()\n if user_metadata_tag:\n doc.metadata[\"source_metadata\"] = user_metadata_tag\n documents.append(doc)\n\n if documents:\n await backend.add_documents(documents)\n self.log(f\"Added {len(documents)} documents to vector store '{self.knowledge_base}'\")\n\n return backend\n\n async def _convert_df_to_data_objects(\n self,\n df_source: pd.DataFrame,\n config_list: list[dict[str, Any]],\n existing_ids: set[str] | None = None,\n ) -> list[Data]:\n \"\"\"Convert DataFrame to Data objects for vector store.\"\"\"\n data_objects: list[Data] = []\n\n if existing_ids is None:\n kb_path = await self._kb_path()\n\n chroma = Chroma(\n persist_directory=str(kb_path),\n collection_name=self.knowledge_base,\n **chroma_langchain_collection_kwargs(),\n )\n\n all_docs = chroma.get()\n\n existing_ids = {metadata.get(\"_id\") for metadata in all_docs[\"metadatas\"] if metadata.get(\"_id\")}\n\n content_cols = []\n identifier_cols = []\n\n for config in config_list:\n col_name = config.get(\"column_name\")\n vectorize = config.get(\"vectorize\") == \"True\" or config.get(\"vectorize\") is True\n identifier = config.get(\"identifier\") == \"True\" or config.get(\"identifier\") is True\n\n if vectorize:\n content_cols.append(col_name)\n if identifier:\n identifier_cols.append(col_name)\n\n for _, row in df_source.iterrows():\n identifier_parts = [str(row[col]) for col in content_cols if col in row and self._scalar_notna(row[col])]\n\n page_content = \" \".join(identifier_parts)\n\n data_dict = {\n \"text\": page_content,\n }\n\n if identifier_cols:\n identifier_parts = [\n str(row[col]) for col in identifier_cols if col in row and self._scalar_notna(row[col])\n ]\n page_content = \" \".join(identifier_parts)\n\n for col in df_source.columns:\n if col not in content_cols and col in row and self._scalar_notna(row[col]):\n value = row[col]\n data_dict[col] = str(value)\n\n page_content_hash = hashlib.sha256(page_content.encode()).hexdigest()\n data_dict[\"_id\"] = page_content_hash\n\n if not self.allow_duplicates and page_content_hash in existing_ids:\n self.log(f\"Skipping duplicate row with hash {page_content_hash}\")\n continue\n\n data_obj = Data(data=data_dict)\n data_objects.append(data_obj)\n\n return data_objects\n\n def is_valid_collection_name(self, name, min_length: int = 3, max_length: int = 63) -> bool:\n \"\"\"Validate collection name.\n\n 1. Contains 3-63 characters\n 2. Starts and ends with alphanumeric character\n 3. Contains only alphanumeric characters, underscores, or hyphens.\n \"\"\"\n if not (min_length <= len(name) <= max_length):\n return False\n\n if not (name[0].isalnum() and name[-1].isalnum()):\n return False\n\n return re.match(r\"^[a-zA-Z0-9_-]+$\", name) is not None\n\n async def _kb_path(self) -> Path | None:\n cached_path = getattr(self, \"_cached_kb_path\", None)\n if cached_path is not None:\n return cached_path\n\n # Lazy import to keep ``lfx`` importable standalone — langflow's\n # user/DB models are not always available at module load time.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching knowledge base path.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n\n kb_root = self._get_kb_root()\n\n self._cached_kb_path = self._resolve_kb_path(kb_root, kb_user, self.knowledge_base)\n\n return self._cached_kb_path\n\n def _resolve_user_metadata_tag(self) -> str:\n \"\"\"Return the JSON-encoded user metadata tag for chunk writes.\"\"\"\n raw = getattr(self, \"metadata_json\", None)\n if not raw:\n return \"\"\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return \"\"\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError) as exc:\n self.log(f\"KnowledgeComponent: metadata_json is not valid JSON ({exc}); skipping metadata stamp.\")\n return \"\"\n if not isinstance(decoded, dict):\n self.log(\"KnowledgeComponent: metadata_json must decode to a JSON object; skipping metadata stamp.\")\n return \"\"\n return json.dumps(decoded, sort_keys=True)\n\n async def build_kb_info(self) -> Data:\n \"\"\"Main ingestion routine → returns a dict with KB metadata.\n\n The annotation is intentionally narrowed to ``Data`` even though the\n cross-mode fallback below may return a ``DataFrame`` from\n ``retrieve_data``. The frontend builds React-Flow handle IDs from\n this output's type list; widening it to ``Data | DataFrame`` makes\n the API advertise ``[\"JSON\", \"Table\"]`` for the ingest output and\n breaks every saved-edge sourceHandle that was generated against\n a single-type handle (BUG-02). Python doesn't enforce return\n annotations at runtime, so the rare fallback path keeps working.\n \"\"\"\n if _is_retrieve_mode(getattr(self, \"mode\", MODE_INGEST)):\n return await self.retrieve_data()\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n run_id: uuid.UUID | None = None\n run_job_id: uuid.UUID | None = None\n run_summary: IngestionSummary | None = None\n run_status: IngestionRunStatus = IngestionRunStatus.SUCCEEDED\n run_error: str | None = None\n kb_record_id: uuid.UUID | None = None\n try:\n input_value = self.input_df[0] if isinstance(self.input_df, list) else self.input_df\n df_source: DataFrame = convert_to_dataframe(input_value, auto_parse=False)\n\n config_list = self._validate_column_config(df_source)\n column_metadata = self._build_column_metadata(config_list, df_source)\n\n kb_path = await self._kb_path()\n if not kb_path:\n msg = \"Knowledge base path is not set. Please create a new knowledge base first.\"\n raise ValueError(msg)\n metadata_path = kb_path / \"embedding_metadata.json\"\n api_key = None\n model_selection = None\n\n if metadata_path.exists():\n settings_service = get_settings_service()\n stored_metadata = json.loads(metadata_path.read_text())\n\n model_selection = stored_metadata.get(\"model_selection\")\n if model_selection:\n model_selection = [model_selection] if isinstance(model_selection, dict) else model_selection\n else:\n embedding_model_name = stored_metadata.get(\"embedding_model\")\n embedding_provider = stored_metadata.get(\"embedding_provider\", \"Unknown\")\n if embedding_model_name:\n try:\n all_options = get_embedding_model_options(user_id=self.user_id)\n match = next(\n (o for o in all_options if o.get(\"name\") == embedding_model_name),\n None,\n )\n if match:\n model_selection = [match]\n else:\n self.log(\n f\"Embedding model '{embedding_model_name}' (provider: {embedding_provider}) \"\n \"from stored metadata is no longer available in the model registry. \"\n \"Please re-create this knowledge base with a supported embedding model.\"\n )\n msg = (\n f\"Embedding model '{embedding_model_name}' is no longer recognized. \"\n \"The knowledge base was created with an older format and the model \"\n \"is not available in the current registry. \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg)\n except ValueError:\n raise\n except Exception: # noqa: BLE001\n self.log(\n f\"Failed to look up embedding model '{embedding_model_name}' in registry. \"\n \"Please re-create this knowledge base with a supported embedding model.\"\n )\n msg = (\n f\"Could not look up embedding model '{embedding_model_name}' \"\n f\"(provider: {embedding_provider}). \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg) # noqa: B904\n\n encrypted_key = stored_metadata.get(\"api_key\")\n if encrypted_key:\n try:\n api_key = decrypt_api_key(encrypted_key, settings_service)\n except (InvalidToken, TypeError, ValueError) as e:\n if not self.api_key:\n log_label = f\"knowledge base '{self.knowledge_base}'\"\n msg = (\n f\"Cannot decrypt the stored embedding API key for {log_label}. \"\n \"This usually means the server's SECRET_KEY changed after the \"\n \"key was saved. To recover, supply the embedding provider API \"\n \"key on the component's 'Embedding Provider API Key' input and \"\n \"re-run ingestion — the key will be re-encrypted with the \"\n \"current SECRET_KEY.\"\n )\n raise KBKeyDecryptError(msg) from e\n logger.warning(\"Stored API key undecryptable; using component-supplied key. Error: %s\", e)\n\n if self.api_key:\n api_key = self.api_key\n if model_selection:\n self._save_embedding_metadata(\n kb_path=kb_path,\n model_selection=model_selection,\n api_key=api_key,\n )\n\n if not model_selection:\n msg = \"No embedding model configuration found. Please create the knowledge base first.\"\n raise ValueError(msg)\n\n embedding_function = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n api_key=api_key,\n chunk_size=self.chunk_size,\n )\n\n run_id, run_job_id, run_summary, kb_record_id = await self._begin_ingestion_run(kb_path)\n if kb_record_id is not None:\n await self._record_kb_status(kb_record_id, \"ingesting\")\n\n backend = await self._create_vector_store(df_source, config_list, embedding_function=embedding_function)\n\n self._save_kb_files(kb_path, config_list)\n\n try:\n if not isinstance(backend, BaseVectorStoreBackend):\n pass\n elif backend.backend_type == BackendType.CHROMA and hasattr(backend, \"raw_langchain_store\"):\n self._update_metadata_metrics(kb_path, backend.raw_langchain_store())\n else:\n await self._update_backend_metadata_metrics(kb_path, backend)\n # Stamp the KB with the file extensions we just ingested so\n # the Knowledge Bases list renders the correct icon for\n # flow-driven ingestion (input_df), matching direct upload.\n # We look at the raw input first because ``convert_to_dataframe``\n # drops Message/Data metadata fields (file_path, mimetype, …)\n # when projecting onto the DataFrame.\n source_types = self._extract_source_types_from_input(input_value)\n source_types |= self._extract_source_types_from_df(df_source)\n self._merge_source_types(kb_path, source_types)\n finally:\n if isinstance(backend, BaseVectorStoreBackend):\n await backend.teardown()\n\n meta: dict[str, Any] = {\n \"kb_id\": str(uuid.uuid4()),\n \"kb_name\": self.knowledge_base,\n \"rows\": len(df_source),\n \"column_metadata\": column_metadata,\n \"path\": str(kb_path),\n \"config_columns\": len(config_list),\n \"timestamp\": datetime.now(tz=timezone.utc).isoformat(),\n }\n\n if run_summary is not None:\n run_summary.record_item(\n IngestionItemResult(\n item_id=self.knowledge_base,\n display_name=f\"{self.knowledge_base} ({len(df_source)} rows)\",\n status=IngestionItemStatus.SUCCEEDED,\n chunks_created=len(df_source),\n ),\n )\n\n if kb_record_id is not None:\n await self._record_kb_stats(kb_record_id, kb_path)\n await self._record_kb_status(kb_record_id, \"ready\")\n\n self.status = f\"✅ KB **{self.knowledge_base}** saved · {len(df_source)} chunks.\"\n\n return Data(data=meta)\n\n except (OSError, ValueError, RuntimeError, KeyError) as e:\n run_status = IngestionRunStatus.FAILED\n run_error = str(e) or e.__class__.__name__\n msg = f\"Error during KB ingestion: {e}\"\n raise RuntimeError(msg) from e\n except Exception as e:\n run_status = IngestionRunStatus.FAILED\n run_error = str(e) or e.__class__.__name__\n raise\n finally:\n if run_id is not None and run_summary is not None:\n await self._finalize_ingestion_run(\n run_id=run_id,\n job_id=run_job_id,\n summary=run_summary,\n status=run_status,\n error_message=run_error,\n )\n if kb_record_id is not None and run_status is IngestionRunStatus.FAILED:\n await self._record_kb_status(kb_record_id, \"failed\", failure_reason=run_error)\n\n async def _begin_ingestion_run(\n self,\n kb_path: Path,\n ) -> tuple[uuid.UUID | None, uuid.UUID | None, IngestionSummary | None, uuid.UUID | None]:\n \"\"\"Create a parent ``Job`` and seed an ingestion-run row.\"\"\"\n if not self.user_id:\n self.log(\"No user_id on component; skipping ingestion-run tracking.\")\n return None, None, None, None\n\n try:\n user_uuid = uuid.UUID(str(self.user_id))\n except (ValueError, TypeError) as exc:\n self.log(f\"Could not coerce user_id={self.user_id!r} to UUID; skipping run tracking ({exc}).\")\n return None, None, None, None\n\n try:\n from langflow.api.utils import ingestion_run_service, knowledge_base_service\n from langflow.services.database.models.jobs.model import JobStatus, JobType\n from langflow.services.deps import get_job_service\n except ImportError as exc:\n self.log(f\"Run-history wiring unavailable; ingestion will not be recorded ({exc}).\")\n return None, None, None, None\n\n try:\n kb_record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base)\n kb_record_id = kb_record.id if kb_record is not None else None\n\n user_metadata = self._parse_user_metadata_dict()\n\n job_id = uuid.uuid4()\n job_service = get_job_service()\n raw_flow_id = getattr(self, \"flow_id\", None)\n flow_id_uuid = uuid.UUID(str(raw_flow_id)) if raw_flow_id else job_id\n await job_service.create_job(\n job_id=job_id,\n flow_id=flow_id_uuid,\n job_type=JobType.INGESTION,\n asset_id=kb_record_id,\n asset_type=\"knowledge_base\",\n user_id=user_uuid,\n )\n await job_service.update_job_status(job_id, JobStatus.IN_PROGRESS)\n\n source = FlowComponentSource(\n user_id=user_uuid,\n source_config={\n \"knowledge_base\": self.knowledge_base,\n \"kb_path\": str(kb_path),\n \"flow_id\": str(flow_id_uuid),\n },\n )\n\n run_id = await ingestion_run_service.create_run(\n kb_name=self.knowledge_base,\n source=source,\n job_id=job_id,\n user_id=user_uuid,\n kb_id=kb_record_id,\n user_metadata=user_metadata,\n )\n if run_id is None:\n return None, job_id, None, kb_record_id\n await ingestion_run_service.mark_running(run_id)\n\n summary = IngestionSummary(\n kb_name=self.knowledge_base,\n source_type=source.source_type.value,\n user_id=user_uuid,\n job_id=job_id,\n source_config=source.describe().get(\"config\") or {},\n user_metadata=user_metadata,\n )\n self.log(f\"Started ingestion run job_id={job_id} kb_name={self.knowledge_base} kb_id={kb_record_id}\")\n except Exception as exc: # noqa: BLE001 — telemetry must never abort ingestion\n self.log(f\"Could not begin ingestion-run tracking: {exc}\")\n return None, None, None, None\n else:\n return run_id, job_id, summary, kb_record_id\n\n async def _finalize_ingestion_run(\n self,\n *,\n run_id: uuid.UUID,\n job_id: uuid.UUID | None,\n summary: IngestionSummary,\n status: IngestionRunStatus,\n error_message: str | None,\n ) -> None:\n \"\"\"Persist the final summary and transition the parent Job.\"\"\"\n try:\n from langflow.api.utils import ingestion_run_service\n from langflow.services.database.models.jobs.model import JobStatus\n from langflow.services.deps import get_job_service\n except ImportError as exc:\n self.log(f\"Run-history wiring unavailable; ingestion-run finalize skipped ({exc}).\")\n return\n\n try:\n await ingestion_run_service.finalize_run(\n run_id,\n summary=summary,\n status=status,\n error_message=error_message,\n )\n if job_id is not None:\n terminal_status = JobStatus.COMPLETED if status is not IngestionRunStatus.FAILED else JobStatus.FAILED\n await get_job_service().update_job_status(job_id, terminal_status, finished_timestamp=True)\n except Exception as exc: # noqa: BLE001 — telemetry must never re-raise\n self.log(f\"Could not finalize ingestion-run tracking: {exc}\")\n\n async def _record_kb_status(\n self,\n kb_record_id: uuid.UUID,\n status_value: str,\n *,\n failure_reason: str | None = None,\n ) -> None:\n \"\"\"Mirror Path A's KB-row status transitions.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n from langflow.services.database.models.knowledge_base.model import KnowledgeBaseStatus\n except ImportError:\n return\n try:\n await knowledge_base_service.update_status(\n kb_record_id,\n status=KnowledgeBaseStatus(status_value),\n failure_reason=failure_reason,\n )\n except Exception as exc: # noqa: BLE001\n self.log(f\"Could not update KB status to {status_value}: {exc}\")\n\n async def _record_kb_stats(self, kb_record_id: uuid.UUID, kb_path: Path) -> None:\n \"\"\"Push freshly-refreshed metrics from embedding_metadata.json onto the DB row.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n except ImportError:\n self.log(\"knowledge_base_service unavailable; KB stats will not sync to DB row.\")\n return\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n self.log(f\"No embedding_metadata.json at {metadata_path}; skipping KB stats sync.\")\n return\n try:\n metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError) as exc:\n self.log(f\"Could not read KB metadata for stats sync: {exc}\")\n return\n chunks = int(metadata.get(\"chunks\", 0) or 0)\n words = int(metadata.get(\"words\", 0) or 0)\n characters = int(metadata.get(\"characters\", 0) or 0)\n try:\n await knowledge_base_service.update_stats(\n kb_record_id,\n chunks=chunks,\n words=words,\n characters=characters,\n size_bytes=int(metadata.get(\"size\", 0) or 0),\n source_types=list(metadata.get(\"source_types\") or []),\n chunk_size=metadata.get(\"chunk_size\"),\n chunk_overlap=metadata.get(\"chunk_overlap\"),\n separator=metadata.get(\"separator\"),\n )\n self.log(f\"Synced KB stats to DB row {kb_record_id}: chunks={chunks} words={words} characters={characters}\")\n except Exception as exc: # noqa: BLE001\n self.log(f\"Could not sync KB stats to DB row {kb_record_id}: {exc}\")\n\n def _parse_user_metadata_dict(self) -> dict[str, Any]:\n \"\"\"Decode ``metadata_json`` to a dict, or ``{}`` on any error.\"\"\"\n raw = getattr(self, \"metadata_json\", None)\n if not raw:\n return {}\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return {}\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError):\n return {}\n return decoded if isinstance(decoded, dict) else {}\n\n # =====================================================================\n # RETRIEVAL CODE PATH\n # =====================================================================\n @property\n def _user_uuid(self) -> uuid.UUID | None:\n \"\"\"Return self.user_id as a UUID, converting from str if necessary.\"\"\"\n if not self.user_id:\n return None\n return self.user_id if isinstance(self.user_id, uuid.UUID) else uuid.UUID(self.user_id)\n\n def _get_kb_metadata(self, kb_path: Path, *, require_api_key: bool = False) -> dict:\n \"\"\"Load the knowledge base's embedding metadata file.\"\"\"\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n return load_kb_metadata(\n kb_path,\n log_label=f\"knowledge base '{self.knowledge_base}'\",\n require_api_key=require_api_key,\n )\n\n async def _resolve_backend(self, *, kb_user: str) -> tuple[str, dict[str, Any]]: # noqa: ARG002\n \"\"\"Return ``(backend_type, backend_config)`` for this KB.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n\n user_uuid = self._user_uuid\n if user_uuid is None:\n return BackendType.CHROMA.value, {}\n record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base)\n except Exception as exc: # noqa: BLE001\n logger.debug(\"KB record lookup failed: %s\", exc)\n return BackendType.CHROMA.value, {}\n\n if record is None:\n return BackendType.CHROMA.value, {}\n return (\n record.backend_type or BackendType.CHROMA.value,\n record.backend_config or {},\n )\n\n def _resolve_model_selection(self, metadata: dict[str, Any]) -> list[dict[str, Any]]:\n \"\"\"Resolve the ``get_embeddings``-compatible model selection from metadata.\"\"\"\n model_selection = metadata.get(\"model_selection\")\n if model_selection:\n selection_list = [model_selection] if isinstance(model_selection, dict) else list(model_selection)\n return [self._hydrate_model_metadata(entry) for entry in selection_list]\n\n embedding_model_name = metadata.get(\"embedding_model\")\n embedding_provider = metadata.get(\"embedding_provider\", \"Unknown\")\n if not embedding_model_name:\n msg = (\n f\"Knowledge base '{self.knowledge_base}' has no embedding model recorded; \"\n \"re-create it with a supported embedding model.\"\n )\n raise ValueError(msg)\n\n match = self._find_catalog_entry(embedding_model_name)\n if match is None:\n msg = (\n f\"Embedding model '{embedding_model_name}' (provider '{embedding_provider}') \"\n \"recorded for this knowledge base is no longer available in the model registry. \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg)\n return [match]\n\n def _hydrate_model_metadata(self, entry: dict[str, Any]) -> dict[str, Any]:\n \"\"\"Fill in ``metadata.embedding_class`` / ``param_mapping`` if missing.\"\"\"\n entry_metadata = entry.get(\"metadata\") or {}\n has_class = bool(entry_metadata.get(\"embedding_class\"))\n has_mapping = bool(entry_metadata.get(\"param_mapping\"))\n if has_class and has_mapping:\n return entry\n\n model_name = entry.get(\"name\")\n if not model_name:\n return entry\n\n catalog_entry = self._find_catalog_entry(model_name)\n if catalog_entry is None:\n return entry\n\n catalog_metadata = catalog_entry.get(\"metadata\") or {}\n merged_metadata = {**catalog_metadata, **entry_metadata}\n return {**entry, \"metadata\": merged_metadata}\n\n def _find_catalog_entry(self, model_name: str) -> dict[str, Any] | None:\n \"\"\"Look up an embedding model by name in the unified-models catalog.\"\"\"\n options = get_embedding_model_options(user_id=self.user_id)\n return next((o for o in options if o.get(\"name\") == model_name), None)\n\n async def retrieve_data(self) -> DataFrame:\n \"\"\"Retrieve data from the selected knowledge base.\n\n Annotation narrowed to ``DataFrame`` to keep this output's handle\n type list at ``[\"Table\"]`` only; widening to a union surfaces\n ``[\"Table\", \"JSON\"]`` on the API and breaks every starter-project\n edge whose sourceHandle was stored against a single-type handle\n (BUG-02). The cross-mode defensive fallback may still return a\n ``Data`` at runtime — Python ignores return annotations, so\n nothing breaks.\n \"\"\"\n if not _is_retrieve_mode(getattr(self, \"mode\", MODE_INGEST)):\n return await self.build_kb_info()\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n # Lazy import: langflow's user/DB models aren't part of lfx's\n # standalone install, so ``lfx run <starter>.json`` can't resolve\n # this symbol at module import time. Deferring to use keeps the\n # component importable in both environments.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching Knowledge Base data.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n kb_path = self._resolve_kb_path(_get_knowledge_bases_root_path(), kb_user, self.knowledge_base)\n\n component_api_key = self.api_key if getattr(self, \"api_key\", None) else None\n needs_stored_key = not component_api_key\n metadata = self._get_kb_metadata(kb_path, require_api_key=needs_stored_key)\n if not metadata:\n msg = f\"Metadata not found for knowledge base: {self.knowledge_base}. Ensure it has been indexed.\"\n raise ValueError(msg)\n\n api_key = component_api_key or metadata.get(\"api_key\")\n model_selection = self._resolve_model_selection(metadata)\n chunk_size = metadata.get(\"chunk_size\")\n embedding_function = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n api_key=api_key,\n chunk_size=chunk_size,\n )\n\n backend_type, backend_config = await self._resolve_backend(kb_user=kb_user)\n backend = create_backend(\n backend_type,\n kb_name=self.knowledge_base,\n kb_path=kb_path,\n backend_config=backend_config,\n embedding_function=embedding_function,\n user_id=self.user_id,\n )\n try:\n user_metadata_filter = _parse_metadata_filter(getattr(self, \"metadata_filter\", None))\n use_scores = bool(self.search_query)\n search_k = self.top_k * 4 if user_metadata_filter else self.top_k\n results = await backend.similarity_search(\n query=self.search_query or \"\",\n k=search_k,\n with_scores=use_scores,\n )\n if user_metadata_filter:\n results = [\n (doc, score) for doc, score in results if _chunk_matches_filter(doc.metadata, user_metadata_filter)\n ]\n results = results[: self.top_k]\n\n embeddings_by_key: dict[tuple[str, str], list[float]] = {}\n if self.include_embeddings and results:\n # Join each retrieved chunk to its stored embedding. We key on\n # ``_id`` when present and fall back to page content otherwise,\n # so KBs populated by direct file upload — whose chunks carry no\n # ``_id`` — still resolve. See ``_embedding_match_key``.\n wanted_keys = {_embedding_match_key(doc.page_content, doc.metadata) for doc, _score in results}\n async for batch in backend.iter_documents(include_embeddings=True):\n for entry in batch:\n if entry.embedding is None:\n continue\n key = _embedding_match_key(entry.content, entry.metadata)\n if key in wanted_keys:\n embeddings_by_key[key] = entry.embedding\n if len(embeddings_by_key) == len(wanted_keys):\n break\n\n data_list: list[Data] = []\n for doc, score in results:\n kwargs: dict[str, Any] = {\"content\": doc.page_content}\n if use_scores:\n kwargs[\"_score\"] = -1 * score\n if self.include_metadata:\n kwargs.update(doc.metadata)\n if self.include_embeddings:\n kwargs[\"_embeddings\"] = embeddings_by_key.get(_embedding_match_key(doc.page_content, doc.metadata))\n data_list.append(Data(**kwargs))\n\n return DataFrame(data=data_list)\n finally:\n await backend.teardown()\n\n\ndef _embedding_match_key(content: str, metadata: dict[str, Any] | None) -> tuple[str, str]:\n \"\"\"Build a stable key for aligning a retrieved chunk with its stored embedding.\n\n Embeddings are gathered separately (via ``iter_documents``) and then joined\n back onto the search results. Component-driven ingestion stamps a\n content-hash ``_id`` on every chunk, but direct file-upload ingestion\n (``KBIngestionHelper.perform_ingestion``) does not. Keying the join purely on\n ``_id`` therefore left every upload-populated KB with ``_embeddings: None``.\n\n We prefer ``_id`` when present (so legitimately distinct chunks that happen to\n share text aren't collapsed) and fall back to the chunk's page content\n otherwise. The content fallback is exact for the embedding use case: two\n chunks with identical text necessarily share the same embedding vector. The\n leading ``\"id\"`` / ``\"content\"`` tag namespaces the two key spaces so a mixed\n KB (some chunks with ``_id``, some without) never cross-matches.\n \"\"\"\n doc_id = metadata.get(\"_id\") if metadata else None\n if doc_id:\n return (\"id\", str(doc_id))\n return (\"content\", content or \"\")\n\n\ndef _parse_metadata_filter(raw: str | None) -> dict[str, list[str]]:\n \"\"\"Decode the ``metadata_filter`` input into a {key: [values]} map.\n\n Empty or malformed input maps to an empty filter so retrieval falls back\n to the unfiltered path. JSON errors are swallowed here rather than raised:\n surfacing component-config errors at the canvas node would break a flow\n run for what is meant to be an optional refinement.\n \"\"\"\n if not raw:\n return {}\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return {}\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError):\n logger.warning(\"KnowledgeComponent: metadata_filter is not valid JSON; ignoring filter.\")\n return {}\n if not isinstance(decoded, dict):\n logger.warning(\"KnowledgeComponent: metadata_filter must be a JSON object; ignoring filter.\")\n return {}\n result: dict[str, list[str]] = {}\n for key, value in decoded.items():\n if not isinstance(key, str):\n continue\n if isinstance(value, list):\n result[key] = [str(entry) for entry in value]\n else:\n result[key] = [str(value)]\n return result\n\n\ndef _chunk_matches_filter(metadata: dict[str, Any] | None, filt: dict[str, list[str]]) -> bool:\n \"\"\"AND across keys, OR within key values, mirroring the chunks endpoint.\"\"\"\n if not filt:\n return True\n if not metadata:\n return False\n raw = metadata.get(\"source_metadata\")\n if not raw:\n return False\n try:\n stored = json.loads(raw) if isinstance(raw, str) else raw\n except json.JSONDecodeError:\n return False\n if not isinstance(stored, dict):\n return False\n for key, expected_values in filt.items():\n actual = stored.get(key)\n if actual is None:\n return False\n actual_set = {str(entry) for entry in actual} if isinstance(actual, list) else {str(actual)}\n if not actual_set & set(expected_values):\n return False\n return True\n" | |||
| "value": "\"\"\"Unified Knowledge component — ingest into or retrieve from a knowledge base.\n\nThis component merges what used to live in ``ingestion.py`` and\n``retrieval.py`` into a single, mode-driven component. A ``TabInput``\n(\"📥 Ingest\" / \"🔍 Retrieve\") drives which inputs and which output are\nvisible. The merged shape gives users one node per KB instead of two\nparallel ones that always shared the same ``knowledge_base`` picker,\nembedding-model metadata, and backend registry.\n\nBoth legacy classes (``KnowledgeIngestionComponent``,\n``KnowledgeBaseComponent``) remain importable as thin subclasses of\n``KnowledgeComponent`` — see the sibling ``ingestion.py`` / ``retrieval.py``\nmodules — so saved flows continue to load unchanged.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport asyncio\nimport hashlib\nimport json\nimport re\nimport uuid\nfrom dataclasses import asdict, dataclass, field\nfrom datetime import datetime, timezone\nfrom typing import TYPE_CHECKING, Any\n\nimport pandas as pd\nfrom cryptography.fernet import InvalidToken\nfrom langchain_chroma import Chroma\nfrom langflow.services.auth.utils import decrypt_api_key, encrypt_api_key\n\nfrom lfx.base.knowledge_bases.backends import BackendType, BaseVectorStoreBackend, create_backend\nfrom lfx.base.knowledge_bases.ingestion_sources.base import (\n IngestionItemResult,\n IngestionItemStatus,\n IngestionRunStatus,\n IngestionSummary,\n)\nfrom lfx.base.knowledge_bases.ingestion_sources.flow_component import FlowComponentSource\nfrom lfx.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases\nfrom lfx.base.models.unified_models import get_embedding_model_options, get_embeddings\nfrom lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs\nfrom lfx.components.files_and_knowledge._kb_paths import (\n KBKeyDecryptError,\n load_kb_metadata,\n)\nfrom lfx.components.files_and_knowledge._kb_paths import (\n get_knowledge_bases_root_path as _get_knowledge_bases_root_path,\n)\nfrom lfx.components.processing.converter import convert_to_dataframe\nfrom lfx.custom import Component\nfrom lfx.io import (\n BoolInput,\n DBProviderInput,\n DropdownInput,\n HandleInput,\n IntInput,\n MessageTextInput,\n ModelInput,\n Output,\n SecretStrInput,\n StrInput,\n TabInput,\n TableInput,\n)\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.schema.dotdict import dotdict\nfrom lfx.schema.table import EditMode\nfrom lfx.services.deps import (\n get_settings_service,\n session_scope,\n)\nfrom lfx.utils.component_utils import set_current_fields, set_field_display\nfrom lfx.utils.validate_cloud import raise_error_if_astra_cloud_disable_component\n\n\ndef _inputs_for_mode(default_mode: str) -> list:\n \"\"\"Return a fresh copy of the canonical inputs list with show flags set for the given mode.\n\n Used by the legacy subclasses so a saved flow keyed on\n ``KnowledgeIngestion`` / ``KnowledgeBase`` lands on a node template\n whose default visibility already matches the pinned mode — no\n \"wait for the user to click the mode tab\" UX gap on load.\n \"\"\"\n always_visible = set(KnowledgeComponent.default_keys) | set(KnowledgeComponent.mode_config[default_mode])\n inputs_copy = []\n for inp in KnowledgeComponent.inputs:\n clone = inp.model_copy(deep=True) if hasattr(inp, \"model_copy\") else inp\n if clone.name == \"mode\":\n clone.value = default_mode\n else:\n clone.show = clone.name in always_visible\n inputs_copy.append(clone)\n return inputs_copy\n\n\nif TYPE_CHECKING:\n from pathlib import Path\n\n# Mode constants. Plain-text labels (no emoji) for consistency with the rest of\n# Langflow's TabInput palette — see ``MemoryComponent`` for the same convention.\nMODE_INGEST = \"Ingest\"\nMODE_RETRIEVE = \"Retrieve\"\n\n\ndef _is_retrieve_mode(value: Any) -> bool:\n \"\"\"Lenient mode check: treats any label containing 'Retrieve' as retrieve mode.\n\n Older saved flows may carry the emoji-prefixed labels (\"📥 Ingest\" /\n \"🔍 Retrieve\") this component used to ship with; substring matching\n keeps those loading without forcing a flow rewrite.\n \"\"\"\n return isinstance(value, str) and \"Retrieve\" in value\n\n\n# Error message used by both the ingest and retrieve paths when the user is\n# running against an Astra cloud environment that disables these flows.\nastra_error_msg = \"Knowledge ingestion and retrieval are not supported in Astra cloud environment.\"\n\n_DEFAULT_OPENSEARCH_CONFIG = {\n \"url_variable\": \"OPENSEARCH_URL\",\n \"username_variable\": \"OPENSEARCH_USERNAME\",\n \"password_variable\": \"OPENSEARCH_PASSWORD\", # pragma: allowlist secret\n \"index_name\": \"\",\n \"vector_field\": \"vector_field\",\n \"text_field\": \"text\",\n}\n\n_DEFAULT_CHROMA_CLOUD_CONFIG = {\n \"mode\": \"cloud\",\n \"tenant_variable\": \"CHROMA_TENANT\",\n \"database_variable\": \"CHROMA_DATABASE\",\n \"api_key_variable\": \"CHROMA_API_KEY\", # pragma: allowlist secret\n}\n\n\nclass KnowledgeComponent(Component):\n \"\"\"One component for both writing into and reading from a Langflow knowledge base.\n\n A ``TabInput`` switches between ingestion and retrieval. The\n ``update_build_config`` / ``update_outputs`` hooks hide the inputs\n and the output that don't apply to the current mode so the canvas\n node stays focused.\n \"\"\"\n\n display_name = \"Knowledge\"\n description = \"Ingest into or retrieve from a Langflow knowledge base.\"\n icon = \"database\"\n name = \"Knowledge\"\n\n # ------ Mode → visible-fields wiring ---------------------------------\n # ``default_keys`` are inputs always visible regardless of mode.\n # ``mode_config`` lists the inputs unique to each mode; everything outside\n # this set is hidden when its mode is not selected.\n default_keys: list[str] = [\"mode\", \"knowledge_base\"]\n mode_config: dict[str, list[str]] = {\n MODE_INGEST: [\n \"input_df\",\n \"column_config\",\n \"chunk_size\",\n \"api_key\",\n \"allow_duplicates\",\n \"metadata_json\",\n ],\n MODE_RETRIEVE: [\n \"search_query\",\n \"top_k\",\n \"include_metadata\",\n \"include_embeddings\",\n \"metadata_filter\",\n ],\n }\n\n def __init__(self, *args, **kwargs) -> None:\n super().__init__(*args, **kwargs)\n self._cached_kb_path: Path | None = None\n\n @dataclass\n class NewKnowledgeBaseInput:\n functionality: str = \"create\"\n fields: dict[str, dict] = field(\n default_factory=lambda: {\n \"data\": {\n \"node\": {\n \"name\": \"create_knowledge_base\",\n \"description\": \"Create new knowledge in Langflow.\",\n \"display_name\": \"Create new Knowledge Base\",\n \"field_order\": [\n \"01_new_kb_name\",\n \"02_embedding_model\",\n \"03_knowledge_backend\",\n ],\n \"template\": {\n \"01_new_kb_name\": StrInput(\n name=\"new_kb_name\",\n display_name=\"Knowledge Name\",\n info=\"Name of the new knowledge to create.\",\n required=True,\n ),\n \"02_embedding_model\": ModelInput(\n name=\"embedding_model\",\n display_name=\"Choose Embedding Model\",\n info=(\n \"Select the embedding model to use for this knowledge base. \"\n \"Langflow uses the configured credentials for that model provider.\"\n ),\n required=True,\n model_type=\"embedding\",\n ),\n \"03_knowledge_backend\": DBProviderInput(\n name=\"knowledge_backend\",\n display_name=\"DB Provider\",\n info=(\n \"Select where this knowledge base stores vectors. \"\n \"OpenSearch uses the global DB Providers settings.\"\n ),\n required=True,\n ),\n },\n },\n }\n }\n )\n\n # ------ Inputs --------------------------------------------------------\n # ``knowledge_base`` is kept at position 0 so legacy tests that check\n # ``component.inputs[0].dialog_inputs`` continue to work.\n inputs = [\n DropdownInput(\n name=\"knowledge_base\",\n display_name=\"Knowledge\",\n info=\"Select the knowledge to load data from.\",\n required=True,\n options=[],\n refresh_button=True,\n real_time_refresh=True,\n dialog_inputs=asdict(NewKnowledgeBaseInput()),\n ),\n TabInput(\n name=\"mode\",\n display_name=\"Mode\",\n options=[MODE_INGEST, MODE_RETRIEVE],\n value=MODE_INGEST,\n info=\"Switch between writing new data into the knowledge base and querying it.\",\n real_time_refresh=True,\n tool_mode=True,\n ),\n # --- Ingest-only inputs (default-shown; hidden when mode == Retrieve) -\n HandleInput(\n name=\"input_df\",\n display_name=\"Input\",\n info=(\n \"Table with all original columns (already chunked / processed). \"\n \"Accepts Message, Data, or DataFrame. If Message or Data is provided, \"\n \"it is converted to a DataFrame automatically.\"\n ),\n input_types=[\"Message\", \"Data\", \"JSON\", \"DataFrame\", \"Table\"],\n required=True,\n dynamic=True,\n show=True,\n ),\n TableInput(\n name=\"column_config\",\n display_name=\"Column Configuration\",\n info=\"Configure column behavior for the knowledge base.\",\n required=True,\n table_schema=[\n {\n \"name\": \"column_name\",\n \"display_name\": \"Column Name\",\n \"type\": \"str\",\n \"description\": \"Name of the column in the source DataFrame\",\n \"edit_mode\": EditMode.INLINE,\n },\n {\n \"name\": \"vectorize\",\n \"display_name\": \"Vectorize\",\n \"type\": \"boolean\",\n \"description\": \"Create embeddings for this column\",\n \"default\": False,\n \"edit_mode\": EditMode.INLINE,\n },\n {\n \"name\": \"identifier\",\n \"display_name\": \"Identifier\",\n \"type\": \"boolean\",\n \"description\": \"Use this column as unique identifier\",\n \"default\": False,\n \"edit_mode\": EditMode.INLINE,\n },\n ],\n value=[\n {\n \"column_name\": \"text\",\n \"vectorize\": True,\n \"identifier\": True,\n },\n ],\n dynamic=True,\n show=True,\n ),\n IntInput(\n name=\"chunk_size\",\n display_name=\"Chunk Size\",\n info=\"Batch size for processing embeddings\",\n advanced=True,\n value=1000,\n dynamic=True,\n show=True,\n ),\n SecretStrInput(\n name=\"api_key\",\n display_name=\"Embedding Provider API Key\",\n info=\"Overrides global provider settings. Leave blank to use your pre-configured API Key.\",\n advanced=True,\n required=False,\n dynamic=True,\n show=True,\n ),\n BoolInput(\n name=\"allow_duplicates\",\n display_name=\"Allow Duplicates\",\n info=\"Allow duplicate rows in the knowledge base\",\n advanced=True,\n value=False,\n dynamic=True,\n show=True,\n ),\n StrInput(\n name=\"metadata_json\",\n display_name=\"Metadata\",\n info=(\n \"Optional JSON object of user metadata applied to every chunk produced by this \"\n 'run (e.g. {\"tag\": \"invoice\", \"year\": \"2026\"}). Same shape as the upload modal '\n \"Metadata section so chunks browser filters + Knowledge retrieval metadata_filter \"\n \"work uniformly across upload, folder, and flow-driven ingestion. Malformed JSON is \"\n \"ignored with a warning rather than failing the run.\"\n ),\n advanced=True,\n required=False,\n dynamic=True,\n show=True,\n ),\n # --- Retrieve-only inputs (default-hidden; shown when mode == Retrieve) -\n MessageTextInput(\n name=\"search_query\",\n display_name=\"Search Query\",\n info=\"Optional search query to filter knowledge base data.\",\n tool_mode=True,\n dynamic=True,\n show=False,\n ),\n IntInput(\n name=\"top_k\",\n display_name=\"Top K Results\",\n info=\"Number of top results to return from the knowledge base.\",\n value=5,\n advanced=True,\n required=False,\n dynamic=True,\n show=False,\n ),\n BoolInput(\n name=\"include_metadata\",\n display_name=\"Include Metadata\",\n info=\"Whether to include all metadata in the output. If false, only content is returned.\",\n value=True,\n advanced=False,\n dynamic=True,\n show=False,\n ),\n BoolInput(\n name=\"include_embeddings\",\n display_name=\"Include Embeddings\",\n info=\"Whether to include embeddings in the output. Only applicable if 'Include Metadata' is enabled.\",\n value=False,\n advanced=True,\n dynamic=True,\n show=False,\n ),\n MessageTextInput(\n name=\"metadata_filter\",\n display_name=\"Metadata Filter\",\n info=(\n \"Optional JSON object of user-metadata key/value pairs. Only chunks \"\n 'whose source_metadata matches every key are returned (e.g. {\"tag\": \"invoice\"} '\n 'or {\"tag\": [\"invoice\", \"audit\"]} for OR-of-values). Backends without '\n \"native filtering apply the match client-side after retrieval.\"\n ),\n advanced=True,\n dynamic=True,\n show=False,\n ),\n ]\n\n # ------ Outputs -------------------------------------------------------\n # Both outputs are declared at the class level so the runtime can\n # dispatch to either method depending on which one the saved flow has\n # wired up. ``update_outputs`` filters the canvas-visible output per\n # selected mode; see ``TypeConverterComponent`` and ``MemoryComponent``\n # for the same pattern.\n #\n # Output names match the legacy ``KnowledgeIngestionComponent``\n # (``dataframe_output``) and ``KnowledgeBaseComponent``\n # (``retrieve_data``) so saved flow edges keyed on those names resolve\n # cleanly against the merged component.\n outputs = [\n Output(\n display_name=\"Results\",\n name=\"dataframe_output\",\n method=\"build_kb_info\",\n types=[\"JSON\"],\n selected=\"JSON\",\n ),\n Output(\n display_name=\"Results\",\n name=\"retrieve_data\",\n method=\"retrieve_data\",\n info=\"Returns the data from the selected knowledge base.\",\n types=[\"Table\"],\n selected=\"Table\",\n ),\n ]\n\n # ------ Mode-driven UI updates ---------------------------------------\n async def update_frontend_node(self, new_frontend_node: dict, current_frontend_node: dict):\n \"\"\"Sync the visible output with the current ``mode`` value on canvas load.\n\n Saved flows hit this path: ``update_outputs`` is normally only triggered\n by ``real_time_refresh`` field edits, so without this re-sync the\n canvas could land with both outputs visible.\n \"\"\"\n await super().update_frontend_node(new_frontend_node, current_frontend_node)\n mode_value = new_frontend_node.get(\"template\", {}).get(\"mode\", {}).get(\"value\", MODE_INGEST)\n self.update_outputs(new_frontend_node, \"mode\", mode_value)\n return new_frontend_node\n\n def update_outputs(self, frontend_node: dict, field_name: str, field_value: Any) -> dict:\n \"\"\"Filter visible outputs to match the selected mode.\n\n Triggered by the ``mode`` ``TabInput`` ``real_time_refresh`` flag.\n The class-level ``outputs`` list always carries both entries so the\n runtime can resolve either ``build_kb_info`` or ``retrieve_data``\n regardless of canvas visibility — we just hide the unused one here.\n \"\"\"\n if field_name != \"mode\":\n return frontend_node\n if _is_retrieve_mode(field_value):\n frontend_node[\"outputs\"] = [\n Output(\n display_name=\"Results\",\n name=\"retrieve_data\",\n method=\"retrieve_data\",\n info=\"Returns the data from the selected knowledge base.\",\n types=[\"Table\"],\n selected=\"Table\",\n )\n ]\n else:\n frontend_node[\"outputs\"] = [\n Output(\n display_name=\"Results\",\n name=\"dataframe_output\",\n method=\"build_kb_info\",\n types=[\"JSON\"],\n selected=\"JSON\",\n )\n ]\n return frontend_node\n\n async def update_build_config(\n self,\n build_config,\n field_value: Any,\n field_name: str | None = None,\n ):\n \"\"\"Refresh KB options, drive the create-KB dialog, and hide off-mode fields.\"\"\"\n # Astra-cloud gate covers both ingest and retrieve paths.\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n # Always populate the create-KB dialog's embedding-model options so the\n # ModelInput renders correctly regardless of which input triggered the\n # refresh.\n try:\n dialog_template = (\n build_config[\"knowledge_base\"]\n .get(\"dialog_inputs\", {})\n .get(\"fields\", {})\n .get(\"data\", {})\n .get(\"node\", {})\n .get(\"template\", {})\n )\n if \"02_embedding_model\" in dialog_template:\n embedding_options = get_embedding_model_options(user_id=self.user_id)\n dialog_template[\"02_embedding_model\"][\"options\"] = embedding_options\n except Exception: # noqa: BLE001\n self.log(\"Failed to populate embedding model options in dialog\")\n\n # KB-picker refresh + create-new-KB flow (lifted from the legacy ingestion\n # component verbatim; relied on by both the canvas refresh button and the\n # dialog-submit path).\n if field_name == \"knowledge_base\":\n # Lazy import keeps lfx importable without langflow installed.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching knowledge base list.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n if isinstance(field_value, dict) and \"01_new_kb_name\" in field_value:\n if not self.is_valid_collection_name(field_value[\"01_new_kb_name\"]):\n msg = f\"Invalid knowledge base name: {field_value['01_new_kb_name']}\"\n raise ValueError(msg)\n\n model_selection = field_value[\"02_embedding_model\"]\n if isinstance(model_selection, dict):\n model_selection = [model_selection]\n\n backend_type, backend_config = self._normalize_backend_selection(\n field_value.get(\"03_knowledge_backend\")\n )\n\n embed_model = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n )\n\n try:\n await asyncio.wait_for(\n asyncio.to_thread(embed_model.embed_query, \"test\"),\n timeout=10,\n )\n except TimeoutError as e:\n msg = \"Embedding validation timed out. Please verify network connectivity and key.\"\n raise ValueError(msg) from e\n except Exception as e:\n msg = f\"Embedding validation failed: {e!s}\"\n raise ValueError(msg) from e\n\n kb_path = self._resolve_kb_path(\n _get_knowledge_bases_root_path(), kb_user, field_value[\"01_new_kb_name\"]\n )\n kb_path.mkdir(parents=True, exist_ok=True)\n\n build_config[\"knowledge_base\"][\"value\"] = field_value[\"01_new_kb_name\"]\n self._save_embedding_metadata(\n kb_path=kb_path,\n model_selection=model_selection,\n backend_type=backend_type,\n backend_config=backend_config,\n )\n await self._create_knowledge_base_record(\n user_id=self.user_id,\n name=field_value[\"01_new_kb_name\"],\n model_selection=model_selection,\n backend_type=backend_type,\n backend_config=backend_config,\n )\n\n build_config[\"knowledge_base\"][\"options\"] = await get_knowledge_bases(\n _get_knowledge_bases_root_path(),\n user_id=self.user_id,\n )\n if build_config[\"knowledge_base\"][\"value\"] not in build_config[\"knowledge_base\"][\"options\"]:\n build_config[\"knowledge_base\"][\"value\"] = None\n\n # Honor the current mode regardless of which field triggered the refresh.\n # Falls back to MODE_INGEST when ``mode`` is missing (legacy nodes).\n current_mode = build_config.get(\"mode\", {}).get(\"value\") if isinstance(build_config, dict) else None\n if field_name == \"mode\":\n current_mode = field_value\n # Map legacy/emoji-prefixed labels onto the current canonical values so\n # flows saved before the label change still toggle visibility correctly.\n if _is_retrieve_mode(current_mode):\n current_mode = MODE_RETRIEVE\n elif current_mode not in self.mode_config:\n current_mode = MODE_INGEST\n return set_current_fields(\n build_config=build_config if isinstance(build_config, dotdict) else dotdict(build_config),\n action_fields=self.mode_config,\n selected_action=current_mode,\n default_fields=self.default_keys,\n func=set_field_display,\n )\n\n # =====================================================================\n # INGESTION CODE PATH\n # =====================================================================\n def _get_kb_root(self) -> Path:\n \"\"\"Return the root directory for knowledge bases.\"\"\"\n return _get_knowledge_bases_root_path()\n\n @staticmethod\n def _resolve_kb_path(kb_root: Path, kb_user: str, kb_name: str) -> Path:\n \"\"\"Resolve the selected KB inside the authenticated user's directory.\"\"\"\n # Lazy import keeps lfx importable without langflow's DB services.\n from langflow.services.memory_base.kb_path_helpers import validate_kb_path\n\n user_root = kb_root / kb_user\n validate_kb_path(kb_root, user_root)\n\n kb_path = user_root / kb_name\n validate_kb_path(user_root, kb_path)\n return kb_path\n\n @staticmethod\n def _scalar_notna(value) -> bool:\n \"\"\"Check if a value is not NA, safely handling arrays and sequences.\n\n ``pd.notna`` returns an array when given an array-like input, which\n cannot be used directly in a boolean context. This helper collapses\n the result to a single scalar ``bool``.\n \"\"\"\n result = pd.notna(value)\n if hasattr(result, \"__iter__\") and not isinstance(result, str):\n import numpy as np\n\n arr = np.asarray(result)\n return arr.size > 0 and arr.all()\n return bool(result)\n\n def _validate_column_config(self, df_source: pd.DataFrame) -> list[dict[str, Any]]:\n \"\"\"Validate column configuration using Structured Output patterns.\"\"\"\n if not self.column_config:\n msg = \"Column configuration cannot be empty\"\n raise ValueError(msg)\n\n config_list = self.column_config if isinstance(self.column_config, list) else []\n\n df_columns = set(df_source.columns)\n for config in config_list:\n col_name = config.get(\"column_name\")\n if col_name not in df_columns:\n msg = f\"Column '{col_name}' not found in DataFrame. Available columns: {sorted(df_columns)}\"\n raise ValueError(msg)\n\n return config_list\n\n def _build_embedding_metadata(\n self,\n model_selection: list[dict[str, Any]],\n api_key: str | None = None,\n backend_type: str = BackendType.CHROMA.value,\n backend_config: dict[str, Any] | None = None,\n ) -> dict[str, Any]:\n \"\"\"Build embedding model metadata from a model selection dict.\"\"\"\n model_dict = model_selection[0] if isinstance(model_selection, list) else model_selection\n embedding_model = model_dict.get(\"name\", \"\")\n embedding_provider = model_dict.get(\"provider\", \"Unknown\")\n\n api_key_to_save = None\n if api_key and hasattr(api_key, \"get_secret_value\"):\n api_key_to_save = api_key.get_secret_value()\n elif isinstance(api_key, str):\n api_key_to_save = api_key\n\n encrypted_api_key = None\n if api_key_to_save:\n settings_service = get_settings_service()\n try:\n encrypted_api_key = encrypt_api_key(api_key_to_save, settings_service=settings_service)\n except (TypeError, ValueError) as e:\n self.log(f\"Could not encrypt API key: {e}\")\n\n return {\n \"embedding_provider\": embedding_provider,\n \"embedding_model\": embedding_model,\n \"model_selection\": model_dict,\n \"api_key\": encrypted_api_key,\n \"api_key_used\": bool(api_key),\n \"chunk_size\": self.chunk_size,\n \"backend_type\": backend_type,\n \"backend_config\": backend_config or {},\n \"created_at\": datetime.now(timezone.utc).isoformat(),\n }\n\n def _save_embedding_metadata(\n self,\n kb_path: Path,\n model_selection: list[dict[str, Any]],\n api_key: str | None = None,\n backend_type: str | None = None,\n backend_config: dict[str, Any] | None = None,\n ) -> None:\n \"\"\"Save embedding model metadata.\"\"\"\n metadata_path = kb_path / \"embedding_metadata.json\"\n existing_metadata: dict[str, Any] = {}\n if metadata_path.exists():\n try:\n existing_metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError):\n existing_metadata = {}\n\n embedding_metadata = self._build_embedding_metadata(\n model_selection,\n api_key,\n backend_type=backend_type or existing_metadata.get(\"backend_type\") or BackendType.CHROMA.value,\n backend_config=backend_config\n if backend_config is not None\n else existing_metadata.get(\"backend_config\") or {},\n )\n metadata_path.write_text(json.dumps(embedding_metadata, indent=2))\n\n def _update_metadata_metrics(self, kb_path: Path, chroma: Chroma) -> None:\n \"\"\"Update embedding_metadata.json with accurate chunk/word/character counts.\"\"\"\n import chromadb.errors\n from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper\n\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n\n try:\n metadata = json.loads(metadata_path.read_text())\n KBAnalysisHelper.update_text_metrics(kb_path, metadata, chroma)\n metadata[\"size\"] = KBStorageHelper.get_directory_size(kb_path)\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError, chromadb.errors.ChromaError) as e:\n self.log(f\"Warning: Could not update metadata metrics: {e}\")\n\n @staticmethod\n def _extract_source_types_from_df(df_source: pd.DataFrame) -> set[str]:\n \"\"\"Pull file extensions out of common path/name columns on the source DataFrame.\n\n The direct-upload ingestion path stores extensions in\n ``embedding_metadata.json[source_types]`` so the KB list can render the\n correct file-type icon. When ingestion happens via a connected\n ``input_df`` (e.g. File → Knowledge) the same field stayed empty and\n the icon defaulted to a blank tile. We look at the well-known columns\n the File / S3 / cloud-storage components produce and collect any\n plausible extension so the icon is consistent across both flows.\n \"\"\"\n candidate_columns = (\"file_path\", \"file_name\", \"filename\", \"source\", \"path\", \"mimetype\")\n extensions: set[str] = set()\n for col in candidate_columns:\n if col not in df_source.columns:\n continue\n for value in df_source[col].dropna():\n ext = KnowledgeComponent._extension_from_value(value)\n # Drop anything that doesn't look like an extension (URL\n # query strings, version segments, etc.) — the icon palette\n # keys off short alphanumeric tokens like \"pdf\"/\"docx\".\n if ext:\n extensions.add(ext)\n return extensions\n\n @staticmethod\n def _extension_from_value(value: Any) -> str | None:\n \"\"\"Return a normalized file-extension token from a path / filename / MIME string.\n\n Accepts values like ``\"report.PDF\"``, ``\"/docs/notes.txt\"``, or\n ``\"application/pdf\"`` and returns ``\"pdf\"`` / ``\"txt\"``. Returns\n ``None`` if no plausible extension can be derived.\n \"\"\"\n extension_length_limit = 10\n if value is None:\n return None\n text = str(value).strip()\n if not text:\n return None\n # MIME types like ``application/pdf`` carry the canonical extension\n # in the subtype slot — preserve them so File / S3 messages keyed\n # only on ``mimetype`` still resolve to an icon.\n if \"/\" in text and \".\" not in text.rsplit(\"/\", 1)[-1]:\n subtype = text.rsplit(\"/\", 1)[-1].strip().lower()\n return subtype if subtype and len(subtype) <= extension_length_limit and subtype.isalnum() else None\n if \".\" not in text:\n return None\n ext = text.rsplit(\".\", 1)[-1].strip().lower()\n if ext and len(ext) <= extension_length_limit and ext.isalnum():\n return ext\n return None\n\n @classmethod\n def _extract_source_types_from_mapping(cls, mapping: Any) -> set[str]:\n \"\"\"Pull extensions from a Message/Data-style ``data`` dict or a plain dict.\"\"\"\n if not isinstance(mapping, dict):\n return set()\n candidate_keys = (\"file_path\", \"file_name\", \"filename\", \"source\", \"path\", \"mimetype\")\n extensions: set[str] = set()\n for key in candidate_keys:\n ext = cls._extension_from_value(mapping.get(key))\n if ext:\n extensions.add(ext)\n return extensions\n\n @classmethod\n def _extract_source_types_from_input(cls, input_value: Any) -> set[str]:\n \"\"\"Pull file extensions out of the raw component input.\n\n ``convert_to_dataframe`` strips Message/Data fields down to ``text``\n when projecting onto a DataFrame, so file metadata attached to a\n File-component \"Raw Content\" output never reaches\n ``_extract_source_types_from_df``. Looking at the raw input first\n keeps the KB icon consistent with direct upload.\n \"\"\"\n if input_value is None:\n return set()\n if isinstance(input_value, list):\n extensions: set[str] = set()\n for item in input_value:\n extensions |= cls._extract_source_types_from_input(item)\n return extensions\n if isinstance(input_value, pd.DataFrame):\n return cls._extract_source_types_from_df(input_value)\n # Message / Data / JSON all expose a ``data`` dict via the lfx schema.\n mapping = getattr(input_value, \"data\", None)\n if mapping is None and isinstance(input_value, dict):\n mapping = input_value\n return cls._extract_source_types_from_mapping(mapping)\n\n def _merge_source_types(self, kb_path: Path, extensions: set[str]) -> None:\n \"\"\"Merge newly observed extensions into the KB's ``source_types`` metadata.\n\n Mirrors the direct-upload path in ``KBIngestionHelper`` so the icon\n rendering on the Knowledge Bases list works regardless of which\n ingestion route was used.\n \"\"\"\n if not extensions:\n return\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n try:\n metadata = json.loads(metadata_path.read_text())\n existing = set(metadata.get(\"source_types\") or [])\n metadata[\"source_types\"] = sorted(existing | extensions)\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError) as e:\n self.log(f\"Warning: Could not update source_types metadata: {e}\")\n\n async def _update_backend_metadata_metrics(self, kb_path: Path, backend: BaseVectorStoreBackend) -> None:\n \"\"\"Update metadata metrics for non-Chroma backends.\"\"\"\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return\n\n try:\n metadata = json.loads(metadata_path.read_text())\n chunks = await backend.count()\n characters = 0\n words = 0\n async for batch in backend.iter_documents():\n for document in batch:\n characters += len(document.content)\n words += len(document.content.split())\n\n metadata[\"chunks\"] = chunks\n metadata[\"characters\"] = characters\n metadata[\"words\"] = words\n metadata[\"avg_chunk_size\"] = characters / chunks if chunks else 0.0\n metadata[\"size\"] = await backend.storage_size_bytes()\n metadata_path.write_text(json.dumps(metadata, indent=2))\n except (OSError, ValueError, TypeError, json.JSONDecodeError) as e:\n self.log(f\"Warning: Could not update backend metadata metrics: {e}\")\n\n @staticmethod\n def _normalize_backend_selection(value: Any) -> tuple[str, dict[str, Any]]:\n \"\"\"Normalize a DBProviderInput value into backend type/config.\"\"\"\n if not value:\n return BackendType.CHROMA.value, {}\n\n if isinstance(value, str):\n backend_type = value if value == BackendType.OPENSEARCH.value else BackendType.CHROMA.value\n return (\n backend_type,\n _DEFAULT_OPENSEARCH_CONFIG.copy() if backend_type == BackendType.OPENSEARCH.value else {},\n )\n\n if not isinstance(value, dict):\n return BackendType.CHROMA.value, {}\n\n backend_type = str(value.get(\"backend_type\") or value.get(\"id\") or BackendType.CHROMA.value)\n\n if backend_type == BackendType.OPENSEARCH.value:\n backend_config = value.get(\"backend_config\") or value.get(\"config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return BackendType.OPENSEARCH.value, {**_DEFAULT_OPENSEARCH_CONFIG, **backend_config}\n\n if backend_type == \"chroma_cloud\":\n backend_config = value.get(\"backend_config\") or value.get(\"config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return BackendType.CHROMA.value, {**_DEFAULT_CHROMA_CLOUD_CONFIG, **backend_config}\n\n return BackendType.CHROMA.value, {}\n\n @staticmethod\n def _get_backend_from_metadata(kb_path: Path) -> tuple[str, dict[str, Any]] | None:\n \"\"\"Read ``(backend_type, backend_config)`` off the on-disk sidecar.\n\n Returns ``None`` — never a Chroma default — when the sidecar is absent or\n unreadable. \"Could not resolve\" has to stay distinguishable from\n \"explicitly local\", or a replica that merely lacks the file writes a\n remote-backed KB into a local store. See ``_resolve_backend_config``.\n \"\"\"\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n return None\n try:\n metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError):\n return None\n\n backend_type = str(metadata.get(\"backend_type\") or BackendType.CHROMA.value)\n backend_config = metadata.get(\"backend_config\") or {}\n if not isinstance(backend_config, dict):\n backend_config = {}\n return backend_type, backend_config\n\n async def _create_knowledge_base_record(\n self,\n *,\n user_id: Any,\n name: str,\n model_selection: list[dict[str, Any]],\n backend_type: str,\n backend_config: dict[str, Any],\n ) -> None:\n \"\"\"Persist the component-created KB in the DB when Langflow is available.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n except ImportError:\n return\n\n try:\n await knowledge_base_service.create_record(\n user_id=user_id,\n name=name,\n model_selection=model_selection,\n column_config=self.column_config if isinstance(self.column_config, list) else [],\n backend_type=backend_type,\n backend_config=backend_config,\n )\n except Exception as exc: # noqa: BLE001\n self.log(f\"Warning: could not persist knowledge base record: {exc}\")\n\n def _save_kb_files(\n self,\n kb_path: Path,\n config_list: list[dict[str, Any]],\n ) -> None:\n \"\"\"Save KB files using File Component storage patterns.\"\"\"\n try:\n kb_path.mkdir(parents=True, exist_ok=True)\n\n cfg_path = kb_path / \"schema.json\"\n if not cfg_path.exists():\n cfg_path.write_text(json.dumps(config_list, indent=2))\n\n except (OSError, TypeError, ValueError) as e:\n self.log(f\"Error saving KB files: {e}\")\n\n def _build_column_metadata(self, config_list: list[dict[str, Any]], df_source: pd.DataFrame) -> dict[str, Any]:\n \"\"\"Build detailed column metadata.\"\"\"\n metadata: dict[str, Any] = {\n \"total_columns\": len(df_source.columns),\n \"mapped_columns\": len(config_list),\n \"unmapped_columns\": len(df_source.columns) - len(config_list),\n \"columns\": [],\n \"summary\": {\"vectorized_columns\": [], \"identifier_columns\": []},\n }\n\n for config in config_list:\n col_name = config.get(\"column_name\")\n vectorize = config.get(\"vectorize\") == \"True\" or config.get(\"vectorize\") is True\n identifier = config.get(\"identifier\") == \"True\" or config.get(\"identifier\") is True\n\n metadata[\"columns\"].append(\n {\n \"name\": col_name,\n \"vectorize\": vectorize,\n \"identifier\": identifier,\n }\n )\n\n if vectorize:\n metadata[\"summary\"][\"vectorized_columns\"].append(col_name)\n if identifier:\n metadata[\"summary\"][\"identifier_columns\"].append(col_name)\n\n return metadata\n\n async def _create_vector_store(\n self,\n df_source: pd.DataFrame,\n config_list: list[dict[str, Any]],\n embedding_function,\n ) -> BaseVectorStoreBackend:\n \"\"\"Create vector store using the configured DB provider.\"\"\"\n vector_store_dir = await self._kb_path()\n if not vector_store_dir:\n msg = \"Knowledge base path is not set. Please create a new knowledge base first.\"\n raise ValueError(msg)\n vector_store_dir.mkdir(parents=True, exist_ok=True)\n\n backend_type, backend_config = await self._resolve_backend_config(vector_store_dir)\n backend = create_backend(\n backend_type,\n kb_name=self.knowledge_base,\n kb_path=vector_store_dir,\n backend_config=backend_config,\n embedding_function=embedding_function,\n user_id=self.user_id,\n )\n await backend.ensure_ready()\n\n existing_ids = None\n if backend_type != BackendType.CHROMA.value and not self.allow_duplicates:\n existing_ids = set()\n async for batch in backend.iter_documents():\n for document in batch:\n doc_id = document.metadata.get(\"_id\")\n if doc_id:\n existing_ids.add(doc_id)\n\n data_objects = await self._convert_df_to_data_objects(df_source, config_list, existing_ids=existing_ids)\n\n user_metadata_tag = self._resolve_user_metadata_tag()\n\n documents = []\n for data_obj in data_objects:\n doc = data_obj.to_lc_document()\n if user_metadata_tag:\n doc.metadata[\"source_metadata\"] = user_metadata_tag\n documents.append(doc)\n\n if documents:\n await backend.add_documents(documents)\n self.log(f\"Added {len(documents)} documents to vector store '{self.knowledge_base}'\")\n\n return backend\n\n async def _convert_df_to_data_objects(\n self,\n df_source: pd.DataFrame,\n config_list: list[dict[str, Any]],\n existing_ids: set[str] | None = None,\n ) -> list[Data]:\n \"\"\"Convert DataFrame to Data objects for vector store.\"\"\"\n data_objects: list[Data] = []\n\n if existing_ids is None:\n kb_path = await self._kb_path()\n\n chroma = Chroma(\n persist_directory=str(kb_path),\n collection_name=self.knowledge_base,\n **chroma_langchain_collection_kwargs(),\n )\n\n all_docs = chroma.get()\n\n existing_ids = {metadata.get(\"_id\") for metadata in all_docs[\"metadatas\"] if metadata.get(\"_id\")}\n\n content_cols = []\n identifier_cols = []\n\n for config in config_list:\n col_name = config.get(\"column_name\")\n vectorize = config.get(\"vectorize\") == \"True\" or config.get(\"vectorize\") is True\n identifier = config.get(\"identifier\") == \"True\" or config.get(\"identifier\") is True\n\n if vectorize:\n content_cols.append(col_name)\n if identifier:\n identifier_cols.append(col_name)\n\n for _, row in df_source.iterrows():\n identifier_parts = [str(row[col]) for col in content_cols if col in row and self._scalar_notna(row[col])]\n\n page_content = \" \".join(identifier_parts)\n\n data_dict = {\n \"text\": page_content,\n }\n\n if identifier_cols:\n identifier_parts = [\n str(row[col]) for col in identifier_cols if col in row and self._scalar_notna(row[col])\n ]\n page_content = \" \".join(identifier_parts)\n\n for col in df_source.columns:\n if col not in content_cols and col in row and self._scalar_notna(row[col]):\n value = row[col]\n data_dict[col] = str(value)\n\n page_content_hash = hashlib.sha256(page_content.encode()).hexdigest()\n data_dict[\"_id\"] = page_content_hash\n\n if not self.allow_duplicates and page_content_hash in existing_ids:\n self.log(f\"Skipping duplicate row with hash {page_content_hash}\")\n continue\n\n data_obj = Data(data=data_dict)\n data_objects.append(data_obj)\n\n return data_objects\n\n def is_valid_collection_name(self, name, min_length: int = 3, max_length: int = 63) -> bool:\n \"\"\"Validate collection name.\n\n 1. Contains 3-63 characters\n 2. Starts and ends with alphanumeric character\n 3. Contains only alphanumeric characters, underscores, or hyphens.\n \"\"\"\n if not (min_length <= len(name) <= max_length):\n return False\n\n if not (name[0].isalnum() and name[-1].isalnum()):\n return False\n\n return re.match(r\"^[a-zA-Z0-9_-]+$\", name) is not None\n\n async def _kb_path(self) -> Path | None:\n cached_path = getattr(self, \"_cached_kb_path\", None)\n if cached_path is not None:\n return cached_path\n\n # Lazy import to keep ``lfx`` importable standalone — langflow's\n # user/DB models are not always available at module load time.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching knowledge base path.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n\n kb_root = self._get_kb_root()\n\n self._cached_kb_path = self._resolve_kb_path(kb_root, kb_user, self.knowledge_base)\n\n return self._cached_kb_path\n\n def _resolve_user_metadata_tag(self) -> str:\n \"\"\"Return the JSON-encoded user metadata tag for chunk writes.\"\"\"\n raw = getattr(self, \"metadata_json\", None)\n if not raw:\n return \"\"\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return \"\"\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError) as exc:\n self.log(f\"KnowledgeComponent: metadata_json is not valid JSON ({exc}); skipping metadata stamp.\")\n return \"\"\n if not isinstance(decoded, dict):\n self.log(\"KnowledgeComponent: metadata_json must decode to a JSON object; skipping metadata stamp.\")\n return \"\"\n return json.dumps(decoded, sort_keys=True)\n\n async def build_kb_info(self) -> Data:\n \"\"\"Main ingestion routine → returns a dict with KB metadata.\n\n The annotation is intentionally narrowed to ``Data`` even though the\n cross-mode fallback below may return a ``DataFrame`` from\n ``retrieve_data``. The frontend builds React-Flow handle IDs from\n this output's type list; widening it to ``Data | DataFrame`` makes\n the API advertise ``[\"JSON\", \"Table\"]`` for the ingest output and\n breaks every saved-edge sourceHandle that was generated against\n a single-type handle (BUG-02). Python doesn't enforce return\n annotations at runtime, so the rare fallback path keeps working.\n \"\"\"\n if _is_retrieve_mode(getattr(self, \"mode\", MODE_INGEST)):\n return await self.retrieve_data()\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n run_id: uuid.UUID | None = None\n run_job_id: uuid.UUID | None = None\n run_summary: IngestionSummary | None = None\n run_status: IngestionRunStatus = IngestionRunStatus.SUCCEEDED\n run_error: str | None = None\n kb_record_id: uuid.UUID | None = None\n try:\n input_value = self.input_df[0] if isinstance(self.input_df, list) else self.input_df\n df_source: DataFrame = convert_to_dataframe(input_value, auto_parse=False)\n\n config_list = self._validate_column_config(df_source)\n column_metadata = self._build_column_metadata(config_list, df_source)\n\n kb_path = await self._kb_path()\n if not kb_path:\n msg = \"Knowledge base path is not set. Please create a new knowledge base first.\"\n raise ValueError(msg)\n metadata_path = kb_path / \"embedding_metadata.json\"\n api_key = None\n model_selection = None\n\n if metadata_path.exists():\n settings_service = get_settings_service()\n stored_metadata = json.loads(metadata_path.read_text())\n\n model_selection = stored_metadata.get(\"model_selection\")\n if model_selection:\n model_selection = [model_selection] if isinstance(model_selection, dict) else model_selection\n else:\n embedding_model_name = stored_metadata.get(\"embedding_model\")\n embedding_provider = stored_metadata.get(\"embedding_provider\", \"Unknown\")\n if embedding_model_name:\n try:\n all_options = get_embedding_model_options(user_id=self.user_id)\n match = next(\n (o for o in all_options if o.get(\"name\") == embedding_model_name),\n None,\n )\n if match:\n model_selection = [match]\n else:\n self.log(\n f\"Embedding model '{embedding_model_name}' (provider: {embedding_provider}) \"\n \"from stored metadata is no longer available in the model registry. \"\n \"Please re-create this knowledge base with a supported embedding model.\"\n )\n msg = (\n f\"Embedding model '{embedding_model_name}' is no longer recognized. \"\n \"The knowledge base was created with an older format and the model \"\n \"is not available in the current registry. \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg)\n except ValueError:\n raise\n except Exception: # noqa: BLE001\n self.log(\n f\"Failed to look up embedding model '{embedding_model_name}' in registry. \"\n \"Please re-create this knowledge base with a supported embedding model.\"\n )\n msg = (\n f\"Could not look up embedding model '{embedding_model_name}' \"\n f\"(provider: {embedding_provider}). \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg) # noqa: B904\n\n encrypted_key = stored_metadata.get(\"api_key\")\n if encrypted_key:\n try:\n api_key = decrypt_api_key(encrypted_key, settings_service)\n except (InvalidToken, TypeError, ValueError) as e:\n if not self.api_key:\n log_label = f\"knowledge base '{self.knowledge_base}'\"\n msg = (\n f\"Cannot decrypt the stored embedding API key for {log_label}. \"\n \"This usually means the server's SECRET_KEY changed after the \"\n \"key was saved. To recover, supply the embedding provider API \"\n \"key on the component's 'Embedding Provider API Key' input and \"\n \"re-run ingestion — the key will be re-encrypted with the \"\n \"current SECRET_KEY.\"\n )\n raise KBKeyDecryptError(msg) from e\n logger.warning(\"Stored API key undecryptable; using component-supplied key. Error: %s\", e)\n\n if self.api_key:\n api_key = self.api_key\n if model_selection:\n self._save_embedding_metadata(\n kb_path=kb_path,\n model_selection=model_selection,\n api_key=api_key,\n )\n\n if not model_selection:\n msg = \"No embedding model configuration found. Please create the knowledge base first.\"\n raise ValueError(msg)\n\n embedding_function = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n api_key=api_key,\n chunk_size=self.chunk_size,\n )\n\n run_id, run_job_id, run_summary, kb_record_id = await self._begin_ingestion_run(kb_path)\n if kb_record_id is not None:\n await self._record_kb_status(kb_record_id, \"ingesting\")\n\n backend = await self._create_vector_store(df_source, config_list, embedding_function=embedding_function)\n\n self._save_kb_files(kb_path, config_list)\n\n try:\n if not isinstance(backend, BaseVectorStoreBackend):\n pass\n elif backend.backend_type == BackendType.CHROMA and hasattr(backend, \"raw_langchain_store\"):\n self._update_metadata_metrics(kb_path, backend.raw_langchain_store())\n else:\n await self._update_backend_metadata_metrics(kb_path, backend)\n # Stamp the KB with the file extensions we just ingested so\n # the Knowledge Bases list renders the correct icon for\n # flow-driven ingestion (input_df), matching direct upload.\n # We look at the raw input first because ``convert_to_dataframe``\n # drops Message/Data metadata fields (file_path, mimetype, …)\n # when projecting onto the DataFrame.\n source_types = self._extract_source_types_from_input(input_value)\n source_types |= self._extract_source_types_from_df(df_source)\n self._merge_source_types(kb_path, source_types)\n finally:\n if isinstance(backend, BaseVectorStoreBackend):\n await backend.teardown()\n\n meta: dict[str, Any] = {\n \"kb_id\": str(uuid.uuid4()),\n \"kb_name\": self.knowledge_base,\n \"rows\": len(df_source),\n \"column_metadata\": column_metadata,\n \"path\": str(kb_path),\n \"config_columns\": len(config_list),\n \"timestamp\": datetime.now(tz=timezone.utc).isoformat(),\n }\n\n if run_summary is not None:\n run_summary.record_item(\n IngestionItemResult(\n item_id=self.knowledge_base,\n display_name=f\"{self.knowledge_base} ({len(df_source)} rows)\",\n status=IngestionItemStatus.SUCCEEDED,\n chunks_created=len(df_source),\n ),\n )\n\n if kb_record_id is not None:\n await self._record_kb_stats(kb_record_id, kb_path)\n await self._record_kb_status(kb_record_id, \"ready\")\n\n self.status = f\"✅ KB **{self.knowledge_base}** saved · {len(df_source)} chunks.\"\n\n return Data(data=meta)\n\n except (OSError, ValueError, RuntimeError, KeyError) as e:\n run_status = IngestionRunStatus.FAILED\n run_error = str(e) or e.__class__.__name__\n msg = f\"Error during KB ingestion: {e}\"\n raise RuntimeError(msg) from e\n except Exception as e:\n run_status = IngestionRunStatus.FAILED\n run_error = str(e) or e.__class__.__name__\n raise\n finally:\n if run_id is not None and run_summary is not None:\n await self._finalize_ingestion_run(\n run_id=run_id,\n job_id=run_job_id,\n summary=run_summary,\n status=run_status,\n error_message=run_error,\n )\n if kb_record_id is not None and run_status is IngestionRunStatus.FAILED:\n await self._record_kb_status(kb_record_id, \"failed\", failure_reason=run_error)\n\n async def _begin_ingestion_run(\n self,\n kb_path: Path,\n ) -> tuple[uuid.UUID | None, uuid.UUID | None, IngestionSummary | None, uuid.UUID | None]:\n \"\"\"Create a parent ``Job`` and seed an ingestion-run row.\"\"\"\n if not self.user_id:\n self.log(\"No user_id on component; skipping ingestion-run tracking.\")\n return None, None, None, None\n\n try:\n user_uuid = uuid.UUID(str(self.user_id))\n except (ValueError, TypeError) as exc:\n self.log(f\"Could not coerce user_id={self.user_id!r} to UUID; skipping run tracking ({exc}).\")\n return None, None, None, None\n\n try:\n from langflow.api.utils import ingestion_run_service, knowledge_base_service\n from langflow.services.database.models.jobs.model import JobStatus, JobType\n from langflow.services.deps import get_job_service\n except ImportError as exc:\n self.log(f\"Run-history wiring unavailable; ingestion will not be recorded ({exc}).\")\n return None, None, None, None\n\n try:\n kb_record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base)\n kb_record_id = kb_record.id if kb_record is not None else None\n\n user_metadata = self._parse_user_metadata_dict()\n\n job_id = uuid.uuid4()\n job_service = get_job_service()\n raw_flow_id = getattr(self, \"flow_id\", None)\n flow_id_uuid = uuid.UUID(str(raw_flow_id)) if raw_flow_id else job_id\n await job_service.create_job(\n job_id=job_id,\n flow_id=flow_id_uuid,\n job_type=JobType.INGESTION,\n asset_id=kb_record_id,\n asset_type=\"knowledge_base\",\n user_id=user_uuid,\n )\n await job_service.update_job_status(job_id, JobStatus.IN_PROGRESS)\n\n source = FlowComponentSource(\n user_id=user_uuid,\n source_config={\n \"knowledge_base\": self.knowledge_base,\n \"kb_path\": str(kb_path),\n \"flow_id\": str(flow_id_uuid),\n },\n )\n\n run_id = await ingestion_run_service.create_run(\n kb_name=self.knowledge_base,\n source=source,\n job_id=job_id,\n user_id=user_uuid,\n kb_id=kb_record_id,\n user_metadata=user_metadata,\n )\n if run_id is None:\n return None, job_id, None, kb_record_id\n await ingestion_run_service.mark_running(run_id)\n\n summary = IngestionSummary(\n kb_name=self.knowledge_base,\n source_type=source.source_type.value,\n user_id=user_uuid,\n job_id=job_id,\n source_config=source.describe().get(\"config\") or {},\n user_metadata=user_metadata,\n )\n self.log(f\"Started ingestion run job_id={job_id} kb_name={self.knowledge_base} kb_id={kb_record_id}\")\n except Exception as exc: # noqa: BLE001 — telemetry must never abort ingestion\n self.log(f\"Could not begin ingestion-run tracking: {exc}\")\n return None, None, None, None\n else:\n return run_id, job_id, summary, kb_record_id\n\n async def _finalize_ingestion_run(\n self,\n *,\n run_id: uuid.UUID,\n job_id: uuid.UUID | None,\n summary: IngestionSummary,\n status: IngestionRunStatus,\n error_message: str | None,\n ) -> None:\n \"\"\"Persist the final summary and transition the parent Job.\"\"\"\n try:\n from langflow.api.utils import ingestion_run_service\n from langflow.services.database.models.jobs.model import JobStatus\n from langflow.services.deps import get_job_service\n except ImportError as exc:\n self.log(f\"Run-history wiring unavailable; ingestion-run finalize skipped ({exc}).\")\n return\n\n try:\n await ingestion_run_service.finalize_run(\n run_id,\n summary=summary,\n status=status,\n error_message=error_message,\n )\n if job_id is not None:\n terminal_status = JobStatus.COMPLETED if status is not IngestionRunStatus.FAILED else JobStatus.FAILED\n await get_job_service().update_job_status(job_id, terminal_status, finished_timestamp=True)\n except Exception as exc: # noqa: BLE001 — telemetry must never re-raise\n self.log(f\"Could not finalize ingestion-run tracking: {exc}\")\n\n async def _record_kb_status(\n self,\n kb_record_id: uuid.UUID,\n status_value: str,\n *,\n failure_reason: str | None = None,\n ) -> None:\n \"\"\"Mirror Path A's KB-row status transitions.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n from langflow.services.database.models.knowledge_base.model import KnowledgeBaseStatus\n except ImportError:\n return\n try:\n await knowledge_base_service.update_status(\n kb_record_id,\n status=KnowledgeBaseStatus(status_value),\n failure_reason=failure_reason,\n )\n except Exception as exc: # noqa: BLE001\n self.log(f\"Could not update KB status to {status_value}: {exc}\")\n\n async def _record_kb_stats(self, kb_record_id: uuid.UUID, kb_path: Path) -> None:\n \"\"\"Push freshly-refreshed metrics from embedding_metadata.json onto the DB row.\"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n except ImportError:\n self.log(\"knowledge_base_service unavailable; KB stats will not sync to DB row.\")\n return\n metadata_path = kb_path / \"embedding_metadata.json\"\n if not metadata_path.exists():\n self.log(f\"No embedding_metadata.json at {metadata_path}; skipping KB stats sync.\")\n return\n try:\n metadata = json.loads(metadata_path.read_text())\n except (OSError, json.JSONDecodeError) as exc:\n self.log(f\"Could not read KB metadata for stats sync: {exc}\")\n return\n chunks = int(metadata.get(\"chunks\", 0) or 0)\n words = int(metadata.get(\"words\", 0) or 0)\n characters = int(metadata.get(\"characters\", 0) or 0)\n try:\n await knowledge_base_service.update_stats(\n kb_record_id,\n chunks=chunks,\n words=words,\n characters=characters,\n size_bytes=int(metadata.get(\"size\", 0) or 0),\n source_types=list(metadata.get(\"source_types\") or []),\n chunk_size=metadata.get(\"chunk_size\"),\n chunk_overlap=metadata.get(\"chunk_overlap\"),\n separator=metadata.get(\"separator\"),\n )\n self.log(f\"Synced KB stats to DB row {kb_record_id}: chunks={chunks} words={words} characters={characters}\")\n except Exception as exc: # noqa: BLE001\n self.log(f\"Could not sync KB stats to DB row {kb_record_id}: {exc}\")\n\n def _parse_user_metadata_dict(self) -> dict[str, Any]:\n \"\"\"Decode ``metadata_json`` to a dict, or ``{}`` on any error.\"\"\"\n raw = getattr(self, \"metadata_json\", None)\n if not raw:\n return {}\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return {}\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError):\n return {}\n return decoded if isinstance(decoded, dict) else {}\n\n # =====================================================================\n # RETRIEVAL CODE PATH\n # =====================================================================\n @property\n def _user_uuid(self) -> uuid.UUID | None:\n \"\"\"Return self.user_id as a UUID, converting from str if necessary.\"\"\"\n if not self.user_id:\n return None\n return self.user_id if isinstance(self.user_id, uuid.UUID) else uuid.UUID(self.user_id)\n\n def _get_kb_metadata(self, kb_path: Path, *, require_api_key: bool = False) -> dict:\n \"\"\"Load the knowledge base's embedding metadata file.\"\"\"\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n return load_kb_metadata(\n kb_path,\n log_label=f\"knowledge base '{self.knowledge_base}'\",\n require_api_key=require_api_key,\n )\n\n async def _backend_from_record(self) -> tuple[str, dict[str, Any]] | None:\n \"\"\"Read ``(backend_type, backend_config)`` off the KB's database row.\n\n Returns ``None`` when there is no row to read — a legacy KB predating the\n record, or bare lfx with no langflow installed. A lookup that *fails* is a\n different situation and is allowed to propagate: guessing a backend after a\n database error is how vectors end up in a store nothing queries.\n \"\"\"\n try:\n from langflow.api.utils import knowledge_base_service\n except ImportError:\n return None\n\n user_uuid = self._user_uuid\n if user_uuid is None:\n return None\n record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base)\n if record is None:\n return None\n return record.backend_type or BackendType.CHROMA.value, record.backend_config or {}\n\n async def _resolve_backend_config(self, kb_path: Path) -> tuple[str, dict[str, Any]]:\n \"\"\"Resolve this KB's backend — database row first, on-disk sidecar second.\n\n Ingestion used to read the sidecar while retrieval read the row. Across\n replicas those disagree: a pod without the sidecar fell back to local\n Chroma and wrote vectors there, while queries followed the row to the\n configured remote backend and found nothing — silently. Both paths resolve\n here now, and a backend that cannot be resolved raises rather than\n defaulting to local storage.\n \"\"\"\n resolved = await self._backend_from_record()\n if resolved is None:\n resolved = self._get_backend_from_metadata(kb_path)\n if resolved is None:\n msg = (\n f\"Cannot determine the vector-store backend for knowledge base \"\n f\"'{self.knowledge_base}': it has no database record and no readable \"\n f\"embedding metadata. Refusing to fall back to local storage, which \"\n f\"would write to a different store than queries read from.\"\n )\n raise ValueError(msg)\n return resolved\n\n def _resolve_model_selection(self, metadata: dict[str, Any]) -> list[dict[str, Any]]:\n \"\"\"Resolve the ``get_embeddings``-compatible model selection from metadata.\"\"\"\n model_selection = metadata.get(\"model_selection\")\n if model_selection:\n selection_list = [model_selection] if isinstance(model_selection, dict) else list(model_selection)\n return [self._hydrate_model_metadata(entry) for entry in selection_list]\n\n embedding_model_name = metadata.get(\"embedding_model\")\n embedding_provider = metadata.get(\"embedding_provider\", \"Unknown\")\n if not embedding_model_name:\n msg = (\n f\"Knowledge base '{self.knowledge_base}' has no embedding model recorded; \"\n \"re-create it with a supported embedding model.\"\n )\n raise ValueError(msg)\n\n match = self._find_catalog_entry(embedding_model_name)\n if match is None:\n msg = (\n f\"Embedding model '{embedding_model_name}' (provider '{embedding_provider}') \"\n \"recorded for this knowledge base is no longer available in the model registry. \"\n \"Please re-create the knowledge base with a supported embedding model.\"\n )\n raise ValueError(msg)\n return [match]\n\n def _hydrate_model_metadata(self, entry: dict[str, Any]) -> dict[str, Any]:\n \"\"\"Fill in ``metadata.embedding_class`` / ``param_mapping`` if missing.\"\"\"\n entry_metadata = entry.get(\"metadata\") or {}\n has_class = bool(entry_metadata.get(\"embedding_class\"))\n has_mapping = bool(entry_metadata.get(\"param_mapping\"))\n if has_class and has_mapping:\n return entry\n\n model_name = entry.get(\"name\")\n if not model_name:\n return entry\n\n catalog_entry = self._find_catalog_entry(model_name)\n if catalog_entry is None:\n return entry\n\n catalog_metadata = catalog_entry.get(\"metadata\") or {}\n merged_metadata = {**catalog_metadata, **entry_metadata}\n return {**entry, \"metadata\": merged_metadata}\n\n def _find_catalog_entry(self, model_name: str) -> dict[str, Any] | None:\n \"\"\"Look up an embedding model by name in the unified-models catalog.\"\"\"\n options = get_embedding_model_options(user_id=self.user_id)\n return next((o for o in options if o.get(\"name\") == model_name), None)\n\n async def retrieve_data(self) -> DataFrame:\n \"\"\"Retrieve data from the selected knowledge base.\n\n Annotation narrowed to ``DataFrame`` to keep this output's handle\n type list at ``[\"Table\"]`` only; widening to a union surfaces\n ``[\"Table\", \"JSON\"]`` on the API and breaks every starter-project\n edge whose sourceHandle was stored against a single-type handle\n (BUG-02). The cross-mode defensive fallback may still return a\n ``Data`` at runtime — Python ignores return annotations, so\n nothing breaks.\n \"\"\"\n if not _is_retrieve_mode(getattr(self, \"mode\", MODE_INGEST)):\n return await self.build_kb_info()\n raise_error_if_astra_cloud_disable_component(astra_error_msg)\n\n # Lazy import: langflow's user/DB models aren't part of lfx's\n # standalone install, so ``lfx run <starter>.json`` can't resolve\n # this symbol at module import time. Deferring to use keeps the\n # component importable in both environments.\n from langflow.services.database.models.user.crud import get_user_by_id\n\n async with session_scope() as db:\n if not self.user_id:\n msg = \"User ID is required for fetching Knowledge Base data.\"\n raise ValueError(msg)\n current_user = await get_user_by_id(db, self.user_id)\n if not current_user:\n msg = f\"User with ID {self.user_id} not found.\"\n raise ValueError(msg)\n kb_user = current_user.username\n kb_path = self._resolve_kb_path(_get_knowledge_bases_root_path(), kb_user, self.knowledge_base)\n\n component_api_key = self.api_key if getattr(self, \"api_key\", None) else None\n needs_stored_key = not component_api_key\n metadata = self._get_kb_metadata(kb_path, require_api_key=needs_stored_key)\n if not metadata:\n msg = f\"Metadata not found for knowledge base: {self.knowledge_base}. Ensure it has been indexed.\"\n raise ValueError(msg)\n\n api_key = component_api_key or metadata.get(\"api_key\")\n model_selection = self._resolve_model_selection(metadata)\n chunk_size = metadata.get(\"chunk_size\")\n embedding_function = get_embeddings(\n model=model_selection,\n user_id=self.user_id,\n api_key=api_key,\n chunk_size=chunk_size,\n )\n\n backend_type, backend_config = await self._resolve_backend_config(kb_path)\n backend = create_backend(\n backend_type,\n kb_name=self.knowledge_base,\n kb_path=kb_path,\n backend_config=backend_config,\n embedding_function=embedding_function,\n user_id=self.user_id,\n )\n try:\n user_metadata_filter = _parse_metadata_filter(getattr(self, \"metadata_filter\", None))\n use_scores = bool(self.search_query)\n search_k = self.top_k * 4 if user_metadata_filter else self.top_k\n results = await backend.similarity_search(\n query=self.search_query or \"\",\n k=search_k,\n with_scores=use_scores,\n )\n if user_metadata_filter:\n results = [\n (doc, score) for doc, score in results if _chunk_matches_filter(doc.metadata, user_metadata_filter)\n ]\n results = results[: self.top_k]\n\n embeddings_by_key: dict[tuple[str, str], list[float]] = {}\n if self.include_embeddings and results:\n # Join each retrieved chunk to its stored embedding. We key on\n # ``_id`` when present and fall back to page content otherwise,\n # so KBs populated by direct file upload — whose chunks carry no\n # ``_id`` — still resolve. See ``_embedding_match_key``.\n wanted_keys = {_embedding_match_key(doc.page_content, doc.metadata) for doc, _score in results}\n async for batch in backend.iter_documents(include_embeddings=True):\n for entry in batch:\n if entry.embedding is None:\n continue\n key = _embedding_match_key(entry.content, entry.metadata)\n if key in wanted_keys:\n embeddings_by_key[key] = entry.embedding\n if len(embeddings_by_key) == len(wanted_keys):\n break\n\n data_list: list[Data] = []\n for doc, score in results:\n kwargs: dict[str, Any] = {\"content\": doc.page_content}\n if use_scores:\n kwargs[\"_score\"] = -1 * score\n if self.include_metadata:\n kwargs.update(doc.metadata)\n if self.include_embeddings:\n kwargs[\"_embeddings\"] = embeddings_by_key.get(_embedding_match_key(doc.page_content, doc.metadata))\n data_list.append(Data(**kwargs))\n\n return DataFrame(data=data_list)\n finally:\n await backend.teardown()\n\n\ndef _embedding_match_key(content: str, metadata: dict[str, Any] | None) -> tuple[str, str]:\n \"\"\"Build a stable key for aligning a retrieved chunk with its stored embedding.\n\n Embeddings are gathered separately (via ``iter_documents``) and then joined\n back onto the search results. Component-driven ingestion stamps a\n content-hash ``_id`` on every chunk, but direct file-upload ingestion\n (``KBIngestionHelper.perform_ingestion``) does not. Keying the join purely on\n ``_id`` therefore left every upload-populated KB with ``_embeddings: None``.\n\n We prefer ``_id`` when present (so legitimately distinct chunks that happen to\n share text aren't collapsed) and fall back to the chunk's page content\n otherwise. The content fallback is exact for the embedding use case: two\n chunks with identical text necessarily share the same embedding vector. The\n leading ``\"id\"`` / ``\"content\"`` tag namespaces the two key spaces so a mixed\n KB (some chunks with ``_id``, some without) never cross-matches.\n \"\"\"\n doc_id = metadata.get(\"_id\") if metadata else None\n if doc_id:\n return (\"id\", str(doc_id))\n return (\"content\", content or \"\")\n\n\ndef _parse_metadata_filter(raw: str | None) -> dict[str, list[str]]:\n \"\"\"Decode the ``metadata_filter`` input into a {key: [values]} map.\n\n Empty or malformed input maps to an empty filter so retrieval falls back\n to the unfiltered path. JSON errors are swallowed here rather than raised:\n surfacing component-config errors at the canvas node would break a flow\n run for what is meant to be an optional refinement.\n \"\"\"\n if not raw:\n return {}\n text = raw.strip() if isinstance(raw, str) else raw\n if not text:\n return {}\n try:\n decoded = json.loads(text)\n except (TypeError, json.JSONDecodeError):\n logger.warning(\"KnowledgeComponent: metadata_filter is not valid JSON; ignoring filter.\")\n return {}\n if not isinstance(decoded, dict):\n logger.warning(\"KnowledgeComponent: metadata_filter must be a JSON object; ignoring filter.\")\n return {}\n result: dict[str, list[str]] = {}\n for key, value in decoded.items():\n if not isinstance(key, str):\n continue\n if isinstance(value, list):\n result[key] = [str(entry) for entry in value]\n else:\n result[key] = [str(value)]\n return result\n\n\ndef _chunk_matches_filter(metadata: dict[str, Any] | None, filt: dict[str, list[str]]) -> bool:\n \"\"\"AND across keys, OR within key values, mirroring the chunks endpoint.\"\"\"\n if not filt:\n return True\n if not metadata:\n return False\n raw = metadata.get(\"source_metadata\")\n if not raw:\n return False\n try:\n stored = json.loads(raw) if isinstance(raw, str) else raw\n except json.JSONDecodeError:\n return False\n if not isinstance(stored, dict):\n return False\n for key, expected_values in filt.items():\n actual = stored.get(key)\n if actual is None:\n return False\n actual_set = {str(entry) for entry in actual} if isinstance(actual, list) else {str(actual)}\n if not actual_set & set(expected_values):\n return False\n return True\n" | |||
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Do not limit client-side metadata filtering to four times top_k.
When metadata_filter is present, retrieval fetches only top_k * 4 unfiltered results and filters them afterward. Matching documents ranked below that arbitrary window are omitted, producing incomplete or empty results despite valid matches. Pass the filter to the backend where supported; otherwise continue fetching until enough matches are found or the store is exhausted.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 8118, The retrieval
flow currently limits client-side metadata filtering to a fixed top_k
multiplier, which can omit valid lower-ranked matches. Update the retrieval
method that handles metadata_filter to pass the filter to backends supporting
native filtering; for unsupported backends, continue paginating or expanding
unfiltered fetches until top_k matches are collected or the store is exhausted.
Preserve existing result ordering and behavior when no filter is provided, using
_parse_metadata_filter and _chunk_matches_filter for client-side filtering.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Avoid local Chroma lookups for Chroma Cloud ingestion.
"chroma_cloud" is normalized to BackendType.CHROMA.value, leaving existing_ids unset. _convert_df_to_data_objects() then constructs a local Chroma(persist_directory=...) to detect duplicates, even though writes go to Chroma Cloud. This can query the wrong store or fail on replicas. Populate duplicate IDs through the resolved backend, or explicitly distinguish local Chroma from Chroma Cloud.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 8118, Update
_convert_df_to_data_objects so Chroma Cloud ingestion does not perform local
Chroma lookups for duplicate detection. Preserve the distinction between local
Chroma and Chroma Cloud after backend normalization, and obtain existing IDs
through the resolved backend when using Chroma Cloud. Ensure existing_ids is
populated before duplicate filtering without constructing Chroma with a local
persist_directory for cloud writes.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Resolve Knowledge embedding configuration from the database before requiring the sidecar.
retrieve_data() loads embedding_metadata.json and resolves the embedding model before calling the new database-backed backend resolver. On a replica with a database record but no local sidecar, retrieval still fails, so remote-backed Knowledge Bases are not actually replica-safe. Resolve model selection from the KB record first and retain the sidecar only as a legacy fallback or for required stored credentials.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 8118, Update
KnowledgeComponent.retrieve_data() to resolve the embedding model and
configuration from the database-backed knowledge-base record before attempting
to load embedding_metadata.json or requiring the local sidecar. Preserve sidecar
loading only as a legacy fallback or when stored credentials are required, so
remote-backed knowledge bases retrieve successfully on replicas without local
sidecar files.
| @@ -9625,7 +9617,7 @@ | |||
| "show": true, | |||
| "title_case": false, | |||
| "type": "code", | |||
| "value": "\"\"\"Memory Base retrieval component.\n\nQueries the auto-provisioned Chroma KB backing a Memory Base, scoped to the\ncurrent flow's request session. Additional option to filter by session_id if the\ndeveloper wants to turn that on. The component will auto filter based on session_id then.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport uuid\nfrom typing import TYPE_CHECKING, Any\n\nimport chromadb\nimport chromadb.api.client\nimport numpy as np\nfrom langchain_chroma import Chroma\nfrom langflow.api.utils.kb_helpers import KBIngestionHelper\nfrom langflow.services.database.models.memory_base.model import MemoryBase\nfrom langflow.services.database.models.user.crud import get_user_by_id\nfrom langflow.services.memory_base.kb_path_helpers import hash_session_id, validate_kb_path\nfrom sqlmodel import select\n\nfrom lfx.base.vectorstores.chroma_security import chroma_langchain_collection_kwargs\nfrom lfx.components.files_and_knowledge._kb_paths import (\n get_knowledge_bases_root_path,\n load_kb_metadata,\n)\nfrom lfx.custom import Component\nfrom lfx.io import BoolInput, DropdownInput, IntInput, MessageTextInput, Output\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.services.deps import session_scope\n\nif TYPE_CHECKING:\n from pathlib import Path\n\n from langflow.services.database.models.user.model import User\n from sqlmodel.ext.asyncio.session import AsyncSession\n\n\ndef _coerce_uuid(value: Any) -> uuid.UUID | None:\n if value is None:\n return None\n if isinstance(value, uuid.UUID):\n return value\n try:\n return uuid.UUID(str(value))\n except (ValueError, TypeError):\n return None\n\n\ndef _distance_to_similarity(distance: float) -> float:\n \"\"\"Chroma returns a distance; flip the sign so larger == more similar.\"\"\"\n return -1 * distance\n\n\ndef _to_python_scalar(value: Any) -> Any:\n \"\"\"Convert numpy scalars (int64, float64, bool_, …) to Python primitives.\n\n Chroma persists integer/float metadata as numpy scalars, which break JSON\n serialization when this component is consumed as an Agent tool — LangChain's\n tool-output path calls ``vars()`` / iterates the value, both of which fail\n on numpy C-extension scalars. Coerce at the boundary so downstream stays\n primitive-only.\n \"\"\"\n if isinstance(value, np.generic):\n return value.item()\n return value\n\n\nclass MemoryBaseComponent(Component):\n display_name = \"Memory Base\"\n description = (\n \"Retrieve long-term memory from a Memory Base attached to this workflow. \"\n \"When 'Filter by Session' is off, queries run across all sessions.\"\n )\n icon = \"brain\"\n name = \"MemoryBase\"\n\n inputs = [\n DropdownInput(\n name=\"memory_base\",\n display_name=\"Memory Base\",\n info=\"Memory Base whose captured conversation history will be searched.\",\n required=True,\n options=[],\n refresh_button=True,\n real_time_refresh=True,\n ),\n MessageTextInput(\n name=\"search_query\",\n display_name=\"Search Query\",\n info=\"Query used for semantic retrieval against the memory base.\",\n tool_mode=True,\n ),\n IntInput(\n name=\"top_k\",\n display_name=\"Top K Results\",\n info=\"Number of top results to return.\",\n value=5,\n advanced=True,\n required=False,\n ),\n BoolInput(\n name=\"include_metadata\",\n display_name=\"Include Metadata\",\n info=\"Include chunk metadata (session_id, sender, timestamp, …) on each row.\",\n value=True,\n advanced=True,\n ),\n BoolInput(\n name=\"filter_by_session\",\n display_name=\"Filter by Session\",\n info=(\n \"If enabled, only memories from the current session will be retrieved. \"\n \"Disable to allow retrieval across every session ingested into this \"\n \"Memory Base (useful for cross-conversation recall).\"\n ),\n value=True,\n advanced=True,\n ),\n ]\n\n outputs = [\n Output(\n name=\"retrieve_data\",\n display_name=\"Results\",\n method=\"retrieve_memory\",\n info=(\n \"Returns matching memory chunks. Scoped to the current session by \"\n \"default; turn 'Filter by Session' off to retrieve across sessions.\"\n ),\n ),\n ]\n\n def _build_where_clause(self, *, session_id: str | None = None) -> dict | None:\n \"\"\"Compose the Chroma ``where`` clause based on opt-in filters and manual params.\n\n Uses the canonical ``$eq`` operator form rather than the implicit\n ``{\"key\": \"value\"}`` shorthand. Both are accepted by chromadb, but the\n explicit form is unambiguous across versions and tooling.\n \"\"\"\n predicates: list[dict] = []\n # Defensive bool() — BoolInput coerces strings, but if this attribute is\n # ever overridden externally with a non-bool value, ``\"false\"`` would be\n # truthy and silently disable the toggle.\n if bool(self.filter_by_session) and session_id:\n predicates.append({\"session_id\": {\"$eq\": str(session_id)}})\n\n if not predicates:\n return None\n if len(predicates) == 1:\n return predicates[0]\n return {\"$and\": predicates}\n\n async def update_build_config(self, build_config, field_value, field_name=None): # noqa: ARG002\n if field_name != \"memory_base\":\n return build_config\n\n flow_id = _coerce_uuid(self._get_runtime_or_frontend_node_attr(\"flow_id\"))\n user_uuid = _coerce_uuid(self.user_id)\n if flow_id is None or user_uuid is None:\n build_config[\"memory_base\"][\"options\"] = []\n build_config[\"memory_base\"][\"value\"] = None\n return build_config\n\n # At design time self.user_id == the flow developer == MB owner, so this\n # filters to the same set a Flow-lookup would return but without relying\n # on the Flow row being persisted yet.\n async with session_scope() as db:\n stmt = select(MemoryBase).where(\n MemoryBase.flow_id == flow_id,\n MemoryBase.user_id == user_uuid,\n )\n mbs = list((await db.exec(stmt)).all())\n\n options = sorted(mb.name for mb in mbs)\n build_config[\"memory_base\"][\"options\"] = options\n if build_config[\"memory_base\"].get(\"value\") not in options:\n build_config[\"memory_base\"][\"value\"] = None\n return build_config\n\n async def _resolve_attached_mb(\n self,\n db: AsyncSession,\n selected: str,\n flow_id: uuid.UUID,\n ) -> tuple[MemoryBase, User]:\n \"\"\"Look up the MB row scoped to the current flow and resolve its owner.\"\"\"\n mb = (\n await db.exec(\n select(MemoryBase).where(\n MemoryBase.name == selected,\n MemoryBase.flow_id == flow_id,\n )\n )\n ).first()\n if mb is None:\n msg = f\"Memory Base '{selected}' is not attached to this flow.\"\n raise ValueError(msg)\n\n owner = await get_user_by_id(db, mb.user_id)\n if owner is None:\n msg = \"Memory Base owner account could not be resolved.\"\n raise ValueError(msg)\n return mb, owner\n\n def _resolve_kb_location(self, owner_username: str, kb_name: str) -> Path:\n \"\"\"Build and validate the on-disk KB path for the given owner.\"\"\"\n kb_root = get_knowledge_bases_root_path()\n kb_path = kb_root / owner_username / kb_name\n try:\n validate_kb_path(kb_root, kb_path)\n except ValueError as exc:\n msg = \"Memory Base path is not accessible.\"\n raise ValueError(msg) from exc\n return kb_path\n\n async def _build_chroma(\n self,\n kb_path: Path,\n owner: User,\n metadata: dict,\n kb_name: str,\n ) -> Chroma:\n \"\"\"Construct a Chroma client wired to the KB's embedding function.\"\"\"\n provider = metadata.get(\"embedding_provider\")\n model = metadata.get(\"embedding_model\")\n embedding_function = await KBIngestionHelper.build_embeddings(provider, model, owner)\n\n chromadb.api.client.SharedSystemClient.clear_system_cache()\n return Chroma(\n persist_directory=str(kb_path),\n embedding_function=embedding_function,\n collection_name=kb_name,\n **chroma_langchain_collection_kwargs(),\n )\n\n def _format_results(self, results: list[tuple]) -> DataFrame:\n \"\"\"Convert Chroma (doc, score) tuples into the component's DataFrame output.\n\n Metadata values are coerced from numpy scalars to Python primitives so the\n resulting DataFrame is JSON-serializable when the component is invoked as\n an Agent tool.\n \"\"\"\n data_list: list[Data] = []\n for doc, score in results:\n kwargs: dict = {\"content\": doc.page_content}\n if self.search_query:\n kwargs[\"_score\"] = _to_python_scalar(_distance_to_similarity(score))\n if self.include_metadata:\n for key, value in (doc.metadata or {}).items():\n kwargs[key] = _to_python_scalar(value)\n data_list.append(Data(**kwargs))\n return DataFrame(data=data_list)\n\n async def retrieve_memory(self) -> DataFrame:\n \"\"\"Retrieve chunks from the selected Memory Base.\n\n Scoped to the current ``session_id`` when ``filter_by_session`` is true; when\n false, every chunk in the Memory Base is queryable so the agent can recall\n context from prior conversations across all sessions.\n \"\"\"\n session_id = getattr(self.graph, \"session_id\", None)\n if bool(self.filter_by_session) and not session_id:\n # Only required when filtering is on, since the value gates the where clause.\n msg = (\n \"A session_id is required on the flow request when 'Filter by Session' \"\n \"is enabled — disable the toggle to allow cross-session retrieval.\"\n )\n raise ValueError(msg)\n\n flow_id = _coerce_uuid(getattr(self.graph, \"flow_id\", None))\n if flow_id is None:\n msg = \"flow_id is not available on the graph context; Memory Base retrieval is unavailable.\"\n raise ValueError(msg)\n\n selected = self.memory_base\n if not selected:\n msg = \"No Memory Base is selected.\"\n raise ValueError(msg)\n\n async with session_scope() as db:\n mb, owner = await self._resolve_attached_mb(db, selected, flow_id)\n owner_username = owner.username\n kb_name = mb.kb_name\n\n kb_path = self._resolve_kb_location(owner_username, kb_name)\n metadata = load_kb_metadata(kb_path, log_label=f\"memory base '{selected}'\")\n if not metadata:\n msg = f\"Memory Base '{selected}' has no embedding metadata on disk.\"\n raise ValueError(msg)\n\n chroma = await self._build_chroma(kb_path, owner, metadata, kb_name)\n where = self._build_where_clause(session_id=session_id)\n\n logger.debug(\n \"MemoryBase retrieval mb=%s session_hash=%s where=%s top_k=%s\",\n selected,\n hash_session_id(session_id) if session_id else \"<none>\",\n where,\n self.top_k,\n )\n\n if not self.search_query:\n # Embedding providers may reject empty input; skip the round-trip entirely.\n return DataFrame(data=[])\n\n results = chroma.similarity_search_with_score(\n query=self.search_query,\n k=self.top_k,\n filter=where,\n )\n return self._format_results(results)\n" | |||
| "value": "\"\"\"Memory Base retrieval component.\n\nQueries the auto-provisioned Chroma KB backing a Memory Base, scoped to the\ncurrent flow's request session. Additional option to filter by session_id if the\ndeveloper wants to turn that on. The component will auto filter based on session_id then.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport uuid\nfrom typing import TYPE_CHECKING, Any\n\nimport numpy as np\nfrom langflow.api.utils.kb_helpers import KBIngestionHelper, resolve_backend_selection, resolve_embedding_selection\nfrom langflow.services.database.models.memory_base.model import MemoryBase\nfrom langflow.services.database.models.user.crud import get_user_by_id\nfrom langflow.services.memory_base.kb_path_helpers import hash_session_id, validate_kb_path\nfrom sqlmodel import select\n\nfrom lfx.base.knowledge_bases.backends import create_backend\nfrom lfx.components.files_and_knowledge._kb_paths import (\n get_knowledge_bases_root_path,\n)\nfrom lfx.custom import Component\nfrom lfx.io import BoolInput, DropdownInput, IntInput, MessageTextInput, Output\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.schema.dataframe import DataFrame\nfrom lfx.services.deps import session_scope\n\nif TYPE_CHECKING:\n from pathlib import Path\n\n from langflow.services.database.models.user.model import User\n from sqlmodel.ext.asyncio.session import AsyncSession\n\n from lfx.base.knowledge_bases.backends.base import BaseVectorStoreBackend\n\n\ndef _coerce_uuid(value: Any) -> uuid.UUID | None:\n if value is None:\n return None\n if isinstance(value, uuid.UUID):\n return value\n try:\n return uuid.UUID(str(value))\n except (ValueError, TypeError):\n return None\n\n\ndef _distance_to_similarity(distance: float) -> float:\n \"\"\"Chroma returns a distance; flip the sign so larger == more similar.\"\"\"\n return -1 * distance\n\n\ndef _to_python_scalar(value: Any) -> Any:\n \"\"\"Convert numpy scalars (int64, float64, bool_, …) to Python primitives.\n\n Chroma persists integer/float metadata as numpy scalars, which break JSON\n serialization when this component is consumed as an Agent tool — LangChain's\n tool-output path calls ``vars()`` / iterates the value, both of which fail\n on numpy C-extension scalars. Coerce at the boundary so downstream stays\n primitive-only.\n \"\"\"\n if isinstance(value, np.generic):\n return value.item()\n return value\n\n\nclass MemoryBaseComponent(Component):\n display_name = \"Memory Base\"\n description = (\n \"Retrieve long-term memory from a Memory Base attached to this workflow. \"\n \"When 'Filter by Session' is off, queries run across all sessions.\"\n )\n icon = \"brain\"\n name = \"MemoryBase\"\n\n inputs = [\n DropdownInput(\n name=\"memory_base\",\n display_name=\"Memory Base\",\n info=\"Memory Base whose captured conversation history will be searched.\",\n required=True,\n options=[],\n refresh_button=True,\n real_time_refresh=True,\n ),\n MessageTextInput(\n name=\"search_query\",\n display_name=\"Search Query\",\n info=\"Query used for semantic retrieval against the memory base.\",\n tool_mode=True,\n ),\n IntInput(\n name=\"top_k\",\n display_name=\"Top K Results\",\n info=\"Number of top results to return.\",\n value=5,\n advanced=True,\n required=False,\n ),\n BoolInput(\n name=\"include_metadata\",\n display_name=\"Include Metadata\",\n info=\"Include chunk metadata (session_id, sender, timestamp, …) on each row.\",\n value=True,\n advanced=True,\n ),\n BoolInput(\n name=\"filter_by_session\",\n display_name=\"Filter by Session\",\n info=(\n \"If enabled, only memories from the current session will be retrieved. \"\n \"Disable to allow retrieval across every session ingested into this \"\n \"Memory Base (useful for cross-conversation recall).\"\n ),\n value=True,\n advanced=True,\n ),\n ]\n\n outputs = [\n Output(\n name=\"retrieve_data\",\n display_name=\"Results\",\n method=\"retrieve_memory\",\n info=(\n \"Returns matching memory chunks. Scoped to the current session by \"\n \"default; turn 'Filter by Session' off to retrieve across sessions.\"\n ),\n ),\n ]\n\n def _build_where_clause(self, *, session_id: str | None = None) -> dict | None:\n \"\"\"Compose the metadata filter based on opt-in filters and manual params.\n\n Emits a flat ``{key: value}`` dict, which is the contract every\n ``BaseVectorStoreBackend`` accepts: Chroma reads it as ``$eq`` shorthand\n and OpenSearch translates it into a bool query. Chroma's explicit\n ``{\"$eq\": ...}`` operator form is NOT portable — a remote backend would\n treat the operator dict as a literal value and silently match nothing.\n \"\"\"\n predicates: dict = {}\n # Defensive bool() — BoolInput coerces strings, but if this attribute is\n # ever overridden externally with a non-bool value, ``\"false\"`` would be\n # truthy and silently disable the toggle.\n if bool(self.filter_by_session) and session_id:\n predicates[\"session_id\"] = str(session_id)\n\n return predicates or None\n\n async def update_build_config(self, build_config, field_value, field_name=None): # noqa: ARG002\n if field_name != \"memory_base\":\n return build_config\n\n flow_id = _coerce_uuid(self._get_runtime_or_frontend_node_attr(\"flow_id\"))\n user_uuid = _coerce_uuid(self.user_id)\n if flow_id is None or user_uuid is None:\n build_config[\"memory_base\"][\"options\"] = []\n build_config[\"memory_base\"][\"value\"] = None\n return build_config\n\n # At design time self.user_id == the flow developer == MB owner, so this\n # filters to the same set a Flow-lookup would return but without relying\n # on the Flow row being persisted yet.\n async with session_scope() as db:\n stmt = select(MemoryBase).where(\n MemoryBase.flow_id == flow_id,\n MemoryBase.user_id == user_uuid,\n )\n mbs = list((await db.exec(stmt)).all())\n\n options = sorted(mb.name for mb in mbs)\n build_config[\"memory_base\"][\"options\"] = options\n if build_config[\"memory_base\"].get(\"value\") not in options:\n build_config[\"memory_base\"][\"value\"] = None\n return build_config\n\n async def _resolve_attached_mb(\n self,\n db: AsyncSession,\n selected: str,\n flow_id: uuid.UUID,\n ) -> tuple[MemoryBase, User]:\n \"\"\"Look up the MB row scoped to the current flow and resolve its owner.\"\"\"\n mb = (\n await db.exec(\n select(MemoryBase).where(\n MemoryBase.name == selected,\n MemoryBase.flow_id == flow_id,\n )\n )\n ).first()\n if mb is None:\n msg = f\"Memory Base '{selected}' is not attached to this flow.\"\n raise ValueError(msg)\n\n owner = await get_user_by_id(db, mb.user_id)\n if owner is None:\n msg = \"Memory Base owner account could not be resolved.\"\n raise ValueError(msg)\n return mb, owner\n\n def _resolve_kb_location(self, owner_username: str, kb_name: str) -> Path:\n \"\"\"Build and validate the on-disk KB path for the given owner.\"\"\"\n kb_root = get_knowledge_bases_root_path()\n kb_path = kb_root / owner_username / kb_name\n try:\n validate_kb_path(kb_root, kb_path)\n except ValueError as exc:\n msg = \"Memory Base path is not accessible.\"\n raise ValueError(msg) from exc\n return kb_path\n\n async def _build_backend(\n self,\n kb_path: Path,\n owner: User,\n kb_name: str,\n ) -> BaseVectorStoreBackend:\n \"\"\"Construct the KB's configured backend, wired to its embedding function.\n\n Both the embedding config and the backend are resolved from the\n ``knowledge_base`` row (sidecar only as a legacy fallback), so a Memory\n Base provisioned on OpenSearch or Chroma Cloud is queried there — and with\n the right embedding model — even on a replica whose local disk never held\n the KB directory or its ``embedding_metadata.json``.\n \"\"\"\n provider, model = await resolve_embedding_selection(user_id=owner.id, kb_name=kb_name, kb_path=kb_path)\n embedding_function = await KBIngestionHelper.build_embeddings(provider, model, owner)\n\n backend_type, backend_config = await resolve_backend_selection(\n user_id=owner.id, kb_name=kb_name, kb_path=kb_path\n )\n backend = create_backend(\n backend_type,\n kb_name=kb_name,\n kb_path=kb_path,\n backend_config=backend_config,\n embedding_function=embedding_function,\n user_id=owner.id,\n )\n await backend.ensure_ready()\n return backend\n\n def _format_results(self, results: list[tuple]) -> DataFrame:\n \"\"\"Convert Chroma (doc, score) tuples into the component's DataFrame output.\n\n Metadata values are coerced from numpy scalars to Python primitives so the\n resulting DataFrame is JSON-serializable when the component is invoked as\n an Agent tool.\n \"\"\"\n data_list: list[Data] = []\n for doc, score in results:\n kwargs: dict = {\"content\": doc.page_content}\n if self.search_query:\n kwargs[\"_score\"] = _to_python_scalar(_distance_to_similarity(score))\n if self.include_metadata:\n for key, value in (doc.metadata or {}).items():\n kwargs[key] = _to_python_scalar(value)\n data_list.append(Data(**kwargs))\n return DataFrame(data=data_list)\n\n async def retrieve_memory(self) -> DataFrame:\n \"\"\"Retrieve chunks from the selected Memory Base.\n\n Scoped to the current ``session_id`` when ``filter_by_session`` is true; when\n false, every chunk in the Memory Base is queryable so the agent can recall\n context from prior conversations across all sessions.\n \"\"\"\n session_id = getattr(self.graph, \"session_id\", None)\n if bool(self.filter_by_session) and not session_id:\n # Only required when filtering is on, since the value gates the where clause.\n msg = (\n \"A session_id is required on the flow request when 'Filter by Session' \"\n \"is enabled — disable the toggle to allow cross-session retrieval.\"\n )\n raise ValueError(msg)\n\n flow_id = _coerce_uuid(getattr(self.graph, \"flow_id\", None))\n if flow_id is None:\n msg = \"flow_id is not available on the graph context; Memory Base retrieval is unavailable.\"\n raise ValueError(msg)\n\n selected = self.memory_base\n if not selected:\n msg = \"No Memory Base is selected.\"\n raise ValueError(msg)\n\n async with session_scope() as db:\n mb, owner = await self._resolve_attached_mb(db, selected, flow_id)\n owner_username = owner.username\n kb_name = mb.kb_name\n\n # ``kb_path`` is still needed for a local-Chroma backend (that's where its\n # vectors live), but embedding + backend now resolve from the DB row, so a\n # missing on-disk ``embedding_metadata.json`` is no longer fatal — a\n # remote-backed Memory Base is fully queryable on a replica with no local\n # KB directory.\n kb_path = self._resolve_kb_location(owner_username, kb_name)\n\n where = self._build_where_clause(session_id=session_id)\n\n logger.debug(\n \"MemoryBase retrieval mb=%s session_hash=%s where=%s top_k=%s\",\n selected,\n hash_session_id(session_id) if session_id else \"<none>\",\n where,\n self.top_k,\n )\n\n if not self.search_query:\n # Embedding providers may reject empty input; skip the round-trip entirely.\n return DataFrame(data=[])\n\n backend = await self._build_backend(kb_path, owner, kb_name)\n try:\n results = await backend.similarity_search(\n query=self.search_query,\n k=self.top_k,\n filter=where,\n with_scores=True,\n )\n finally:\n await backend.teardown()\n return self._format_results(results)\n" | |||
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Normalize boolean input before applying the session filter.
bool("false") evaluates to True, so serialized or externally overridden values can keep session filtering enabled when the user selected false. Normalize string values explicitly before both _build_where_clause() and retrieve_memory() evaluate filter_by_session.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 9620, Normalize
filter_by_session through a shared boolean conversion that recognizes string
values such as "false" as false, then use the normalized result in both
_build_where_clause() and retrieve_memory(). Ensure session_id validation and
where-clause construction follow the actual user-selected boolean state rather
than Python truthiness.
Comment on lines
+1565
to
+1607
| async def _backend_from_record(self) -> tuple[str, dict[str, Any]] | None: | ||
| """Read ``(backend_type, backend_config)`` off the KB's database row. | ||
|
|
||
| Returns ``None`` when there is no row to read — a legacy KB predating the | ||
| record, or bare lfx with no langflow installed. A lookup that *fails* is a | ||
| different situation and is allowed to propagate: guessing a backend after a | ||
| database error is how vectors end up in a store nothing queries. | ||
| """ | ||
| try: | ||
| from langflow.api.utils import knowledge_base_service | ||
| except ImportError: | ||
| return None | ||
|
|
||
| user_uuid = self._user_uuid | ||
| if user_uuid is None: | ||
| return BackendType.CHROMA.value, {} | ||
| record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base) | ||
| except Exception as exc: # noqa: BLE001 | ||
| logger.debug("KB record lookup failed: %s", exc) | ||
| return BackendType.CHROMA.value, {} | ||
|
|
||
| user_uuid = self._user_uuid | ||
| if user_uuid is None: | ||
| return None | ||
| record = await knowledge_base_service.get_by_user_and_name(user_uuid, self.knowledge_base) | ||
| if record is None: | ||
| return BackendType.CHROMA.value, {} | ||
| return ( | ||
| record.backend_type or BackendType.CHROMA.value, | ||
| record.backend_config or {}, | ||
| ) | ||
| return None | ||
| return record.backend_type or BackendType.CHROMA.value, record.backend_config or {} | ||
|
|
||
| async def _resolve_backend_config(self, kb_path: Path) -> tuple[str, dict[str, Any]]: | ||
| """Resolve this KB's backend — database row first, on-disk sidecar second. | ||
|
|
||
| Ingestion used to read the sidecar while retrieval read the row. Across | ||
| replicas those disagree: a pod without the sidecar fell back to local | ||
| Chroma and wrote vectors there, while queries followed the row to the | ||
| configured remote backend and found nothing — silently. Both paths resolve | ||
| here now, and a backend that cannot be resolved raises rather than | ||
| defaulting to local storage. | ||
| """ | ||
| resolved = await self._backend_from_record() | ||
| if resolved is None: | ||
| resolved = self._get_backend_from_metadata(kb_path) | ||
| if resolved is None: | ||
| msg = ( | ||
| f"Cannot determine the vector-store backend for knowledge base " | ||
| f"'{self.knowledge_base}': it has no database record and no readable " | ||
| f"embedding metadata. Refusing to fall back to local storage, which " | ||
| f"would write to a different store than queries read from." | ||
| ) | ||
| raise ValueError(msg) | ||
| return resolved |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd _kb_paths.py src/lfx --exec cat -n {}Repository: langflow-ai/langflow
Length of output: 5098
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== knowledge.py outline matches =="
ast-grep outline src/lfx/src/lfx/components/files_and_knowledge/knowledge.py --match KnowledgeBase --view expanded || true
echo "== relevant knowledge.py sections =="
sed -n '1480,1730p' src/lfx/src/lfx/components/files_and_knowledge/knowledge.py
echo "== _get_kb_metadata definitions/usages =="
rg -n "def _get_kb_metadata|async def _get_kb_metadata|load_kb_metadata|_get_kb_metadata\\(" src/lfx/src/lfx/components/files_and_knowledge -S
echo "== backend_type/config in KB model/DB =="
rg -n "backend_type|backend_config|embedding_metadata|embedding_metadata\\.json|embed_model|embedding_model|model" src/lfx/src/lfx/components/files_and_knowledge -SRepository: langflow-ai/langflow
Length of output: 38163
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate KnowledgeBase model/schema files =="
fd -i 'knowledge_base' src/lfx/langflow src --type f -x sh -c 'echo "--- $1"; sed -n "1,220p" "$1"' sh {}
echo "== inspect relevant memory_retrieval helpers =="
sed -n '180,260p' src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py
echo "== inspect model selection helpers in knowledge.py ingestion =="
sed -n '1160,1285p' src/lfx/src/lfx/components/files_and_knowledge/knowledge.py
echo "== check DB migration/stat fields around embedding/model/backend =="
rg -n "KnowledgeBase|knowledge_base|backend_type|backend_config|model_selection|embedding_model|embedding_metadata" src/lfx -S --glob '!**/__pycache__/**' | sed -n '1,220p'Repository: langflow-ai/langflow
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== exact KnowledgeBase model files =="
git ls-files | rg 'services/database/models/knowledge_base|knowledge_base.*model|knowledge_base.*\.py' | grep -E 'src/lfx/(langflow|src)/langflow' | sed -n '1,80p'
echo "== KnowledgeBase model and update_stats fields =="
for f in $(git ls-files | rg 'src/lfx/langflow/services/database/models/knowledge_base/model|src/lfx/langflow/services/database/models/knowledge_base.*\.py'); do
echo "--- ${f#"$PWD"/}"
rg -n 'class KnowledgeBase|backend_type|backend_config|model_selection|embedding_model|embedding_metadata|chunks|update_stats|stats|model_config|metadata' "$f" -A 10 -B 5
done
echo "== focused knowledge.py model selection ingestion =="
sed -n '1185,1285p' src/lfx/src/lfx/components/files_and_knowledge/knowledge.pyRepository: langflow-ai/langflow
Length of output: 196
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repository knowledge base files =="
git ls-files | rg 'knowledge_base' | sed -n '1,120p'
echo "== KnowledgeBase model and update_stats fields =="
while IFS= read -r f; do
echo "--- ${f}"
rg -n 'class KnowledgeBase|backend_type|backend_config|model_selection|embedding_model|embedding_metadata|chunks|update_stats|stats|model_config|metadata' "$f" -A 10 -B 5 || true
done < <(git ls-files | rg 'knowledge_base')Repository: langflow-ai/langflow
Length of output: 3217
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== knowledge base model/schema/service fields =="
for f in \
src/backend/base/langflow/services/database/models/knowledge_base/model.py \
src/backend/base/langflow/schema/knowledge_base.py
do
if [ -f "$f" ]; then
echo "--- $f"
rg -n 'backend_type|backend_config|model_selection|embedding_model|embedding_metadata|chunks|words|characters|size_bytes|update_stats|Metadata|Pydantic|SQLModel|Table|BaseModel' "$f" -A 12 -B 6 || true
fi
done
echo "== knowledge_base_service relevant definitions =="
if [ -f src/backend/base/langflow/api/utils/knowledge_base_service.py ]; then
echo "--- src/backend/base/langflow/api/utils/knowledge_base_service.py"
rg -n 'get_by_user_and_name|update_stats|KnowledgeBase|backend_type|backend_config|chunks|words|characters|size_bytes|model_selection|embedding_model|metadata|JSON' src/backend/base/langflow/api/utils/knowledge_base_service.py -A 20 -B 8 || true
fi
echo "== lfx component paths relevant imports for backend model =="
rg -n 'class Backend|backend_type|backend_config|BackendType|create_backend' src/lfx/src/lfx/base/knowledge_bases -S -A 8 -B 6 || trueRepository: langflow-ai/langflow
Length of output: 50377
Make embedding-model resolution DB-first like backend resolution.
load_kb_metadata only reads embedding_metadata.json from the KB directory, so retrieve_data can fail on a replica that has the KB row but no local sidecar. Use read_metadata/get_by_user_and_name for model_selection, chunk_size, chunk_overlap, and separator, falling back to disk only for legacy rows.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/files_and_knowledge/knowledge.py` around lines
1565 - 1607, Update retrieve_data’s embedding-model metadata resolution to query
the database record via knowledge_base_service.get_by_user_and_name first,
reusing read_metadata where appropriate for model_selection, chunk_size,
chunk_overlap, and separator. If no record exists, preserve the existing
embedding_metadata.json sidecar fallback for legacy knowledge bases, but do not
require local metadata when a database row is available.
Comment on lines
+216
to
+245
| async def _build_backend( | ||
| self, | ||
| kb_path: Path, | ||
| owner: User, | ||
| metadata: dict, | ||
| kb_name: str, | ||
| ) -> Chroma: | ||
| """Construct a Chroma client wired to the KB's embedding function.""" | ||
| provider = metadata.get("embedding_provider") | ||
| model = metadata.get("embedding_model") | ||
| ) -> BaseVectorStoreBackend: | ||
| """Construct the KB's configured backend, wired to its embedding function. | ||
|
|
||
| Both the embedding config and the backend are resolved from the | ||
| ``knowledge_base`` row (sidecar only as a legacy fallback), so a Memory | ||
| Base provisioned on OpenSearch or Chroma Cloud is queried there — and with | ||
| the right embedding model — even on a replica whose local disk never held | ||
| the KB directory or its ``embedding_metadata.json``. | ||
| """ | ||
| provider, model = await resolve_embedding_selection(user_id=owner.id, kb_name=kb_name, kb_path=kb_path) | ||
| embedding_function = await KBIngestionHelper.build_embeddings(provider, model, owner) | ||
|
|
||
| chromadb.api.client.SharedSystemClient.clear_system_cache() | ||
| return Chroma( | ||
| persist_directory=str(kb_path), | ||
| backend_type, backend_config = await resolve_backend_selection( | ||
| user_id=owner.id, kb_name=kb_name, kb_path=kb_path | ||
| ) | ||
| backend = create_backend( | ||
| backend_type, | ||
| kb_name=kb_name, | ||
| kb_path=kb_path, | ||
| backend_config=backend_config, | ||
| embedding_function=embedding_function, | ||
| collection_name=kb_name, | ||
| **chroma_langchain_collection_kwargs(), | ||
| user_id=owner.id, | ||
| ) | ||
| await backend.ensure_ready() | ||
| return backend |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Backend can leak if ensure_ready() fails after construction.
create_backend(...) followed by await backend.ensure_ready() isn't wrapped in try/except; if ensure_ready() raises (e.g. a failed connection to a remote backend), the constructed backend is never returned to retrieve_memory, so backend.teardown() (called only in the caller's finally, at Line 326) is never reached for this partially-initialized instance. This mirrors the same pattern in the embedded copies of this component and in KnowledgeComponent._create_vector_store — see the consolidated comment.
async def similarity_search(...): await self.ensure_ready() ... confirms ensure_ready() can perform real connection setup that should be torn down on failure.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py` around
lines 216 - 245, Update _build_backend to clean up the constructed backend when
await backend.ensure_ready() raises. Wrap the readiness call and return path
with exception handling that invokes backend.teardown() on failure, then
re-raises the original exception so retrieve_memory’s existing finally behavior
remains unchanged for successfully returned backends.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Backend construction/setup can leak resources on failure across multiple sites. All four sites create a vector-store backend and call its setup (ensure_ready(), and in the ingestion case iter_documents/add_documents) outside of any protection that guarantees teardown() runs if that step fails — the shared root cause is a missing try/except-with-teardown (or try/finally scoped around construction) at each site.
src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py#L216-L245: wrapcreate_backend(...)+await backend.ensure_ready()in_build_backendso a failure duringensure_ready()tears the partially-built backend down before re-raising.src/backend/base/langflow/initial_setup/starter_projects/Memory Chatbot.json#L1628-L1628: this JSON embeds the identical_build_backendsource; regenerate/re-sync this starter project's embedded component code once the source fix lands.src/backend/base/langflow/initial_setup/starter_projects/Price Deal Finder.json#L2288-L2288: same embedded copy; regenerate once the source fix lands.src/backend/base/langflow/initial_setup/starter_projects/Vector Store RAG.json#L1314-L1314: in the embeddedKnowledgeComponent._create_vector_store, wrapcreate_backend/ensure_ready/iter_documents/add_documentsin try/except that callsbackend.teardown()on failure before re-raising, sincebuild_kb_info's outertry/finallyteardown only starts after this method already returns successfully.
📍 Affects 4 files
src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py#L216-L245(this comment)src/backend/base/langflow/initial_setup/starter_projects/Memory Chatbot.json#L1628-L1628src/backend/base/langflow/initial_setup/starter_projects/Price Deal Finder.json#L2288-L2288src/backend/base/langflow/initial_setup/starter_projects/Vector Store RAG.json#L1314-L1314
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py` around
lines 216 - 245, Backend setup failures can leave partially initialized vector
stores undisposed. In
src/lfx/src/lfx/components/files_and_knowledge/memory_retrieval.py#L216-L245,
update _build_backend to tear down the created backend if ensure_ready() fails,
then re-raise; regenerate or re-sync the identical embedded _build_backend
source in src/backend/base/langflow/initial_setup/starter_projects/Memory
Chatbot.json#L1628-L1628 and Price Deal Finder.json#L2288-L2288. In
src/backend/base/langflow/initial_setup/starter_projects/Vector Store
RAG.json#L1314-L1314, update KnowledgeComponent._create_vector_store to tear
down the backend on failure across create_backend, ensure_ready, iter_documents,
or add_documents before re-raising.
Contributor
Frontend Unit Test Coverage ReportCoverage Summary
Unit Test Results
|
dkaushik94
changed the title
dkaushik94
changed the title
dkaushik94
force-pushed
the
mb-refactor-with-kb
branch
from
July 27, 2026 15:57
b06bf74 to
c5a65c7
Compare
github-actions
Bot
added
enhancement
github-actions
Bot
added
enhancement
github-actions
Bot
added
enhancement
github-actions
Bot
added
enhancement
erichare
self-requested a review
July 28, 2026 14:39
erichare
requested changes
Jul 28, 2026
erichare
left a comment
erichare
left a comment
Collaborator
There was a problem hiding this comment.
@dkaushik94 this is awesome. A few issues to address, could you take a look?
-
[P1] Isolate OpenSearch storage per Memory Base. Creation [copies the global backend configuration unchanged](
), including the single globallangflow/src/frontend/src/modals/createMemoryModal/useCreateMemoryModal.ts
Lines 242 to 245 in 3829f4a
[OPENSEARCH_INDEX_NAME](https://github.com/langflow-ai/langflow/blob/3829f4ac0f3df21cef7169f88afd60e1647c9b10/src/frontend/src/constants/dbProviderConstants.ts#L313-L318). The backend then [uses that index directly]() without alangflow/src/lfx/src/lfx/base/knowledge_bases/backends/opensearch.py
Lines 216 to 223 in 3829f4a
kb_namenamespace. Multiple Memory Bases therefore share vectors, allowing cross-MB retrieval and collection-level deletion. Derive a unique index per KB or enforce a KB namespace on every operation. -
[P1] Translate Memory Retrieval filters for OpenSearch.
_build_where_clause()emits{"session_id": ...}, butOpenSearchBackendforwards it unchanged as OpenSearch query DSL. LangChain expects a clause such as{"bool":{"must":[{"term":{"metadata.session_id":"..."}}]}}. Consequently, the default session-filtered retrieval path fails for OpenSearch-backed Memory Bases. -
[P1] Delete remote vectors before dropping their configuration row.
[MemoryBaseService.delete()](https://github.com/langflow-ai/langflow/blob/3829f4ac0f3df21cef7169f88afd60e1647c9b10/src/backend/base/langflow/services/memory_base/service.py#L299-L331)deletes the Memory Base and authoritative KB row, then only removes local disk storage. OpenSearch indexes and Chroma Cloud collections remain indefinitely, and the configuration needed to clean them up is gone. Resolve the backend and calldelete_collection()before removing the rows. -
[P1] Use the database Memory Base marker in bulk deletion. The PR explicitly [stops writing the sidecar](
), but the bulk endpoint still [checks sidecar metadata](). A direct bulk-delete request can therefore remove a Memory Base’s KB row/storage while leaving the Memory Base record dangling. Run the DB-backedlangflow/src/backend/base/langflow/api/v1/knowledge_bases.py
Lines 2206 to 2238 in 3829f4a
_assert_kb_not_memory_baseguard before path resolution and orphan cleanup.
Collaborator
# 1. provision the vector-store collection
await initialize_kb(kb_name=kb_name, ...)
# 2. INSERT the knowledge_base row (its own session)
await _create_kb_record_for_memory_base(kb_name=kb_name, ...)
# 3. uniqueness check + INSERT memory_base row (separate session_scope)
existing = ... # ValueError if name taken
mb = MemoryBase(...); db.add(mb); await db.commit() # IntegrityError possibleAssumption: steps 1–2 only run for a create that will succeed.
Fix options: Given the DB row is now the source of truth, an orphaned row is worse than the old disk-only sidecar — worth closing before merge. |
Collaborator
try:
await backend.ensure_ready()
_ = backend.vector_store
except Exception as exc: # noqa: BLE001 — provisioning is best-effort
await logger.awarning("Initial %s setup for %s failed: %s", backend_type, kb_name, exc)Assumption: "every backend creates the collection lazily on first write anyway, so a provisioning failure must not block creation." That holds for local Chroma. Fix: when the selected backend is non-Chroma, surface an |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Unify Knowledge Bases and Memory Bases on shared vector-store backends
Summary
Memory Bases were hardwired to local Chroma — a directory plus an on-disk
sidecar file. Because their backend and embedding config could only be recovered
by reading that local directory, a replica that never touched it silently fell
back to local Chroma: a remote-backed Memory Base was indistinguishable from a
local one, and writes/queries could diverge to different stores with no error.
This PR makes every Memory Base a first-class Knowledge Base backed by a
knowledge_baserow that is the single source of truth for its backend,embedding, and cached stats. That one change unlocks the pluggable backends KBs
already support (Chroma local/cloud, OpenSearch, MongoDB, Postgres, Astra) for
Memory Bases, and makes all resolution DB-driven and replica-safe.
What changed
resolve_backend_selection/resolve_embedding_selectionread theknowledge_baserow first and treat theon-disk sidecar as a legacy fallback only — so resolution no longer depends on
local disk.
MemoryBaseCreategainsbackend_type/backend_config(defaultchroma, empty config). These are persisted on thebacking
knowledge_baserow, not on thememory_basetable — so there isno schema migration. Deleting a Memory Base now also removes that row.
(
create_backend) instead of opening Chroma directly. The write helper is nowwrite_documents_to_backendand metadata sync issync_kb_stats_to_record, soa Memory Base on OpenSearch/Chroma Cloud ingests to that store from any replica.
backfill_memory_base_rowsreconciles pre-existingMemory Bases that lack a row. It is sourced from the
memory_basetable (nodisk access), touches only orphans via a single
NOT EXISTSquery, and neverraises — so it's cheap after the first run and safe on a replica.
(
DBProviderInput), required before submit, with matching i18n strings andpayload types.
OpenSearchBackend.teardown()now closes theasync client LangChain's
OpenSearchVectorSearchcreates internally. Previouslythe reference was dropped without closing it, leaking the aiohttp
connector/socket and emitting
Unclosed connector/unclosed transportResourceWarnings after ingestion.Compatibility & risk
backend_type="chroma", empty config) keep existing callers on localChroma — no behavior change unless a different backend is selected.
knowledge_baserow.and skipped).
Testing
Unit tests added/updated across Memory Base CRUD, the ingestion task, memory
retrieval, knowledge components, connector endpoints, and the Create Memory modal
hook.
make formatis clean (ruff + biome).E2E testing evidence:
The following data was written for MemoryBases into Chroma Cloud and a hosted OpenSearch Cluster for testing using testing queries.
Summary by CodeRabbit
New Features
Bug Fixes