feat: Make knowledge bases user-stored and support global vars (#9458)

* feat: Make knowledge bases user-stored

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Fix ruff error

* [autofix.ci] apply automated fixes

* Reuse code

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Don't show options by default

* [autofix.ci] apply automated fixes

* Pass in the Langflow API key if set

* [autofix.ci] apply automated fixes

* Update files.py

* [autofix.ci] apply automated fixes

* Properly handle secret retrieval

* [autofix.ci] apply automated fixes

* Update src/backend/base/langflow/base/data/kb_utils.py

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

* Update src/backend/base/langflow/base/data/kb_utils.py

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

* Update src/backend/base/langflow/components/data/kb_ingest.py

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Feedback from review

* [autofix.ci] apply automated fixes

* Fix other uses of incorrect user

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Feedback from review 2

* [autofix.ci] apply automated fixes

* Update kb_ingest.py

* [autofix.ci] apply automated fixes

* Update tests

* [autofix.ci] apply automated fixes

* Update kb_ingest.py

* [autofix.ci] apply automated fixes

* Fix mypy issues

* [autofix.ci] apply automated fixes

* Update kb_utils.py

* Update test_kb_ingest.py

* Fix tests

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

src/backend/base/langflow/api/v1/knowledge_bases.py

@@ -9,6 +9,7 @@ from langchain_chroma import Chroma
 from loguru import logger
 from pydantic import BaseModel
 
+from langflow.api.utils import CurrentActiveUser
 from langflow.services.deps import get_settings_service
 
 router = APIRouter(tags=["Knowledge Bases"], prefix="/knowledge_bases")
@@ -290,17 +291,19 @@ def get_kb_metadata(kb_path: Path) -> dict:
 
 @router.get("", status_code=HTTPStatus.OK)
 @router.get("/", status_code=HTTPStatus.OK)
-async def list_knowledge_bases() -> list[KnowledgeBaseInfo]:
+async def list_knowledge_bases(current_user: CurrentActiveUser) -> list[KnowledgeBaseInfo]:
     """List all available knowledge bases."""
     try:
         kb_root_path = get_kb_root_path()
+        kb_user = current_user.username
+        kb_path = kb_root_path / kb_user
 
-        if not kb_root_path.exists():
+        if not kb_path.exists():
             return []
 
         knowledge_bases = []
 
-        for kb_dir in kb_root_path.iterdir():
+        for kb_dir in kb_path.iterdir():
             if not kb_dir.is_dir() or kb_dir.name.startswith("."):
                 continue
 
@@ -340,11 +343,12 @@ async def list_knowledge_bases() -> list[KnowledgeBaseInfo]:
 
 
 @router.get("/{kb_name}", status_code=HTTPStatus.OK)
-async def get_knowledge_base(kb_name: str) -> KnowledgeBaseInfo:
+async def get_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> KnowledgeBaseInfo:
     """Get detailed information about a specific knowledge base."""
     try:
         kb_root_path = get_kb_root_path()
-        kb_path = kb_root_path / kb_name
+        kb_user = current_user.username
+        kb_path = kb_root_path / kb_user / kb_name
 
         if not kb_path.exists() or not kb_path.is_dir():
             raise HTTPException(status_code=404, detail=f"Knowledge base '{kb_name}' not found")
@@ -374,11 +378,12 @@ async def get_knowledge_base(kb_name: str) -> KnowledgeBaseInfo:
 
 
 @router.delete("/{kb_name}", status_code=HTTPStatus.OK)
-async def delete_knowledge_base(kb_name: str) -> dict[str, str]:
+async def delete_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> dict[str, str]:
     """Delete a specific knowledge base."""
     try:
         kb_root_path = get_kb_root_path()
-        kb_path = kb_root_path / kb_name
+        kb_user = current_user.username
+        kb_path = kb_root_path / kb_user / kb_name
 
         if not kb_path.exists() or not kb_path.is_dir():
             raise HTTPException(status_code=404, detail=f"Knowledge base '{kb_name}' not found")
@@ -396,15 +401,17 @@ async def delete_knowledge_base(kb_name: str) -> dict[str, str]:
 
 @router.delete("", status_code=HTTPStatus.OK)
 @router.delete("/", status_code=HTTPStatus.OK)
-async def delete_knowledge_bases_bulk(request: BulkDeleteRequest) -> dict[str, object]:
+async def delete_knowledge_bases_bulk(request: BulkDeleteRequest, current_user: CurrentActiveUser) -> dict[str, object]:
     """Delete multiple knowledge bases."""
     try:
         kb_root_path = get_kb_root_path()
+        kb_user = current_user.username
+        kb_user_path = kb_root_path / kb_user
         deleted_count = 0
         not_found_kbs = []
 
         for kb_name in request.kb_names:
-            kb_path = kb_root_path / kb_name
+            kb_path = kb_user_path / kb_name
 
             if not kb_path.exists() or not kb_path.is_dir():
                 not_found_kbs.append(kb_name)

src/backend/base/langflow/api/v2/files.py

@@ -123,7 +123,9 @@ async def upload_user_file(
             unique_filename = new_filename
         else:
             # For normal files, ensure unique name by appending a count if necessary
-            stmt = select(UserFile).where(col(UserFile.name).like(f"{root_filename}%"))
+            stmt = select(UserFile).where(
+                col(UserFile.name).like(f"{root_filename}%"), UserFile.user_id == current_user.id
+            )
             existing_files = await session.exec(stmt)
             files = existing_files.all()  # Fetch all matching records
 

src/backend/base/langflow/base/data/kb_utils.py

@@ -1,5 +1,10 @@
 import math
 from collections import Counter
+from pathlib import Path
+from uuid import UUID
+
+from langflow.services.database.models.user.crud import get_user_by_id
+from langflow.services.deps import session_scope
 
 
 def compute_tfidf(documents: list[str], query_terms: list[str]) -> list[float]:
@@ -102,3 +107,31 @@ def compute_bm25(documents: list[str], query_terms: list[str], k1: float = 1.2,
         scores.append(doc_score)
 
     return scores
+
+
+async def get_knowledge_bases(kb_root: Path, user_id: UUID | str) -> list[str]:
+    """Retrieve a list of available knowledge bases.
+
+    Returns:
+        A list of knowledge base names.
+    """
+    if not kb_root.exists():
+        return []
+
+    # Get the current user
+    async with session_scope() as db:
+        if not user_id:
+            msg = "User ID is required for fetching knowledge bases."
+            raise ValueError(msg)
+        user_id = UUID(user_id) if isinstance(user_id, str) else user_id
+        current_user = await get_user_by_id(db, user_id)
+        if not current_user:
+            msg = f"User with ID {user_id} not found."
+            raise ValueError(msg)
+        kb_user = current_user.username
+    kb_path = kb_root / kb_user
+
+    if not kb_path.exists():
+        return []
+
+    return [str(d.name) for d in kb_path.iterdir() if not d.name.startswith(".") and d.is_dir()]

src/backend/base/langflow/components/agents/mcp_component.py

@@ -16,16 +16,15 @@ from langflow.base.mcp.util import (
 )
 from langflow.custom.custom_component.component_with_cache import ComponentWithCache
 from langflow.inputs.inputs import InputTypes  # noqa: TC001
-from langflow.io import DropdownInput, McpInput, MessageTextInput, Output, SecretStrInput
+from langflow.io import DropdownInput, McpInput, MessageTextInput, Output
 from langflow.io.schema import flatten_schema, schema_to_langflow_inputs
 from langflow.logging import logger
 from langflow.schema.dataframe import DataFrame
 from langflow.schema.message import Message
 
 # Import get_server from the backend API
-from langflow.services.auth.utils import create_user_longterm_token, get_current_user
 from langflow.services.database.models.user.crud import get_user_by_id
-from langflow.services.deps import get_session, get_settings_service, get_storage_service
+from langflow.services.deps import get_settings_service, get_storage_service, session_scope
 
 
 class MCPToolsComponent(ComponentWithCache):
@@ -96,13 +95,6 @@ class MCPToolsComponent(ComponentWithCache):
             show=False,
             tool_mode=False,
         ),
-        SecretStrInput(
-            name="api_key",
-            display_name="Langflow API Key",
-            info="Langflow API key for authentication when fetching MCP servers and tools.",
-            required=False,
-            advanced=True,
-        ),
     ]
 
     outputs = [
@@ -161,19 +153,11 @@ class MCPToolsComponent(ComponentWithCache):
             return self.tools, {"name": server_name, "config": server_config_from_value}
 
         try:
-            async for db in get_session():
+            async with session_scope() as db:
-                # TODO: In 1.6, this may need to be removed or adjusted
+                if not self.user_id:
-                # Try to get the super user token, if possible
+                    msg = "User ID is required for fetching MCP tools."
-                if self.api_key:
+                    raise ValueError(msg)
-                    current_user = await get_current_user(
+                current_user = await get_user_by_id(db, self.user_id)
-                        token=None,
-                        query_param=self.api_key,
-                        header_param=None,
-                        db=db,
-                    )
-                else:
-                    user_id, _ = await create_user_longterm_token(db)
-                    current_user = await get_user_by_id(db, user_id)
 
                 # Try to get server config from DB/API
                 server_config = await get_server(
@@ -184,39 +168,38 @@ class MCPToolsComponent(ComponentWithCache):
                     settings_service=get_settings_service(),
                 )
 
-                # If get_server returns empty but we have a config, use it
+            # If get_server returns empty but we have a config, use it
-                if not server_config and server_config_from_value:
+            if not server_config and server_config_from_value:
-                    server_config = server_config_from_value
+                server_config = server_config_from_value
 
-                if not server_config:
+            if not server_config:
-                    self.tools = []
+                self.tools = []
-                    return [], {"name": server_name, "config": server_config}
+                return [], {"name": server_name, "config": server_config}
 
-                _, tool_list, tool_cache = await update_tools(
+            _, tool_list, tool_cache = await update_tools(
-                    server_name=server_name,
+                server_name=server_name,
-                    server_config=server_config,
+                server_config=server_config,
-                    mcp_stdio_client=self.stdio_client,
+                mcp_stdio_client=self.stdio_client,
-                    mcp_sse_client=self.sse_client,
+                mcp_sse_client=self.sse_client,
-                )
+            )
 
-                self.tool_names = [tool.name for tool in tool_list if hasattr(tool, "name")]
+            self.tool_names = [tool.name for tool in tool_list if hasattr(tool, "name")]
-                self._tool_cache = tool_cache
+            self._tool_cache = tool_cache
-                self.tools = tool_list
+            self.tools = tool_list
-                # Cache the result using shared cache
+            # Cache the result using shared cache
-                cache_data = {
+            cache_data = {
-                    "tools": tool_list,
+                "tools": tool_list,
-                    "tool_names": self.tool_names,
+                "tool_names": self.tool_names,
-                    "tool_cache": tool_cache,
+                "tool_cache": tool_cache,
-                    "config": server_config,
+                "config": server_config,
-                }
+            }
 
-                # Safely update the servers cache
+            # Safely update the servers cache
-                current_servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
+            current_servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
-                if isinstance(current_servers_cache, dict):
+            if isinstance(current_servers_cache, dict):
-                    current_servers_cache[server_name] = cache_data
+                current_servers_cache[server_name] = cache_data
-                    safe_cache_set(self._shared_component_cache, "servers", current_servers_cache)
+                safe_cache_set(self._shared_component_cache, "servers", current_servers_cache)
 
-                return tool_list, {"name": server_name, "config": server_config}
         except (TimeoutError, asyncio.TimeoutError) as e:
             msg = f"Timeout updating tool list: {e!s}"
             logger.exception(msg)
@@ -225,6 +208,8 @@ class MCPToolsComponent(ComponentWithCache):
             msg = f"Error updating tool list: {e!s}"
             logger.exception(msg)
             raise ValueError(msg) from e
+        else:
+            return tool_list, {"name": server_name, "config": server_config}
 
     async def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None) -> dict:
         """Toggle the visibility of connection-specific fields based on the selected mode."""

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

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+import contextlib
 import hashlib
 import json
 import re
@@ -14,6 +16,7 @@ from cryptography.fernet import InvalidToken
 from langchain_chroma import Chroma
 from loguru import logger
 
+from langflow.base.data.kb_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
@@ -21,7 +24,8 @@ from langflow.schema.data import Data
 from langflow.schema.dotdict import dotdict  # noqa: TC001
 from langflow.schema.table import EditMode
 from langflow.services.auth.utils import decrypt_api_key, encrypt_api_key
-from langflow.services.deps import get_settings_service
+from langflow.services.database.models.user.crud import get_user_by_id
+from langflow.services.deps import get_settings_service, get_variable_service, session_scope
 
 HUGGINGFACE_MODEL_NAMES = ["sentence-transformers/all-MiniLM-L6-v2", "sentence-transformers/all-mpnet-base-v2"]
 COHERE_MODEL_NAMES = ["embed-english-v3.0", "embed-multilingual-v3.0"]
@@ -43,6 +47,10 @@ class KBIngestionComponent(Component):
     icon = "database"
     name = "KBIngestion"
 
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._cached_kb_path: Path | None = None
+
     @dataclass
     class NewKnowledgeBaseInput:
         functionality: str = "create"
@@ -76,7 +84,7 @@ class KBIngestionComponent(Component):
                                 display_name="API Key",
                                 info="Provider API key for embedding model",
                                 required=True,
-                                load_from_db=True,
+                                load_from_db=False,
                             ),
                         },
                     },
