fix: Knowledge base component refactor

src/backend/base/langflow/components/data/__init__.py

@@ -3,8 +3,6 @@
 from .directory import DirectoryComponent
 from .file import FileComponent
 from .json_to_data import JSONToDataComponent
-from .kb_ingest import KBIngestionComponent
-from .kb_retrieval import KBRetrievalComponent
 from .news_search import NewsSearchComponent
 from .rss import RSSReaderComponent
 from .sql_executor import SQLComponent
@@ -18,8 +16,6 @@
     "DirectoryComponent",
     "FileComponent",
     "JSONToDataComponent",
-    "KBIngestionComponent",
-    "KBRetrievalComponent",
     "NewsSearchComponent",
     "RSSReaderComponent",
     "SQLComponent",

src/backend/base/langflow/components/knowledge_bases/__init__.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from langflow.components._importing import import_mod
+
+if TYPE_CHECKING:
+    from langflow.components.knowledge_bases.ingestion import KnowledgeIngestionComponent
+    from langflow.components.knowledge_bases.retrieval import KnowledgeRetrievalComponent
+
+_dynamic_imports = {
+    "KnowledgeIngestionComponent": "ingestion",
+    "KnowledgeRetrievalComponent": "retrieval",
+}
+
+__all__ = ["KnowledgeIngestionComponent", "KnowledgeRetrievalComponent"]
+
+
+def __getattr__(attr_name: str) -> Any:
+    """Lazily import input/output components on attribute access."""
+    if attr_name not in _dynamic_imports:
+        msg = f"module '{__name__}' has no attribute '{attr_name}'"
+        raise AttributeError(msg)
+    try:
+        result = import_mod(attr_name, _dynamic_imports[attr_name], __spec__.parent)
+    except (ModuleNotFoundError, ImportError, AttributeError) as e:
+        msg = f"Could not import '{attr_name}' from '{__name__}': {e}"
+        raise AttributeError(msg) from e
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)

src/backend/base/langflow/components/data/kb_ingest.py → src/backend/base/langflow/components/knowledge_bases/ingestion.py

@@ -16,10 +16,10 @@
 from langchain_chroma import Chroma
 from loguru import logger
 
-from langflow.base.data.kb_utils import get_knowledge_bases
+from langflow.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases
 from langflow.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
 from langflow.custom import Component
-from langflow.io import BoolInput, DataFrameInput, DropdownInput, IntInput, Output, SecretStrInput, StrInput, TableInput
+from langflow.io import BoolInput, DropdownInput, HandleInput, IntInput, Output, SecretStrInput, StrInput, TableInput
 from langflow.schema.data import Data
 from langflow.schema.dotdict import dotdict  # noqa: TC001
 from langflow.schema.table import EditMode
@@ -38,14 +38,14 @@
 KNOWLEDGE_BASES_ROOT_PATH = Path(knowledge_directory).expanduser()
 
 
-class KBIngestionComponent(Component):
+class KnowledgeIngestionComponent(Component):
     """Create or append to Langflow Knowledge from a DataFrame."""
 
     # ------ UI metadata ---------------------------------------------------
     display_name = "Knowledge Ingestion"
     description = "Create or update knowledge in Langflow."
-    icon = "database"
-    name = "KBIngestion"
+    icon = "upload"
+    name = "KnowledgeIngestion"
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
@@ -101,12 +101,18 @@ class NewKnowledgeBaseInput:
             required=True,
             options=[],
             refresh_button=True,
+            real_time_refresh=True,
             dialog_inputs=asdict(NewKnowledgeBaseInput()),
         ),
