databao-context-engine 0.1.4.dev1__py3-none-any.whl → 0.1.6__py3-none-any.whl

databao_context_engine/__init__.py

@@ -8,6 +8,13 @@ from databao_context_engine.datasources.check_config import (
 from databao_context_engine.datasources.datasource_context import DatasourceContext
 from databao_context_engine.datasources.types import Datasource, DatasourceId
 from databao_context_engine.init_project import init_dce_project, init_or_get_dce_project
+from databao_context_engine.llm import (
+    OllamaError,
+    OllamaPermanentError,
+    OllamaTransientError,
+    download_ollama_models_if_needed,
+    install_ollama_if_needed,
+)
 from databao_context_engine.plugin_loader import DatabaoContextPluginLoader
 from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
@@ -16,7 +23,7 @@ from databao_context_engine.pluginlib.build_plugin import (
     DatasourceType,
 )
 from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
-from databao_context_engine.project.info import DceInfo, get_databao_context_engine_info
+from databao_context_engine.project.info import DceInfo, DceProjectInfo, get_databao_context_engine_info
 from databao_context_engine.project.init_project import InitErrorReason, InitProjectError
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
@@ -35,6 +42,7 @@ __all__ = [
     "DatasourceType",
     "get_databao_context_engine_info",
     "DceInfo",
+    "DceProjectInfo",
     "init_dce_project",
     "init_or_get_dce_project",
     "InitErrorReason",
@@ -44,4 +52,9 @@ __all__ = [
     "BuildPlugin",
     "BuildDatasourcePlugin",
     "BuildFilePlugin",
+    "install_ollama_if_needed",
+    "download_ollama_models_if_needed",
+    "OllamaError",
+    "OllamaTransientError",
+    "OllamaPermanentError",
 ]

databao_context_engine/build_sources/build_runner.py

@@ -13,7 +13,7 @@ from databao_context_engine.datasources.datasource_discovery import discover_dat
 from databao_context_engine.datasources.types import DatasourceId
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
 from databao_context_engine.plugins.plugin_loader import load_plugins
-from databao_context_engine.project.layout import get_output_dir
+from databao_context_engine.project.layout import ProjectLayout
 
 logger = logging.getLogger(__name__)
 
@@ -36,7 +36,7 @@ class BuildContextResult:
 
 
 def build(
-    project_dir: Path,
+    project_layout: ProjectLayout,
     *,
     build_service: BuildService,
 ) -> list[BuildContextResult]:
@@ -55,15 +55,15 @@ def build(
     """
     plugins = load_plugins()
 
-    datasources = discover_datasources(project_dir)
+    datasources = discover_datasources(project_layout)
 
     if not datasources:
-        logger.info("No sources discovered under %s", project_dir)
+        logger.info("No sources discovered under %s", project_layout.src_dir)
         return []
 
     number_of_failed_builds = 0
     build_result = []
-    reset_all_results(get_output_dir(project_dir))
+    reset_all_results(project_layout.output_dir)
     for discovered_datasource in datasources:
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -87,14 +87,14 @@ def build(
                 plugin=plugin,
             )
 
-            output_dir = get_output_dir(project_dir)
+            output_dir = project_layout.output_dir
 
             context_file_path = export_build_result(output_dir, result)
             append_result_to_all_results(output_dir, result)
 
             build_result.append(
                 BuildContextResult(
-                    datasource_id=DatasourceId.from_string_repr(result.datasource_id),
+                    datasource_id=discovered_datasource.datasource_id,
                     datasource_type=DatasourceType(full_type=result.datasource_type),
                     context_built_at=result.context_built_at,
                     context_file_path=context_file_path,

databao_context_engine/build_sources/build_wiring.py

@@ -1,5 +1,4 @@
 import logging
-from pathlib import Path
 
 from duckdb import DuckDBPyConnection
 
@@ -12,7 +11,7 @@ from databao_context_engine.llm.factory import (
     create_ollama_embedding_provider,
     create_ollama_service,
 )
-from databao_context_engine.project.layout import ensure_project_dir
+from databao_context_engine.project.layout import ProjectLayout
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 from databao_context_engine.services.factories import create_chunk_embedding_service
 from databao_context_engine.storage.connection import open_duckdb_connection
@@ -22,7 +21,9 @@ from databao_context_engine.system.properties import get_db_path
 logger = logging.getLogger(__name__)
 
 
-def build_all_datasources(project_dir: Path, chunk_embedding_mode: ChunkEmbeddingMode) -> list[BuildContextResult]:
+def build_all_datasources(
+    project_layout: ProjectLayout, chunk_embedding_mode: ChunkEmbeddingMode
+) -> list[BuildContextResult]:
     """Build the context for all datasources in the project.
 
     - Instantiates the build service
@@ -31,21 +32,18 @@ def build_all_datasources(project_dir: Path, chunk_embedding_mode: ChunkEmbeddin
     Returns:
         A list of all the contexts built.
     """
-    ensure_project_dir(project_dir)
-
-    logger.debug(f"Starting to build datasources in project {project_dir.resolve()}")
+    logger.debug(f"Starting to build datasources in project {project_layout.project_dir.resolve()}")
 
     # Think about alternative solutions. This solution will mirror the current behaviour
     # The current behaviour only builds what is currently in the /src folder
     # This will need to change in the future when we can pick which datasources to build
-    db_path = get_db_path(project_dir)
+    db_path = get_db_path(project_layout.project_dir)
     db_path.parent.mkdir(parents=True, exist_ok=True)
     if db_path.exists():
         db_path.unlink()
 
+    migrate(db_path)
     with open_duckdb_connection(db_path) as conn:
-        migrate(db_path)
-
         ollama_service = create_ollama_service()
         embedding_provider = create_ollama_embedding_provider(ollama_service)
         description_provider = (
@@ -60,7 +58,7 @@ def build_all_datasources(project_dir: Path, chunk_embedding_mode: ChunkEmbeddin
             chunk_embedding_mode=chunk_embedding_mode,
         )
         return build(
-            project_dir=project_dir,
+            project_layout=project_layout,
             build_service=build_service,
         )
 

databao_context_engine/build_sources/plugin_execution.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, cast
 
-from databao_context_engine.datasources.types import DatasourceId, PreparedConfig, PreparedDatasource
+from databao_context_engine.datasources.types import PreparedConfig, PreparedDatasource
 from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     BuildFilePlugin,
@@ -39,11 +39,8 @@ class BuiltDatasourceContext:
 
 def execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> BuiltDatasourceContext:
     built_context = _execute(prepared_datasource, plugin)
-
-    datasource_id = DatasourceId.from_datasource_config_file_path(prepared_datasource.path)
-
     return BuiltDatasourceContext(
-        datasource_id=str(datasource_id),
+        datasource_id=str(prepared_datasource.datasource_id),
         datasource_type=prepared_datasource.datasource_type.full_type,
         context_built_at=datetime.now(),
         context=built_context,
@@ -61,10 +58,10 @@ def _execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> An
             config=prepared_datasource.config,
             datasource_name=prepared_datasource.datasource_name,
         )
-    else:
-        file_plugin = cast(BuildFilePlugin, plugin)
-        return execute_file_plugin(
-            plugin=file_plugin,
-            datasource_type=prepared_datasource.datasource_type,
-            file_path=prepared_datasource.path,
-        )
+
+    file_plugin = cast(BuildFilePlugin, plugin)
+    return execute_file_plugin(
+        plugin=file_plugin,
+        datasource_type=prepared_datasource.datasource_type,
+        file_path=prepared_datasource.path,
+    )

databao_context_engine/cli/add_datasource_config.py

@@ -1,5 +1,4 @@
 import os
-from collections import defaultdict
 from pathlib import Path
 from typing import Any
 
@@ -28,12 +27,10 @@ def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
     )
     datasource_name = click.prompt("Datasource name?", type=str)
 
-    is_datasource_existing = project_manager.datasource_config_exists(
-        datasource_type=datasource_type, datasource_name=datasource_name
-    )
-    if is_datasource_existing:
+    datasource_id = project_manager.datasource_config_exists(datasource_name=datasource_name)
+    if datasource_id is not None:
         click.confirm(
-            f"A config file already exists for this datasource ({datasource_type.config_folder}/{datasource_name}). Do you want to overwrite it?",
+            f"A config file already exists for this datasource {datasource_id.relative_path_to_config_file()}. Do you want to overwrite it?",
             abort=True,
             default=False,
         )
@@ -52,25 +49,15 @@ def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
 
 
 def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType:
-    supported_types_by_folder = _group_supported_types_by_folder(supported_datasource_types)
-
-    all_config_folders = sorted(supported_types_by_folder.keys())
-    config_folder = click.prompt(
-        "What type of datasource do you want to add?",
-        type=click.Choice(all_config_folders),
-        default=all_config_folders[0] if len(all_config_folders) == 1 else None,
-    )
-    click.echo(f"Selected type: {config_folder}")
-
-    all_subtypes_for_folder = sorted(supported_types_by_folder[config_folder])
+    all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types])
     config_type = click.prompt(
-        "What is the subtype of this datasource?",
-        type=click.Choice(all_subtypes_for_folder),
-        default=all_subtypes_for_folder[0] if len(all_subtypes_for_folder) == 1 else None,
+        "What type of datasource do you want to add?",
+        type=click.Choice(all_datasource_types),
+        default=all_datasource_types[0] if len(all_datasource_types) == 1 else None,
     )
-    click.echo(f"Selected subtype: {config_type}")
+    click.echo(f"Selected type: {config_type}")
 
-    return DatasourceType.from_main_and_subtypes(config_folder, config_type)
+    return DatasourceType(full_type=config_type)
 
 
 def _ask_for_config_details(config_file_structure: list[ConfigPropertyDefinition]) -> dict[str, Any]:
@@ -146,11 +133,3 @@ def _build_config_content_from_properties(
                 config_content[config_file_property.property_key] = property_value
 
     return config_content
-
-
-def _group_supported_types_by_folder(all_plugin_types: set[DatasourceType]) -> dict[str, list[str]]:
-    main_to_subtypes = defaultdict(list)
-    for datasource_type in all_plugin_types:
-        main_to_subtypes[datasource_type.main_type].append(datasource_type.subtype)
-
-    return main_to_subtypes

databao_context_engine/cli/commands.py

@@ -1,3 +1,4 @@
+import os
 import sys
 from pathlib import Path
 from typing import Literal
@@ -13,11 +14,11 @@ from databao_context_engine import (
     InitErrorReason,
     InitProjectError,
     init_dce_project,
+    install_ollama_if_needed,
 )
 from databao_context_engine.cli.datasources import add_datasource_config_cli, check_datasource_connection_cli
 from databao_context_engine.cli.info import echo_info
 from databao_context_engine.config.logging import configure_logging
-from databao_context_engine.llm.install import resolve_ollama_bin
 from databao_context_engine.mcp.mcp_runner import McpTransport, run_mcp_server
 
 
@@ -33,7 +34,7 @@ from databao_context_engine.mcp.mcp_runner import McpTransport, run_mcp_server
 @click.pass_context
 def dce(ctx: Context, verbose: bool, quiet: bool, project_dir: str | None) -> None:
     if verbose and quiet:
-        print("Arguments --quiet and --verbose can not be used together", file=sys.stderr)
+        click.echo("Arguments --quiet and --verbose can not be used together", file=sys.stderr)
         exit(1)
 
     if project_dir is None:
@@ -79,7 +80,7 @@ def init(ctx: Context) -> None:
     click.echo(f"Project initialized successfully at {project_dir.resolve()}")
 
     try:
-        resolve_ollama_bin()
+        install_ollama_if_needed()
     except RuntimeError as e:
         click.echo(str(e), err=True)
 
@@ -171,29 +172,44 @@ def build(
 )
 @click.option(
     "-o",
-    "--output-format",
-    type=click.Choice(["file", "streamed"]),
-    default="file",
-    help="The output format [file (default), streamed]",
+    "--output-file",
+    type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
+    help="If provided, the results are written as YAML into the file given as a path, rather than printed to the console.",
+)
+@click.option(
+    "-d",
+    "--datasource-id",
+    "datasource_ids_as_str",
+    type=click.STRING,
+    multiple=True,
+    help="A datasourceId to restrict the search to. If not provided, search across all datasources. This option can be provided multiple times",
 )
 @click.pass_context
 def retrieve(
     ctx: Context,
     retrieve_text: tuple[str, ...],
     limit: int | None,
-    output_format: Literal["file", "streamed"],
+    output_file: str | None,
+    datasource_ids_as_str: tuple[str, ...] | None,
 ) -> None:
     """Search the project's built context for the most relevant chunks."""
     text = " ".join(retrieve_text)
 
+    datasource_ids = (
+        [DatasourceId.from_string_repr(datasource_id) for datasource_id in datasource_ids_as_str]
+        if datasource_ids_as_str is not None
+        else None
+    )
+
     databao_engine = DatabaoContextEngine(project_dir=ctx.obj["project_dir"])
 
-    retrieve_results = databao_engine.search_context(
-        retrieve_text=text, limit=limit, export_to_file=output_format == "file"
-    )
+    retrieve_results = databao_engine.search_context(retrieve_text=text, limit=limit, datasource_ids=datasource_ids)
 
-    if output_format == "streamed":
-        display_texts = [context_search_result.context_result for context_search_result in retrieve_results]
+    display_texts = [context_search_result.context_result for context_search_result in retrieve_results]
+    if output_file is not None:
+        Path(output_file).expanduser().write_text(f"---{os.linesep}".join(display_texts))
+        click.echo(f"Found {len(retrieve_results)} results, written to {output_file}")
+    else:
         click.echo("\n".join(display_texts))
 
 

databao_context_engine/databao_context_engine.py

@@ -88,17 +88,14 @@ class DatabaoContextEngine:
     def get_all_contexts_formatted(self) -> str:
         all_contexts = self.get_all_contexts()
 
-        all_results = os.linesep.join(
+        return os.linesep.join(
             [f"{get_context_header_for_datasource(context.datasource_id)}{context.context}" for context in all_contexts]
         )
 
-        return all_results
-
     def search_context(
         self,
         retrieve_text: str,
         limit: int | None = None,
-        export_to_file: bool = False,
         datasource_ids: list[DatasourceId] | None = None,
     ) -> list[ContextSearchResult]:
         """Search in the avaialable context for the closest matches to the given text.
@@ -106,20 +103,13 @@ class DatabaoContextEngine:
         Args:
             retrieve_text: The text to search for in the contexts.
             limit: The maximum number of results to return. If None is provided, a default limit of 10 will be used.
-            export_to_file: Whether the results should be exported to a file as a side-effect. If True, the results will be exported in a file in the run directory.
-            datasource_ids: Not Implemented yet: providing this argument changes nothing to the search
+            datasource_ids: If provided, the search results will only come from the datasources with these IDs.
 
         Returns:
             A list of the results found for the search, sorted by distance.
         """
-        # TODO: Filter with datasource_ids
-        # TODO: When no run_name is required, we can extract the "export_to_file" side-effect and let the caller (the CLI) do it themselves
-
         results = retrieve_embeddings(
-            project_layout=self._project_layout,
-            retrieve_text=retrieve_text,
-            limit=limit,
-            export_to_file=export_to_file,
+            project_layout=self._project_layout, retrieve_text=retrieve_text, limit=limit, datasource_ids=datasource_ids
         )
 
         return [

databao_context_engine/databao_context_project_manager.py

@@ -4,10 +4,6 @@ from typing import Any, overload
 
 from databao_context_engine.build_sources import BuildContextResult, build_all_datasources
 from databao_context_engine.databao_context_engine import DatabaoContextEngine
-from databao_context_engine.datasources.add_config import (
-    create_datasource_config_file,
-    get_datasource_id_for_config_file,
-)
 from databao_context_engine.datasources.check_config import (
     CheckDatasourceConnectionResult,
 )
@@ -17,7 +13,14 @@ from databao_context_engine.datasources.check_config import (
 from databao_context_engine.datasources.datasource_discovery import get_datasource_list
 from databao_context_engine.datasources.types import Datasource, DatasourceId
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
-from databao_context_engine.project.layout import ensure_datasource_config_file_doesnt_exist, ensure_project_dir
+from databao_context_engine.project.layout import (
+    ProjectLayout,
+    ensure_project_dir,
+)
+from databao_context_engine.project.layout import (
+    create_datasource_config_file as create_datasource_config_file_internal,
+)
+from databao_context_engine.serialization.yaml import to_yaml_string
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
 
@@ -45,6 +48,7 @@ class DatabaoContextProjectManager:
     """
 
     project_dir: Path
+    _project_layout: ProjectLayout
 
     def __init__(self, project_dir: Path) -> None:
         """Initialize the DatabaoContextProjectManager.
@@ -52,7 +56,7 @@ class DatabaoContextProjectManager:
         Args:
             project_dir: The root directory of the Databao Context Project.
         """
-        ensure_project_dir(project_dir=project_dir)
+        self._project_layout = ensure_project_dir(project_dir=project_dir)
         self.project_dir = project_dir
 
     def get_configured_datasource_list(self) -> list[Datasource]:
@@ -64,7 +68,7 @@ class DatabaoContextProjectManager:
         Returns:
             The list of datasources configured in the project.
         """
-        return get_datasource_list(self.project_dir)
+        return get_datasource_list(self._project_layout)
 
     def build_context(
         self,
@@ -83,7 +87,7 @@ class DatabaoContextProjectManager:
             The list of all built results.
         """
         # TODO: Filter which datasources to build by datasource_ids
-        return build_all_datasources(project_dir=self.project_dir, chunk_embedding_mode=chunk_embedding_mode)
+        return build_all_datasources(project_layout=self._project_layout, chunk_embedding_mode=chunk_embedding_mode)
 
     def check_datasource_connection(
         self, datasource_ids: list[DatasourceId] | None = None
@@ -97,7 +101,9 @@ class DatabaoContextProjectManager:
             The list of all connection check results, sorted by datasource id.
         """
         return sorted(
-            check_datasource_connection_internal(project_dir=self.project_dir, datasource_ids=datasource_ids).values(),
+            check_datasource_connection_internal(
+                project_layout=self._project_layout, datasource_ids=datasource_ids
+            ).values(),
             key=lambda result: str(result.datasource_id),
         )
 
@@ -121,8 +127,8 @@ class DatabaoContextProjectManager:
             The path to the created datasource configuration file.
         """
         # TODO: Before creating the datasource, validate the config content based on which plugin will be used
-        config_file_path = create_datasource_config_file(
-            project_dir=self.project_dir,
+        config_file = _create_datasource_config_file(
+            project_layout=self._project_layout,
             datasource_type=datasource_type,
             datasource_name=datasource_name,
             config_content=config_content,
@@ -130,48 +136,51 @@ class DatabaoContextProjectManager:
         )
 
         return DatasourceConfigFile(
-            datasource_id=DatasourceId.from_datasource_config_file_path(config_file_path),
-            config_file_path=config_file_path,
+            datasource_id=DatasourceId.from_datasource_config_file_path(self._project_layout, config_file),
+            config_file_path=config_file,
         )
 
     @overload
-    def datasource_config_exists(self, *, datasource_type: DatasourceType, datasource_name: str) -> bool: ...
+    def datasource_config_exists(self, *, datasource_name: str) -> DatasourceId | None: ...
     @overload
-    def datasource_config_exists(self, *, datasource_id: DatasourceId) -> bool: ...
+    def datasource_config_exists(self, *, datasource_id: DatasourceId) -> DatasourceId | None: ...
 
     def datasource_config_exists(
         self,
         *,
-        datasource_type: DatasourceType | None = None,
         datasource_name: str | None = None,
         datasource_id: DatasourceId | None = None,
-    ) -> bool:
+    ) -> DatasourceId | None:
         """Check if a datasource configuration file already exists in the project.
 
         Args:
-            datasource_type: The type of the datasource. This must be used in conjunction with datasource_name.
-            datasource_name: The name of the datasource. This must be used in conjunction with datasource_type.
+            datasource_name: The name of the datasource.
             datasource_id: The id of the datasource. If provided, datasource_type and datasource_name will be ignored.
 
         Returns:
-            True if there is already a datasource configuration file for this datasource, False otherwise.
+            datasource_id if there is already a datasource configuration file for this datasource, None otherwise.
 
         Raises:
             ValueError: If the wrong set of arguments is provided.
         """
-        if datasource_type is not None and datasource_name is not None:
-            datasource_id = get_datasource_id_for_config_file(datasource_type, datasource_name)
-        elif datasource_id is None:
+        if datasource_name is not None:
+            datasource_id = DatasourceId.from_string_repr(f"{datasource_name}.yaml")
+            config_file = self._project_layout.src_dir / datasource_id.relative_path_to_config_file()
+            if config_file.is_file():
+                return datasource_id
+            return None
+
+        if datasource_id is None:
             raise ValueError("Either datasource_id or both datasource_type and datasource_name must be provided")
 
         try:
-            ensure_datasource_config_file_doesnt_exist(
-                project_dir=self.project_dir,
-                datasource_id=datasource_id,
-            )
-            return False
+            config_file = self._project_layout.src_dir.joinpath(datasource_id.relative_path_to_config_file())
+            if config_file.is_file():
+                raise ValueError(f"A config file already exists for {str(datasource_id)}")
+
+            return datasource_id
         except ValueError:
-            return True
+            return None
 
     def get_engine_for_project(self) -> DatabaoContextEngine:
         """Instantiate a DatabaoContextEngine for the project.
@@ -180,3 +189,21 @@ class DatabaoContextProjectManager:
             A DatabaoContextEngine instance for the project.
         """
         return DatabaoContextEngine(project_dir=self.project_dir)
+
+
+def _create_datasource_config_file(
+    project_layout: ProjectLayout,
+    datasource_type: DatasourceType,
+    datasource_name: str,
+    config_content: dict[str, Any],
+    overwrite_existing: bool,
+) -> Path:
+    last_datasource_name = datasource_name.split("/")[-1]
+    basic_config = {"type": datasource_type.full_type, "name": last_datasource_name}
+
+    return create_datasource_config_file_internal(
+        project_layout,
+        f"{datasource_name}.yaml",
+        to_yaml_string(basic_config | config_content),
+        overwrite_existing=overwrite_existing,
+    )

databao_context_engine/datasources/check_config.py

@@ -2,7 +2,6 @@ import logging
 import os
 from dataclasses import dataclass
 from enum import Enum
-from pathlib import Path
 
 from pydantic import ValidationError
 
@@ -15,7 +14,7 @@ from databao_context_engine.datasources.types import DatasourceId, PreparedConfi
 from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, NotSupportedError
 from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource
 from databao_context_engine.plugins.plugin_loader import load_plugins
-from databao_context_engine.project.layout import ensure_project_dir
+from databao_context_engine.project.layout import ProjectLayout
 
 logger = logging.getLogger(__name__)
 
@@ -55,21 +54,19 @@ class CheckDatasourceConnectionResult:
 
 
 def check_datasource_connection(
-    project_dir: Path, *, datasource_ids: list[DatasourceId] | None = None
+    project_layout: ProjectLayout, *, datasource_ids: list[DatasourceId] | None = None
 ) -> dict[DatasourceId, CheckDatasourceConnectionResult]:
-    ensure_project_dir(project_dir)
-
     if datasource_ids:
         logger.info(f"Validating datasource(s): {datasource_ids}")
-        datasources_to_traverse = get_datasource_descriptors(project_dir, datasource_ids)
+        datasources_to_traverse = get_datasource_descriptors(project_layout, datasource_ids)
     else:
-        datasources_to_traverse = discover_datasources(project_dir)
+        datasources_to_traverse = discover_datasources(project_layout)
 
     plugins = load_plugins(exclude_file_plugins=True)
 
     result = {}
     for discovered_datasource in datasources_to_traverse:
-        result_key = DatasourceId.from_datasource_config_file_path(discovered_datasource.path)
+        result_key = DatasourceId.from_datasource_config_file_path(project_layout, discovered_datasource.path)
 
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -127,16 +124,16 @@ def _get_validation_result_from_error(datasource_id: DatasourceId, e: Exception)
             summary="Config file is invalid",
             full_message=str(e),
         )
-    elif isinstance(e, NotImplementedError | NotSupportedError):
+    if isinstance(e, NotImplementedError | NotSupportedError):
         return CheckDatasourceConnectionResult(
             datasource_id=datasource_id,
             connection_status=DatasourceConnectionStatus.UNKNOWN,
             summary="Plugin doesn't support validating its config",
         )
-    else:
-        return CheckDatasourceConnectionResult(
-            datasource_id=datasource_id,
-            connection_status=DatasourceConnectionStatus.INVALID,
-            summary="Connection with the datasource can not be established",
-            full_message=str(e),
-        )
+
+    return CheckDatasourceConnectionResult(
+        datasource_id=datasource_id,
+        connection_status=DatasourceConnectionStatus.INVALID,
+        summary="Connection with the datasource can not be established",
+        full_message=str(e),
+    )