@@ -91,11 +99,7 @@ class KBIngestionComponent(Component):
             display_name="Knowledge",
             info="Select the knowledge to load data from.",
             required=True,
-            options=[
+            options=[],
-                str(d.name) for d in KNOWLEDGE_BASES_ROOT_PATH.iterdir() if not d.name.startswith(".") and d.is_dir()
-            ]
-            if KNOWLEDGE_BASES_ROOT_PATH.exists()
-            else [],
             refresh_button=True,
             dialog_inputs=asdict(NewKnowledgeBaseInput()),
         ),
@@ -329,22 +333,23 @@ class KBIngestionComponent(Component):
 
         return metadata
 
-    def _create_vector_store(
+    async def _create_vector_store(
         self, df_source: pd.DataFrame, config_list: list[dict[str, Any]], embedding_model: str, api_key: str
     ) -> None:
         """Create vector store following Local DB component pattern."""
         try:
             # Set up vector store directory
-            base_dir = self._get_kb_root()
+            vector_store_dir = await self._kb_path()
-
+            if not vector_store_dir:
-            vector_store_dir = base_dir / self.knowledge_base
+                msg = "Knowledge base path is not set. Please create a new knowledge base first."
+                raise ValueError(msg)
             vector_store_dir.mkdir(parents=True, exist_ok=True)
 
             # Create embeddings model
             embedding_function = self._build_embeddings(embedding_model, api_key)
 
             # Convert DataFrame to Data objects (following Local DB pattern)
-            data_objects = self._convert_df_to_data_objects(df_source, config_list)
+            data_objects = await self._convert_df_to_data_objects(df_source, config_list)
 
             # Create vector store
             chroma = Chroma(
@@ -367,16 +372,18 @@ class KBIngestionComponent(Component):
         except (OSError, ValueError, RuntimeError) as e:
             self.log(f"Error creating vector store: {e}")
 
-    def _convert_df_to_data_objects(self, df_source: pd.DataFrame, config_list: list[dict[str, Any]]) -> list[Data]:
+    async def _convert_df_to_data_objects(
+        self, df_source: pd.DataFrame, config_list: list[dict[str, Any]]
+    ) -> list[Data]:
         """Convert DataFrame to Data objects for vector store."""
         data_objects: list[Data] = []
 
         # Set up vector store directory
-        base_dir = self._get_kb_root()
+        kb_path = await self._kb_path()
 
         # If we don't allow duplicates, we need to get the existing hashes
         chroma = Chroma(
-            persist_directory=str(base_dir / self.knowledge_base),
+            persist_directory=str(kb_path),
             collection_name=self.knowledge_base,
         )
 
@@ -466,10 +473,34 @@ class KBIngestionComponent(Component):
         # Check allowed characters (condition 3)
         return re.match(r"^[a-zA-Z0-9_-]+$", name) is not None
 
+    async def _kb_path(self) -> Path | None:
+        # Check if we already have the path cached
+        cached_path = getattr(self, "_cached_kb_path", None)
+        if cached_path is not None:
+            return cached_path
+
+        # If not cached, compute it
+        async with session_scope() as db:
+            if not self.user_id:
+                msg = "User ID is required for fetching knowledge base path."
+                raise ValueError(msg)
+            current_user = await get_user_by_id(db, self.user_id)
+            if not current_user:
+                msg = f"User with ID {self.user_id} not found."
+                raise ValueError(msg)
+            kb_user = current_user.username
+
+        kb_root = self._get_kb_root()
+
+        # Cache the result
+        self._cached_kb_path = kb_root / kb_user / self.knowledge_base
+
+        return self._cached_kb_path
+
     # ---------------------------------------------------------------------
     #                         OUTPUT METHODS
     # ---------------------------------------------------------------------
-    def build_kb_info(self) -> Data:
+    async def build_kb_info(self) -> Data:
         """Main ingestion routine → returns a dict with KB metadata."""
         try:
             # Get source DataFrame
@@ -479,11 +510,11 @@ class KBIngestionComponent(Component):
             config_list = self._validate_column_config(df_source)
             column_metadata = self._build_column_metadata(config_list, df_source)
 
-            # Prepare KB folder (using File Component patterns)
-            kb_root = self._get_kb_root()
-            kb_path = kb_root / self.knowledge_base
-
             # Read the embedding info from the knowledge base folder
+            kb_path = await self._kb_path()
+            if not kb_path:
+                msg = "Knowledge base path is not set. Please create a new knowledge base first."
+                raise ValueError(msg)
             metadata_path = kb_path / "embedding_metadata.json"
 
             # If the API key is not provided, try to read it from the metadata file
@@ -506,7 +537,7 @@ class KBIngestionComponent(Component):
                 )
 
             # Create vector store following Local DB component pattern
-            self._create_vector_store(df_source, config_list, embedding_model=embedding_model, api_key=api_key)
+            await self._create_vector_store(df_source, config_list, embedding_model=embedding_model, api_key=api_key)
 
             # Save KB files (using File Component storage patterns)
             self._save_kb_files(kb_path, config_list)
@@ -532,40 +563,77 @@ class KBIngestionComponent(Component):
             self.status = f"❌ KB ingestion failed: {e}"
             return Data(data={"error": str(e), "kb_name": self.knowledge_base})
 
-    def _get_knowledge_bases(self) -> list[str]:
+    async def _get_api_key_variable(self, field_value: dict[str, Any]):
-        """Retrieve a list of available knowledge bases.
+        async with session_scope() as db:
+            if not self.user_id:
+                msg = "User ID is required for fetching global variables."
+                raise ValueError(msg)
+            current_user = await get_user_by_id(db, self.user_id)
+            if not current_user:
+                msg = f"User with ID {self.user_id} not found."
+                raise ValueError(msg)
+            variable_service = get_variable_service()
 
-        Returns:
+            # Process the api_key field variable
-            A list of knowledge base names.
+            return await variable_service.get_variable(
-        """
+                user_id=current_user.id,
-        # Return the list of directories in the knowledge base root path
+                name=field_value["03_api_key"],
-        kb_root_path = self._get_kb_root()
+                field="",
+                session=db,
+            )
 