-        DataFrameInput(
+        HandleInput(
             name="input_df",
-            display_name="Data",
-            info="Table with all original columns (already chunked / processed).",
+            display_name="Input",
+            info=(
+                "Table with all original columns (already chunked / processed). "
+                "Accepts Data or DataFrame. If Data is provided, it is converted to a DataFrame automatically."
+            ),
+            input_types=["Data", "DataFrame"],
+            is_list=True,
             required=True,
         ),
         TableInput(
@@ -504,7 +510,13 @@ async def build_kb_info(self) -> Data:
         """Main ingestion routine → returns a dict with KB metadata."""
         try:
             # Get source DataFrame
-            df_source: pd.DataFrame = self.input_df
+            if isinstance(self.input_df, Data):
+                df_source: pd.DataFrame = self.input_df.to_dataframe()
+            elif isinstance(self.input_df, list) and all(isinstance(item, Data) for item in self.input_df):
+                # If input_df is a list of Data objects, concatenate them into a single DataFrame
+                df_source: pd.DataFrame = pd.concat([item.to_dataframe() for item in self.input_df], ignore_index=True)
+            else:
+                df_source: pd.DataFrame = self.input_df
 
             # Validate column configuration (using Structured Output patterns)
             config_list = self._validate_column_config(df_source)
@@ -559,9 +571,8 @@ async def build_kb_info(self) -> Data:
             return Data(data=meta)
 
         except (OSError, ValueError, RuntimeError, KeyError) as e:
-            self.log(f"Error in KB ingestion: {e}")
-            self.status = f"❌ KB ingestion failed: {e}"
-            return Data(data={"error": str(e), "kb_name": self.knowledge_base})
+            msg = f"Error during KB ingestion: {e}"
+            raise RuntimeError(msg) from e
 
     async def _get_api_key_variable(self, field_value: dict[str, Any]):
         async with session_scope() as db:

src/backend/base/langflow/components/data/kb_retrieval.py → src/backend/base/langflow/components/knowledge_bases/retrieval.py

@@ -7,7 +7,7 @@
 from loguru import logger
 from pydantic import SecretStr
 
-from langflow.base.data.kb_utils import get_knowledge_bases
+from langflow.base.knowledge_bases.knowledge_base_utils import get_knowledge_bases
 from langflow.custom import Component
 from langflow.io import BoolInput, DropdownInput, IntInput, MessageTextInput, Output, SecretStrInput
 from langflow.schema.data import Data
@@ -24,11 +24,11 @@
 KNOWLEDGE_BASES_ROOT_PATH = Path(knowledge_directory).expanduser()
 
 
-class KBRetrievalComponent(Component):
+class KnowledgeRetrievalComponent(Component):
     display_name = "Knowledge Retrieval"
     description = "Search and retrieve data from knowledge."
-    icon = "database"
-    name = "KBRetrieval"
+    icon = "download"
+    name = "KnowledgeRetrieval"
 
     inputs = [
         DropdownInput(
@@ -51,6 +51,7 @@ class KBRetrievalComponent(Component):
             name="search_query",
             display_name="Search Query",
             info="Optional search query to filter knowledge base data.",
+            tool_mode=True,
         ),
         IntInput(
             name="top_k",
@@ -63,17 +64,24 @@ class KBRetrievalComponent(Component):
         BoolInput(
             name="include_metadata",
             display_name="Include Metadata",
-            info="Whether to include all metadata and embeddings in the output. If false, only content is returned.",
+            info="Whether to include all metadata in the output. If false, only content is returned.",
             value=True,
             advanced=False,
         ),
+        BoolInput(
+            name="include_embeddings",
+            display_name="Include Embeddings",
+            info="Whether to include embeddings in the output. Only applicable if 'Include Metadata' is enabled.",
+            value=False,
+            advanced=True,
+        ),
     ]
 
     outputs = [
         Output(
-            name="chroma_kb_data",
+            name="retrieve_data",
             display_name="Results",
-            method="get_chroma_kb_data",
+            method="retrieve_data",
             info="Returns the data from the selected knowledge base.",
         ),
     ]
@@ -162,7 +170,7 @@ def _build_embeddings(self, metadata: dict):
         msg = f"Embedding provider '{provider}' is not supported for retrieval."
         raise NotImplementedError(msg)
 
-    async def get_chroma_kb_data(self) -> DataFrame:
+    async def retrieve_data(self) -> DataFrame:
         """Retrieve data from the selected knowledge base by reading the Chroma collection.
 
         Returns:
@@ -212,16 +220,16 @@ async def get_chroma_kb_data(self) -> DataFrame:
             # For each result, make it a tuple to match the expected output format
             results = [(doc, 0) for doc in results]  # Assign a dummy score of 0
 
-        # If metadata is enabled, get embeddings for the results
+        # If include_embeddings is enabled, get embeddings for the results
         id_to_embedding = {}
-        if self.include_metadata and results:
+        if self.include_embeddings and results:
             doc_ids = [doc[0].metadata.get("_id") for doc in results if doc[0].metadata.get("_id")]
 
             # Only proceed if we have valid document IDs
             if doc_ids:
                 # Access underlying client to get embeddings
                 collection = chroma._client.get_collection(name=self.knowledge_base)
-                embeddings_result = collection.get(where={"_id": {"$in": doc_ids}}, include=["embeddings", "metadatas"])
+                embeddings_result = collection.get(where={"_id": {"$in": doc_ids}}, include=["metadatas", "embeddings"])
 
                 # Create a mapping from document ID to embedding
                 for i, metadata in enumerate(embeddings_result.get("metadatas", [])):
@@ -231,20 +239,16 @@ async def get_chroma_kb_data(self) -> DataFrame:
         # Build output data based on include_metadata setting
         data_list = []
         for doc in results:
+            kwargs = {
+                "content": doc[0].page_content,
+            }
+            if self.search_query:
+                kwargs["_score"] = -1 * doc[1]
             if self.include_metadata:
                 # Include all metadata, embeddings, and content
-                kwargs = {
-                    "content": doc[0].page_content,
-                    **doc[0].metadata,
-                }
-                if self.search_query:
-                    kwargs["_score"] = -1 * doc[1]
+                kwargs.update(doc[0].metadata)
+            if self.include_embeddings:
                 kwargs["_embeddings"] = id_to_embedding.get(doc[0].metadata.get("_id"))
-            else:
-                # Only include content
-                kwargs = {
-                    "content": doc[0].page_content,
-                }
 
             data_list.append(Data(**kwargs))