Unified KB and MB to work with common vector db backends.

Author: dkaushik94

.gitignore

@@ -268,6 +268,7 @@ src/backend/langflow/frontend/
 src/backend/base/langflow/frontend/
 .docker
 scratchpad*
+scratch/
 chroma*/*
 stuff/*
 src/frontend/playwright-report/index.html

src/backend/base/langflow/api/utils/kb_helpers.py

@@ -100,6 +100,115 @@ def chunk_text_for_ingestion(
     return splitter.split_text(text)
 
 
+def _coerce_backend_config_value(value: Any) -> dict[str, Any]:
+    """Normalize a stored ``backend_config`` into a plain dict."""
+    if isinstance(value, dict):
+        return dict(value)
+    return {}
+
+
+async def resolve_backend_selection(
+    *,
+    user_id: uuid.UUID,
+    kb_name: str,
+    kb_path: Path,
+) -> tuple[str, dict[str, Any]]:
+    """Resolve ``(backend_type, backend_config)`` for a KB — DB row first, sidecar second.
+
+    The ``knowledge_base`` row is authoritative. The on-disk sidecar is a legacy
+    fallback for KBs predating the row (or dropped in from an export), not a
+    default: when neither source resolves this raises rather than assuming local
+    storage. A replica that merely lacks the sidecar would otherwise write a
+    remote-backed KB into a local Chroma directory while queries followed the row
+    to the configured cluster and returned nothing, with no error anywhere.
+
+    Shared by Knowledge Bases and Memory Bases so the two cannot drift apart.
+    """
+    from langflow.api.utils import knowledge_base_service
+
+    record = await knowledge_base_service.get_by_user_and_name(user_id, kb_name)
+    if record is not None:
+        return (
+            record.backend_type or BackendType.CHROMA.value,
+            _coerce_backend_config_value(record.backend_config),
+        )
+
+    # A sidecar with no ``backend_type`` is an explicit legacy local KB; only the
+    # *absence* of both sources is ambiguous.
+    metadata_file = kb_path / "embedding_metadata.json"
+    if await asyncio.to_thread(metadata_file.exists):
+        metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
+        return (
+            str(metadata.get("backend_type") or BackendType.CHROMA.value),
+            _coerce_backend_config_value(metadata.get("backend_config")),
+        )
+
+    msg = (
+        f"Cannot determine the vector-store backend for '{kb_name}': it has no "
+        f"knowledge_base record and no embedding metadata on disk. Refusing to fall "
+        f"back to local storage, which would write to a different store than queries "
+        f"read from."
+    )
+    raise ValueError(msg)
+
+
+# Default embedding used when neither the DB row nor the sidecar records one.
+# Matches the historical fallback in ``resolve_embedding`` so behavior is
+# unchanged for callers that relied on it.
+_DEFAULT_EMBEDDING_PROVIDER = "OpenAI"
+_DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
+
+
+async def resolve_embedding_selection(
+    *,
+    user_id: uuid.UUID,
+    kb_name: str,
+    kb_path: Path | None = None,
+) -> tuple[str, str]:
+    """Resolve ``(embedding_provider, embedding_model)`` for a KB — DB row first, sidecar second.
+
+    Counterpart to :func:`resolve_backend_selection`: the ``knowledge_base`` row's
+    ``model_selection`` is authoritative, the on-disk ``embedding_metadata.json``
+    sidecar is a legacy fallback (KBs predating the row / shared with file-based
+    KBs), and a hardcoded default is the last resort.
+
+    Memory Bases always have a row (created at MB-create and backfilled for
+    pre-existing ones), so they resolve entirely from the DB and never touch the
+    sidecar — which is what lets a Memory Base work on a replica whose local disk
+    never held the KB directory. ``kb_path`` is therefore optional: pass ``None``
+    from a fully DB-driven path and the sidecar branch is skipped entirely (no
+    dependency on ``knowledge_bases_dir`` or the local filesystem). The sidecar
+    branch stays for Knowledge Bases, which still write it.
+
+    Shared by Knowledge Bases and Memory Bases so the two cannot drift apart.
+    """
+    from langflow.api.utils import knowledge_base_service
+    from langflow.api.utils.knowledge_base_service import get_embedding_model, get_embedding_provider
+
+    record = await knowledge_base_service.get_by_user_and_name(user_id, kb_name)
+    if record is not None:
+        model = get_embedding_model(record.model_selection)
+        if model:
+            return get_embedding_provider(record.model_selection), model
+
+    # Legacy fallback: read the sidecar. Only attempted when a ``kb_path`` is
+    # supplied and there is no row (or a row with no model recorded) — i.e. never
+    # for a properly provisioned Memory Base.
+    if kb_path is not None:
+        metadata_file = kb_path / "embedding_metadata.json"
+        if await asyncio.to_thread(metadata_file.exists):
+            metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
+            provider = str(metadata.get("embedding_provider") or "").strip()
+            model = str(metadata.get("embedding_model") or "").strip()
+            if model and model.lower() != "unknown":
+                return (
+                    provider if provider and provider.lower() != "unknown" else _DEFAULT_EMBEDDING_PROVIDER,
+                    model,
+                )
+
+    return _DEFAULT_EMBEDDING_PROVIDER, _DEFAULT_EMBEDDING_MODEL
+
+
 class KBStorageHelper:
     """Helper class for Knowledge Base storage and path management."""
 
@@ -981,62 +1090,6 @@ async def cleanup_chroma_chunks_by_job(
         finally:
             await backend.teardown()
 
-    @staticmethod
-    async def write_documents_to_chroma(
-        *,
-        documents: list[Document],
-        chroma: Chroma,
-        task_job_id: uuid.UUID,
-        job_service: JobService,
-    ) -> int:
-        """Write pre-built Documents into an open Chroma collection.
-
-        This is the shared primitive used by both file-based KB ingestion
-        (``perform_ingestion``) and message-based Memory Base ingestion.
-
-        Documents must already be chunked and have their metadata populated
-        by the caller — this method only handles the batched write, cancellation
-        checking, and retry logic.
-
-        Args:
-            documents: LangChain Document objects ready for embedding.
-            chroma: An already-constructed ``Chroma`` instance pointing at the
-                target collection.
-            task_job_id: Job ID used to poll for cancellation.
-            job_service: Service for checking job status.
-
-        Returns:
-            Number of documents successfully written.  If the job is cancelled
-            mid-batch this will be less than ``len(documents)``.
-
-        Raises:
-            Exception: Re-raises any non-cancellation write failure after the
-                retry budget is exhausted.
-        """
-        written = 0
-        for i in range(0, len(documents), INGESTION_BATCH_SIZE):
-            if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
-                return written
-
-            batch = documents[i : i + INGESTION_BATCH_SIZE]
-            for attempt in range(MAX_RETRY_ATTEMPTS):
-                if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
-                    return written
-                try:
-                    await chroma.aadd_documents(batch)
-                    break
-                except Exception as e:
-                    if attempt == MAX_RETRY_ATTEMPTS - 1:
-                        raise
-                    wait = (attempt + 1) * EXPONENTIAL_BACKOFF_MULTIPLIER
-                    await logger.awarning("Write failed, retrying in %ds: %s", wait, e)
-                    await asyncio.sleep(wait)
-
-            written += len(batch)
-            await asyncio.sleep(0.01)
-
-        return written
-
     @staticmethod
     async def write_documents_to_backend(
         *,
@@ -1047,12 +1100,12 @@ async def write_documents_to_backend(
     ) -> int:
         """Write pre-built Documents through a ``BaseVectorStoreBackend``.
 
-        Backend-agnostic counterpart to :meth:`write_documents_to_chroma`.
-        Used by the multi-backend KB ingestion path so Mongo/Astra/
-        Postgres/OpenSearch ingestions share the same batching,
-        cancellation-checking, and exponential-backoff retry logic that
-        Memory Base's Chroma path gets from
-        :meth:`write_documents_to_chroma`.
+        The single write primitive for every ingestion path — file-based
+        Knowledge Bases and message-based Memory Bases alike — so batching,
+        cancellation checking, and exponential-backoff retry behave identically
+        whatever backend the KB is on. A Chroma-only counterpart used to exist
+        for Memory Bases; it was removed so the local-only path cannot be
+        reintroduced by accident.
 
         Documents must already be chunked with metadata populated.
 

src/backend/base/langflow/api/utils/knowledge_base_service.py

@@ -214,6 +214,67 @@ async def backfill_all_users_from_disk(*, kb_root: Path | None = None) -> int:
     return inserted
 
 
+async def backfill_memory_base_rows() -> int:
+    """Ensure every Memory Base has a backing ``knowledge_base`` row.
+
+    Memory Bases are DB-driven: their backend + embedding config resolve from the
+    ``knowledge_base`` row, never an on-disk sidecar. Memory Bases created before
+    that row existed have only a ``memory_base`` row, so this startup reconcile
+    creates the missing ``knowledge_base`` row from the ``memory_base`` table
+    itself — no disk access, so it works on a replica whose filesystem never held
+    the KB directory.
+
+    Only Memory Bases *lacking* a ``knowledge_base`` row are fetched (a single
+    ``NOT EXISTS`` query), so after the first run this returns zero rows and does
+    almost nothing — startup cost is proportional to the backlog, not the total
+    Memory Base count. ``backend_type="chroma"`` with empty config is the correct
+    reconstruction: a Memory Base with no row predates the backend selector, when
+    every Memory Base was local Chroma, and the ``memory_base`` table records no
+    backend to recover. Memory Bases created through the selector get their real
+    backend from ``_create_kb_record_for_memory_base`` and never reach this path.
+
+    Returns the number of rows inserted. Never raises — per-MB failures are logged
+    and skipped so one bad row doesn't block startup.
+    """
+    from sqlalchemy import exists
+
+    from langflow.services.database.models.memory_base.model import MemoryBase
+    from langflow.services.memory_base.embedding_helpers import infer_embedding_provider
+
+    # Fetch only Memory Bases with no matching knowledge_base row (by owner + name).
+    orphan_stmt = (
+        select(MemoryBase)
+        .where(MemoryBase.kb_name != "")
+        .where(
+            ~exists(
+                select(KnowledgeBaseRecord.id)
+                .where(KnowledgeBaseRecord.user_id == MemoryBase.user_id)
+                .where(KnowledgeBaseRecord.name == MemoryBase.kb_name)
+            )
+        )
+    )
+    async with session_scope() as session:
+        orphans = list((await session.exec(orphan_stmt)).all())
+
+    inserted = 0
+    for mb in orphans:
+        try:
+            provider = infer_embedding_provider(mb.embedding_model)
+            await create_record(
+                user_id=mb.user_id,
+                name=mb.kb_name,
+                model_selection={"name": mb.embedding_model, "provider": provider},
+                backend_type="chroma",
+                backend_config={},
+                source_types=["memory"],
+            )
+            inserted += 1
+        except Exception as exc:  # noqa: BLE001
+            await logger.aerror("memory-base backfill: failed to upsert KB row for %s: %s", mb.kb_name, exc)
+
+    return inserted
+
+
 async def update_stats(
     record_id: UUID,
     *,