-        if not kb_root_path.exists():
+    async def update_build_config(
-            return []
+        self,
-
+        build_config: dotdict,
-        return [str(d.name) for d in kb_root_path.iterdir() if not d.name.startswith(".") and d.is_dir()]
+        field_value: Any,
-
+        field_name: str | None = None,
-    def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None) -> dotdict:
+    ) -> dotdict:
         """Update build configuration based on provider selection."""
         # Create a new knowledge base
         if field_name == "knowledge_base":
+            async with session_scope() as db:
+                if not self.user_id:
+                    msg = "User ID is required for fetching knowledge base list."
+                    raise ValueError(msg)
+                current_user = await get_user_by_id(db, self.user_id)
+                if not current_user:
+                    msg = f"User with ID {self.user_id} not found."
+                    raise ValueError(msg)
+                kb_user = current_user.username
             if isinstance(field_value, dict) and "01_new_kb_name" in field_value:
                 # Validate the knowledge base name - Make sure it follows these rules:
                 if not self.is_valid_collection_name(field_value["01_new_kb_name"]):
                     msg = f"Invalid knowledge base name: {field_value['01_new_kb_name']}"
                     raise ValueError(msg)
 
-                # We need to test the API Key one time against the embedding model
+                api_key = field_value.get("03_api_key", None)
-                embed_model = self._build_embeddings(
+                with contextlib.suppress(Exception):
-                    embedding_model=field_value["02_embedding_model"], api_key=field_value["03_api_key"]
+                    # If the API key is a variable, resolve it
-                )
+                    api_key = await self._get_api_key_variable(field_value)
 
-                # Try to generate a dummy embedding to validate the API key
+                # Make sure api_key is a string
-                embed_model.embed_query("test")
+                if not isinstance(api_key, str):
+                    msg = "API key must be a string."
+                    raise ValueError(msg)
+
+                # We need to test the API Key one time against the embedding model
+                embed_model = self._build_embeddings(embedding_model=field_value["02_embedding_model"], api_key=api_key)
+
+                # Try to generate a dummy embedding to validate the API key without blocking the event loop
+                try:
+                    await asyncio.wait_for(
+                        asyncio.to_thread(embed_model.embed_query, "test"),
+                        timeout=10,
+                    )
+                except TimeoutError as e:
+                    msg = "Embedding validation timed out. Please verify network connectivity and key."
+                    raise ValueError(msg) from e
+                except Exception as e:
+                    msg = f"Embedding validation failed: {e!s}"
+                    raise ValueError(msg) from e
 
                 # Create the new knowledge base directory
-                kb_path = KNOWLEDGE_BASES_ROOT_PATH / field_value["01_new_kb_name"]
+                kb_path = KNOWLEDGE_BASES_ROOT_PATH / kb_user / field_value["01_new_kb_name"]
                 kb_path.mkdir(parents=True, exist_ok=True)
 
                 # Save the embedding metadata
@@ -573,11 +641,16 @@ class KBIngestionComponent(Component):
                 self._save_embedding_metadata(
                     kb_path=kb_path,
                     embedding_model=field_value["02_embedding_model"],
-                    api_key=field_value["03_api_key"],
+                    api_key=api_key,
                 )
 
             # Update the knowledge base options dynamically
-            build_config["knowledge_base"]["options"] = self._get_knowledge_bases()
+            build_config["knowledge_base"]["options"] = await get_knowledge_bases(
+                KNOWLEDGE_BASES_ROOT_PATH,
+                user_id=self.user_id,
+            )
+
+            # If the selected knowledge base is not available, reset it
             if build_config["knowledge_base"]["value"] not in build_config["knowledge_base"]["options"]:
                 build_config["knowledge_base"]["value"] = None
 

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

@@ -5,13 +5,16 @@ from typing import Any
 from cryptography.fernet import InvalidToken
 from langchain_chroma import Chroma
 from loguru import logger
+from pydantic import SecretStr
 
+from langflow.base.data.kb_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
 from langflow.schema.dataframe import DataFrame
 from langflow.services.auth.utils import decrypt_api_key
-from langflow.services.deps import get_settings_service
+from langflow.services.database.models.user.crud import get_user_by_id
+from langflow.services.deps import get_settings_service, session_scope
 
 settings = get_settings_service().settings
 knowledge_directory = settings.knowledge_bases_dir
@@ -33,11 +36,7 @@ class KBRetrievalComponent(Component):
             display_name="Knowledge",
             info="Select the knowledge to load data from.",
             required=True,
-            options=[
+            options=[],
-                str(d.name) for d in KNOWLEDGE_BASES_ROOT_PATH.iterdir() if not d.name.startswith(".") and d.is_dir()
-            ]
-            if KNOWLEDGE_BASES_ROOT_PATH.exists()
-            else [],
             refresh_button=True,
             real_time_refresh=True,
         ),
@@ -79,21 +78,13 @@ class KBRetrievalComponent(Component):
         ),
     ]
 
-    def _get_knowledge_bases(self) -> list[str]:
+    async def update_build_config(self, build_config, field_value, field_name=None):  # noqa: ARG002
-        """Retrieve a list of available knowledge bases.
-
-        Returns:
-            A list of knowledge base names.
-        """
-        if not KNOWLEDGE_BASES_ROOT_PATH.exists():
-            return []
-
-        return [str(d.name) for d in KNOWLEDGE_BASES_ROOT_PATH.iterdir() if not d.name.startswith(".") and d.is_dir()]
-
-    def update_build_config(self, build_config, field_value, field_name=None):  # noqa: ARG002
         if field_name == "knowledge_base":
             # Update the knowledge base options dynamically
-            build_config["knowledge_base"]["options"] = self._get_knowledge_bases()
+            build_config["knowledge_base"]["options"] = await get_knowledge_bases(
+                KNOWLEDGE_BASES_ROOT_PATH,
+                user_id=self.user_id,  # Use the user_id from the component context
+            )
 
             # If the selected knowledge base is not available, reset it
             if build_config["knowledge_base"]["value"] not in build_config["knowledge_base"]["options"]:
@@ -129,15 +120,12 @@ class KBRetrievalComponent(Component):
 
     def _build_embeddings(self, metadata: dict):
         """Build embedding model from metadata."""
+        runtime_api_key = self.api_key.get_secret_value() if isinstance(self.api_key, SecretStr) else self.api_key
         provider = metadata.get("embedding_provider")
         model = metadata.get("embedding_model")
-        api_key = metadata.get("api_key")
+        api_key = runtime_api_key or metadata.get("api_key")
         chunk_size = metadata.get("chunk_size")
 
-        # If user provided a key in the input, it overrides the stored one.
-        if self.api_key and self.api_key.get_secret_value():
-            api_key = self.api_key.get_secret_value()
-
         # Handle various providers
         if provider == "OpenAI":
             from langchain_openai import OpenAIEmbeddings
@@ -174,13 +162,23 @@ class KBRetrievalComponent(Component):
         msg = f"Embedding provider '{provider}' is not supported for retrieval."
         raise NotImplementedError(msg)
 
-    def get_chroma_kb_data(self) -> DataFrame:
+    async def get_chroma_kb_data(self) -> DataFrame:
         """Retrieve data from the selected knowledge base by reading the Chroma collection.
 
         Returns:
             A DataFrame containing the data rows from the knowledge base.
         """
-        kb_path = KNOWLEDGE_BASES_ROOT_PATH / self.knowledge_base
+        # Get the current user
+        async with session_scope() as db:
+            if not self.user_id:
+                msg = "User ID is required for fetching Knowledge Base data."
+                raise ValueError(msg)
+            current_user = await get_user_by_id(db, self.user_id)
+            if not current_user:
+                msg = f"User with ID {self.user_id} not found."
+                raise ValueError(msg)
+            kb_user = current_user.username
+        kb_path = KNOWLEDGE_BASES_ROOT_PATH / kb_user / self.knowledge_base
 
         metadata = self._get_kb_metadata(kb_path)
         if not metadata:

src/backend/base/langflow/components/processing/save_file.py

@@ -1,7 +1,6 @@
 import json
 from collections.abc import AsyncIterator, Iterator
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import orjson
 import pandas as pd
@@ -10,16 +9,12 @@ from fastapi.encoders import jsonable_encoder
 
 from langflow.api.v2.files import upload_user_file
 from langflow.custom import Component
-from langflow.io import DropdownInput, HandleInput, SecretStrInput, StrInput
+from langflow.io import DropdownInput, HandleInput, StrInput
 from langflow.schema import Data, DataFrame, Message
-from langflow.services.auth.utils import create_user_longterm_token, get_current_user
 from langflow.services.database.models.user.crud import get_user_by_id
-from langflow.services.deps import get_session, get_settings_service, get_storage_service
+from langflow.services.deps import get_settings_service, get_storage_service, session_scope
 from langflow.template.field.base import Output
 
-if TYPE_CHECKING:
-    from langflow.services.database.models.user.model import User
-
 
 class SaveToFileComponent(Component):
     display_name = "Save File"
@@ -55,13 +50,6 @@ class SaveToFileComponent(Component):
             value="",
             advanced=True,
         ),
-        SecretStrInput(
-            name="api_key",
-            display_name="Langflow API Key",
-            info="Langflow API key for authentication when saving the file.",
-            required=False,
-            advanced=True,
-        ),
     ]
 
     outputs = [Output(display_name="File Path", name="message", method="save_to_file")]
@@ -148,25 +136,11 @@ class SaveToFileComponent(Component):
             raise FileNotFoundError(msg)
 
         with file_path.open("rb") as f:
-            async for db in get_session():
+            async with session_scope() as db:
-                # TODO: In 1.6, this may need to be removed or adjusted
+                if not self.user_id:
-                # Try to get the super user token, if possible
+                    msg = "User ID is required for file saving."
-                current_user: User | None = None
-                if self.api_key:
-                    current_user = await get_current_user(
-                        token="",
-                        query_param=self.api_key,
-                        header_param="",
-                        db=db,
-                    )
-                else:
-                    user_id, _ = await create_user_longterm_token(db)
-                    current_user = await get_user_by_id(db, user_id)
-
-                # Fail if the user is not found
-                if not current_user:
-                    msg = "User not found. Please provide a valid API key or ensure the user exists."
                     raise ValueError(msg)
+                current_user = await get_user_by_id(db, self.user_id)
 
                 await upload_user_file(
                     file=UploadFile(filename=file_path.name, file=f, size=file_path.stat().st_size),

src/backend/tests/unit/components/data/test_kb_ingest.py

@@ -1,8 +1,10 @@
 import json
-from unittest.mock import MagicMock, patch
+import uuid
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pandas as pd
 import pytest
+from langflow.base.data.kb_utils import get_knowledge_bases
 from langflow.components.data.kb_ingest import KBIngestionComponent
 from langflow.schema.data import Data
 
@@ -21,8 +23,43 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         with patch("langflow.components.data.kb_ingest.KNOWLEDGE_BASES_ROOT_PATH", tmp_path):
             yield
 
+    class MockUser:
+        def __init__(self, user_id):
+            self.id = user_id
+            self.username = "langflow"
+
     @pytest.fixture
-    def default_kwargs(self, tmp_path):
+    def mock_user_data(self):
+        """Create mock user data that persists for the test function."""
+        mock_uuid = uuid.uuid4()
+        mock_user = self.MockUser(mock_uuid)
+        return {"user_id": mock_uuid, "user": mock_user.username, "user_obj": mock_user}
+
+    @pytest.fixture(autouse=True)
+    def setup_mocks(self, mock_user_data):
+        """Mock the component's user_id attribute and User object."""
+        with (
+            patch.object(KBIngestionComponent, "user_id", mock_user_data["user_id"]),
+            patch(
+                "langflow.components.data.kb_ingest.get_user_by_id",
+                new_callable=AsyncMock,
+                return_value=mock_user_data["user_obj"],
+            ),
+            patch(
+                "langflow.base.data.kb_utils.get_user_by_id",
+                new_callable=AsyncMock,
+                return_value=mock_user_data["user_obj"],
+            ),
+        ):
+            yield
+
+    @pytest.fixture
+    def mock_user_id(self, mock_user_data):
+        """Get the mock user data."""
+        return {"user_id": mock_user_data["user_id"], "user": mock_user_data["user"]}
+
+    @pytest.fixture
+    def default_kwargs(self, tmp_path, mock_user_id):
         """Return default kwargs for component instantiation."""
         # Create a sample DataFrame
         data_df = pd.DataFrame(
@@ -38,8 +75,8 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
 
         # Create knowledge base directory
         kb_name = "test_kb"
-        kb_path = tmp_path / kb_name
+        kb_path = tmp_path / mock_user_id["user"] / kb_name
-        kb_path.mkdir(exist_ok=True)
+        kb_path.mkdir(parents=True, exist_ok=True)
 
         # Create embedding metadata file
         metadata = {
@@ -206,7 +243,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         assert "text" in metadata["summary"]["vectorized_columns"]
         assert "category" in metadata["summary"]["identifier_columns"]
 
-    def test_convert_df_to_data_objects(self, component_class, default_kwargs):
+    async def test_convert_df_to_data_objects(self, component_class, default_kwargs):
         """Test converting DataFrame to Data objects."""
         component = component_class(**default_kwargs)
         data_df = default_kwargs["input_df"]
@@ -218,7 +255,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
             mock_chroma_instance.get.return_value = {"metadatas": []}
             mock_chroma.return_value = mock_chroma_instance
 
-            data_objects = component._convert_df_to_data_objects(data_df, config_list)
+            data_objects = await component._convert_df_to_data_objects(data_df, config_list)
 
         assert len(data_objects) == 2
         assert all(isinstance(obj, Data) for obj in data_objects)
@@ -230,7 +267,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         assert "category" in first_obj.data
         assert "_id" in first_obj.data
 
-    def test_convert_df_to_data_objects_no_duplicates(self, component_class, default_kwargs):
+    async def test_convert_df_to_data_objects_no_duplicates(self, component_class, default_kwargs):
         """Test converting DataFrame to Data objects with duplicate prevention."""
         default_kwargs["allow_duplicates"] = False
         component = component_class(**default_kwargs)
@@ -251,7 +288,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
                 mock_hash_obj.hexdigest.side_effect = [existing_hash, "different_hash"]
                 mock_hash.return_value = mock_hash_obj
 
-                data_objects = component._convert_df_to_data_objects(data_df, config_list)
+                data_objects = await component._convert_df_to_data_objects(data_df, config_list)
 
         # Should only return one object (second row) since first is duplicate
         assert len(data_objects) == 1
@@ -274,7 +311,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
 
     @patch("langflow.components.data.kb_ingest.json.loads")
     @patch("langflow.components.data.kb_ingest.decrypt_api_key")
-    def test_build_kb_info_success(self, mock_decrypt, mock_json_loads, component_class, default_kwargs):
+    async def test_build_kb_info_success(self, mock_decrypt, mock_json_loads, component_class, default_kwargs):
         """Test successful KB info building."""
         component = component_class(**default_kwargs)
 
@@ -287,7 +324,7 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
 
         # Mock vector store creation
         with patch.object(component, "_create_vector_store"), patch.object(component, "_save_kb_files"):
-            result = component.build_kb_info()
+            result = await component.build_kb_info()
 
         assert isinstance(result, Data)
         assert "kb_id" in result.data
@@ -295,32 +332,21 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         assert "rows" in result.data
         assert result.data["rows"] == 2
 
-    def test_get_knowledge_bases(self, component_class, default_kwargs, tmp_path):
+    async def test_get_knowledge_bases(self, tmp_path, mock_user_id):
         """Test getting list of knowledge bases."""
-        component = component_class(**default_kwargs)
-
         # Create additional test directories
-        (tmp_path / "kb1").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb1").mkdir(parents=True, exist_ok=True)
-        (tmp_path / "kb2").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb2").mkdir(parents=True, exist_ok=True)
-        (tmp_path / ".hidden").mkdir()  # Should be ignored
+        (tmp_path / mock_user_id["user"] / ".hidden").mkdir(parents=True, exist_ok=True)  # Should be ignored
 
-        kb_list = component._get_knowledge_bases()
+        kb_list = await get_knowledge_bases(tmp_path, user_id=mock_user_id["user_id"])
 
         assert "test_kb" in kb_list
         assert "kb1" in kb_list
         assert "kb2" in kb_list
         assert ".hidden" not in kb_list
 
-    @patch("langflow.components.data.kb_ingest.Path.exists")
+    async def test_update_build_config_new_kb(self, component_class, default_kwargs):
-    def test_get_knowledge_bases_no_path(self, mock_exists, component_class, default_kwargs):
-        """Test getting knowledge bases when path doesn't exist."""
-        component = component_class(**default_kwargs)
-        mock_exists.return_value = False
-
-        kb_list = component._get_knowledge_bases()
-        assert kb_list == []
-
-    def test_update_build_config_new_kb(self, component_class, default_kwargs):
         """Test updating build config for new knowledge base creation."""
         component = component_class(**default_kwargs)
 
@@ -329,26 +355,24 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         field_value = {
             "01_new_kb_name": "new_test_kb",
             "02_embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
-            "03_api_key": None,
+            "03_api_key": "abc123",  # Mock API key
         }
 
         # Mock embedding validation
         with (
             patch.object(component, "_build_embeddings") as mock_build_emb,
             patch.object(component, "_save_embedding_metadata"),
-            patch.object(component, "_get_knowledge_bases") as mock_get_kbs,
         ):
             mock_embeddings = MagicMock()
             mock_embeddings.embed_query.return_value = [0.1, 0.2, 0.3]
             mock_build_emb.return_value = mock_embeddings
-            mock_get_kbs.return_value = ["new_test_kb"]
 
-            result = component.update_build_config(build_config, field_value, "knowledge_base")
+            result = await component.update_build_config(build_config, field_value, "knowledge_base")
 
         assert result["knowledge_base"]["value"] == "new_test_kb"
         assert "new_test_kb" in result["knowledge_base"]["options"]
 
-    def test_update_build_config_invalid_kb_name(self, component_class, default_kwargs):
+    async def test_update_build_config_invalid_kb_name(self, component_class, default_kwargs):
         """Test updating build config with invalid KB name."""
         component = component_class(**default_kwargs)
 
@@ -360,4 +384,4 @@ class TestKBIngestionComponent(ComponentTestBaseWithoutClient):
         }
 
         with pytest.raises(ValueError, match="Invalid knowledge base name"):
-            component.update_build_config(build_config, field_value, "knowledge_base")
+            await component.update_build_config(build_config, field_value, "knowledge_base")

src/backend/tests/unit/components/data/test_kb_retrieval.py

@@ -1,10 +1,13 @@
 import contextlib
 import json
+import uuid
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from langflow.base.data.kb_utils import get_knowledge_bases
 from langflow.components.data.kb_retrieval import KBRetrievalComponent
+from pydantic import SecretStr
 
 from tests.base import ComponentTestBaseWithoutClient
 
@@ -21,13 +24,48 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         with patch("langflow.components.data.kb_retrieval.KNOWLEDGE_BASES_ROOT_PATH", tmp_path):
             yield
 
+    class MockUser:
+        def __init__(self, user_id):
+            self.id = user_id
+            self.username = "langflow"
+
     @pytest.fixture
-    def default_kwargs(self, tmp_path):
+    def mock_user_data(self):
+        """Create mock user data that persists for the test function."""
+        mock_uuid = uuid.uuid4()
+        mock_user = self.MockUser(mock_uuid)
+        return {"user_id": mock_uuid, "user": mock_user.username, "user_obj": mock_user}
+
+    @pytest.fixture(autouse=True)
+    def setup_mocks(self, mock_user_data):
+        """Mock the component's user_id attribute and User object."""
+        with (
+            patch.object(KBRetrievalComponent, "user_id", mock_user_data["user_id"]),
+            patch(
+                "langflow.components.data.kb_retrieval.get_user_by_id",
+                new_callable=AsyncMock,
+                return_value=mock_user_data["user_obj"],
+            ),
+            patch(
+                "langflow.base.data.kb_utils.get_user_by_id",
+                new_callable=AsyncMock,
+                return_value=mock_user_data["user_obj"],
+            ),
+        ):
+            yield
+
+    @pytest.fixture
+    def mock_user_id(self, mock_user_data):
+        """Get the mock user data."""
+        return {"user_id": mock_user_data["user_id"], "user": mock_user_data["user"]}
+
+    @pytest.fixture
+    def default_kwargs(self, tmp_path, mock_user_id):
         """Return default kwargs for component instantiation."""
         # Create knowledge base directory structure
         kb_name = "test_kb"
-        kb_path = tmp_path / kb_name
+        kb_path = tmp_path / mock_user_id["user"] / kb_name
-        kb_path.mkdir(exist_ok=True)
+        kb_path.mkdir(parents=True, exist_ok=True)
 
         # Create embedding metadata file
         metadata = {
@@ -55,61 +93,50 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         # This is a new component, so it doesn't exist in older versions
         return []
 
-    def test_get_knowledge_bases(self, component_class, default_kwargs, tmp_path):
+    async def test_get_knowledge_bases(self, tmp_path, mock_user_id):
         """Test getting list of knowledge bases."""
-        component = component_class(**default_kwargs)
-
         # Create additional test directories
-        (tmp_path / "kb1").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb1").mkdir(parents=True, exist_ok=True)
-        (tmp_path / "kb2").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb2").mkdir(parents=True, exist_ok=True)
-        (tmp_path / ".hidden").mkdir()  # Should be ignored
+        (tmp_path / mock_user_id["user"] / ".hidden").mkdir(parents=True, exist_ok=True)  # Should be ignored
 
-        kb_list = component._get_knowledge_bases()
+        kb_list = await get_knowledge_bases(tmp_path, user_id=mock_user_id["user_id"])
 
         assert "test_kb" in kb_list
         assert "kb1" in kb_list
         assert "kb2" in kb_list
         assert ".hidden" not in kb_list
 
-    @patch("langflow.components.data.kb_retrieval.Path.exists")
+    async def test_update_build_config(self, component_class, default_kwargs, tmp_path, mock_user_id):
-    def test_get_knowledge_bases_no_path(self, mock_exists, component_class, default_kwargs):
-        """Test getting knowledge bases when path doesn't exist."""
-        component = component_class(**default_kwargs)
-        mock_exists.return_value = False
-
-        kb_list = component._get_knowledge_bases()
-        assert kb_list == []
-
-    def test_update_build_config(self, component_class, default_kwargs, tmp_path):
         """Test updating build configuration."""
         component = component_class(**default_kwargs)
 
         # Create additional KB directories
-        (tmp_path / "kb1").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb1").mkdir(parents=True, exist_ok=True)
-        (tmp_path / "kb2").mkdir()
+        (tmp_path / mock_user_id["user"] / "kb2").mkdir(parents=True, exist_ok=True)
 
         build_config = {"knowledge_base": {"value": "test_kb", "options": []}}
 
-        result = component.update_build_config(build_config, None, "knowledge_base")
+        result = await component.update_build_config(build_config, None, "knowledge_base")
 
         assert "test_kb" in result["knowledge_base"]["options"]
         assert "kb1" in result["knowledge_base"]["options"]
         assert "kb2" in result["knowledge_base"]["options"]
 
-    def test_update_build_config_invalid_kb(self, component_class, default_kwargs):
+    async def test_update_build_config_invalid_kb(self, component_class, default_kwargs):
         """Test updating build config when selected KB is not available."""
         component = component_class(**default_kwargs)
 
         build_config = {"knowledge_base": {"value": "nonexistent_kb", "options": ["test_kb"]}}
 
-        result = component.update_build_config(build_config, None, "knowledge_base")
+        result = await component.update_build_config(build_config, None, "knowledge_base")
 
         assert result["knowledge_base"]["value"] is None
 
-    def test_get_kb_metadata_success(self, component_class, default_kwargs):
+    def test_get_kb_metadata_success(self, component_class, default_kwargs, mock_user_id):
         """Test successful metadata loading."""
         component = component_class(**default_kwargs)
-        kb_path = Path(default_kwargs["kb_root_path"]) / default_kwargs["knowledge_base"]
+        kb_path = Path(default_kwargs["kb_root_path"]) / mock_user_id["user"] / default_kwargs["knowledge_base"]
 
         with patch("langflow.components.data.kb_retrieval.decrypt_api_key") as mock_decrypt:
             mock_decrypt.return_value = "decrypted_key"
@@ -120,21 +147,21 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         assert metadata["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2"
         assert "chunk_size" in metadata
 
-    def test_get_kb_metadata_no_file(self, component_class, default_kwargs, tmp_path):
+    def test_get_kb_metadata_no_file(self, component_class, default_kwargs, tmp_path, mock_user_id):
         """Test metadata loading when file doesn't exist."""
         component = component_class(**default_kwargs)
-        nonexistent_path = tmp_path / "nonexistent"
+        nonexistent_path = tmp_path / mock_user_id["user"] / "nonexistent"
-        nonexistent_path.mkdir()
+        nonexistent_path.mkdir(parents=True, exist_ok=True)
 
         metadata = component._get_kb_metadata(nonexistent_path)
 
         assert metadata == {}
 
-    def test_get_kb_metadata_json_error(self, component_class, default_kwargs, tmp_path):
+    def test_get_kb_metadata_json_error(self, component_class, default_kwargs, tmp_path, mock_user_id):
         """Test metadata loading with invalid JSON."""
         component = component_class(**default_kwargs)
-        kb_path = tmp_path / "invalid_json_kb"
+        kb_path = tmp_path / mock_user_id["user"] / "invalid_json_kb"
-        kb_path.mkdir()
+        kb_path.mkdir(parents=True, exist_ok=True)
 
         # Create invalid JSON file
         (kb_path / "embedding_metadata.json").write_text("invalid json content")
@@ -143,11 +170,11 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
 
         assert metadata == {}
 
-    def test_get_kb_metadata_decrypt_error(self, component_class, default_kwargs, tmp_path):
+    def test_get_kb_metadata_decrypt_error(self, component_class, default_kwargs, tmp_path, mock_user_id):
         """Test metadata loading with decryption error."""
         component = component_class(**default_kwargs)
-        kb_path = tmp_path / "decrypt_error_kb"
+        kb_path = tmp_path / mock_user_id["user"] / "decrypt_error_kb"
-        kb_path.mkdir()
+        kb_path.mkdir(parents=True, exist_ok=True)
 
         # Create metadata with encrypted key
         metadata = {
@@ -274,10 +301,8 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
 
     def test_build_embeddings_with_user_api_key(self, component_class, default_kwargs):
         """Test that user-provided API key overrides stored one."""
-        # Create a mock secret input
+        # Use a real SecretStr object instead of a mock
-
+        mock_secret = SecretStr("user-provided-key")
-        mock_secret = MagicMock()
-        mock_secret.get_secret_value.return_value = "user-provided-key"
 
         default_kwargs["api_key"] = mock_secret
         component = component_class(**default_kwargs)
@@ -285,7 +310,7 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         metadata = {
             "embedding_provider": "OpenAI",
             "embedding_model": "text-embedding-ada-002",
-            "api_key": "stored-key",
+            "api_key": "stored-key",  # This should be overridden by the user-provided key
             "chunk_size": 1000,
         }
 
@@ -295,14 +320,17 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
 
             component._build_embeddings(metadata)
 
+            # The user-provided key should override the stored key in metadata
             mock_openai.assert_called_once_with(
-                model="text-embedding-ada-002", api_key="user-provided-key", chunk_size=1000
+                model="text-embedding-ada-002",
+                api_key="user-provided-key",  # Should use the user-provided key, not "stored-key"
+                chunk_size=1000,
             )
 
-    def test_get_chroma_kb_data_no_metadata(self, component_class, default_kwargs, tmp_path):
+    async def test_get_chroma_kb_data_no_metadata(self, component_class, default_kwargs, tmp_path, mock_user_id):
         """Test retrieving data when metadata is missing."""
         # Remove metadata file
-        kb_path = tmp_path / default_kwargs["knowledge_base"]
+        kb_path = tmp_path / mock_user_id["user"] / default_kwargs["knowledge_base"]
         metadata_file = kb_path / "embedding_metadata.json"
         if metadata_file.exists():
             metadata_file.unlink()
@@ -310,7 +338,7 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         component = component_class(**default_kwargs)
 
         with pytest.raises(ValueError, match="Metadata not found for knowledge base"):
-            component.get_chroma_kb_data()
+            await component.get_chroma_kb_data()
 
     def test_get_chroma_kb_data_path_construction(self, component_class, default_kwargs):
         """Test that get_chroma_kb_data constructs the correct paths."""
@@ -331,7 +359,7 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
         assert hasattr(component, "top_k")
         assert hasattr(component, "include_embeddings")
 
-    def test_get_chroma_kb_data_method_exists(self, component_class, default_kwargs):
+    async def test_get_chroma_kb_data_method_exists(self, component_class, default_kwargs):
         """Test that get_chroma_kb_data method exists and can be called."""
         component = component_class(**default_kwargs)
 
@@ -349,7 +377,7 @@ class TestKBRetrievalComponent(ComponentTestBaseWithoutClient):
 
             # This is a unit test focused on the component's internal logic
             with contextlib.suppress(Exception):
-                component.get_chroma_kb_data()
+                await component.get_chroma_kb_data()
 
             # Verify internal methods were called
             mock_get_metadata.assert_called_once()