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

databao_context_engine/__init__.py

@@ -1,16 +1,23 @@
-from databao_context_engine.build_sources.internal.build_runner import BuildContextResult
+from databao_context_engine.build_sources import BuildContextResult
+from databao_context_engine.databao_context_engine import ContextSearchResult, DatabaoContextEngine
 from databao_context_engine.databao_context_project_manager import DatabaoContextProjectManager, DatasourceConfigFile
-from databao_context_engine.databao_engine import ContextSearchResult, DatabaoContextEngine
-from databao_context_engine.datasource_config.check_config import (
+from databao_context_engine.datasources.check_config import (
     CheckDatasourceConnectionResult,
     DatasourceConnectionStatus,
 )
-from databao_context_engine.datasource_config.datasource_context import DatasourceContext
+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.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.plugin_loader import DatabaoContextPluginLoader
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    BuildFilePlugin,
+    BuildPlugin,
+    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.init_project import InitErrorReason, InitProjectError
-from databao_context_engine.project.types import Datasource, DatasourceId
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
 __all__ = [
@@ -32,4 +39,9 @@ __all__ = [
     "init_or_get_dce_project",
     "InitErrorReason",
     "InitProjectError",
+    "DatabaoContextPluginLoader",
+    "ConfigPropertyDefinition",
+    "BuildPlugin",
+    "BuildDatasourcePlugin",
+    "BuildFilePlugin",
 ]

databao_context_engine/build_sources/__init__.py

@@ -0,0 +1,4 @@
+from databao_context_engine.build_sources.build_runner import BuildContextResult
+from databao_context_engine.build_sources.build_wiring import build_all_datasources
+
+__all__ = ["build_all_datasources", "BuildContextResult"]

databao_context_engine/build_sources/{internal/build_runner.py → build_runner.py}

@@ -3,22 +3,32 @@ from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
 
-from databao_context_engine.build_sources.internal.build_service import BuildService
-from databao_context_engine.build_sources.internal.export_results import (
+from databao_context_engine.build_sources.build_service import BuildService
+from databao_context_engine.build_sources.export_results import (
     append_result_to_all_results,
-    create_run_dir,
     export_build_result,
+    reset_all_results,
 )
+from databao_context_engine.datasources.datasource_discovery import discover_datasources, prepare_source
+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.datasource_discovery import discover_datasources, prepare_source
-from databao_context_engine.project.types import DatasourceId
+from databao_context_engine.project.layout import get_output_dir
 
 logger = logging.getLogger(__name__)
 
 
 @dataclass
 class BuildContextResult:
+    """Result of a single datasource context build.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        datasource_type: The type of the datasource.
+        context_built_at: The timestamp when the context was built.
+        context_file_path: The path to the generated context file.
+    """
+
     datasource_id: DatasourceId
     datasource_type: DatasourceType
     context_built_at: datetime
@@ -29,16 +39,19 @@ def build(
     project_dir: Path,
     *,
     build_service: BuildService,
-    project_id: str,
-    dce_version: str,
 ) -> list[BuildContextResult]:
-    """
-    Build entrypoint.
+    """Build the context for all datasources in the project.
+
+    Unless you already have access to BuildService, this should not be called directly.
+    Instead, internal callers should go through the build_wiring module or directly use DatabaoContextProjectManager.build_context().
 
     1) Load available plugins
     2) Discover sources
     3) Create a run
     4) For each source, call process_source
+
+    Returns:
+        A list of all the contexts built.
     """
     plugins = load_plugins()
 
@@ -48,11 +61,9 @@ def build(
         logger.info("No sources discovered under %s", project_dir)
         return []
 
-    run = None
-    run_dir = None
-
     number_of_failed_builds = 0
     build_result = []
+    reset_all_results(get_output_dir(project_dir))
     for discovered_datasource in datasources:
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -71,19 +82,15 @@ def build(
                 number_of_failed_builds += 1
                 continue
 
-            if run is None or run_dir is None:
-                # Initialiase the run as soon as we found a source
-                run = build_service.start_run(project_id=project_id, dce_version=dce_version)
-                run_dir = create_run_dir(project_dir, run.run_name)
-
             result = build_service.process_prepared_source(
-                run_id=run.run_id,
                 prepared_source=prepared_source,
                 plugin=plugin,
             )
 
-            context_file_path = export_build_result(run_dir, result)
-            append_result_to_all_results(run_dir, result)
+            output_dir = get_output_dir(project_dir)
+
+            context_file_path = export_build_result(output_dir, result)
+            append_result_to_all_results(output_dir, result)
 
             build_result.append(
                 BuildContextResult(
@@ -99,9 +106,6 @@ def build(
 
             number_of_failed_builds += 1
 
-    if run is not None:
-        build_service.finalize_run(run_id=run.run_id)
-
     logger.debug(
         "Successfully built %d datasources. %s",
         len(build_result),

databao_context_engine/build_sources/build_service.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import logging
+
+from databao_context_engine.build_sources.plugin_execution import BuiltDatasourceContext, execute
+from databao_context_engine.datasources.types import PreparedDatasource
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildPlugin,
+)
+from databao_context_engine.serialization.yaml import to_yaml_string
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingService
+
+logger = logging.getLogger(__name__)
+
+
+class BuildService:
+    def __init__(
+        self,
+        *,
+        chunk_embedding_service: ChunkEmbeddingService,
+    ) -> None:
+        self._chunk_embedding_service = chunk_embedding_service
+
+    def process_prepared_source(
+        self,
+        *,
+        prepared_source: PreparedDatasource,
+        plugin: BuildPlugin,
+    ) -> BuiltDatasourceContext:
+        """Process a single source to build its context.
+
+        1) Execute the plugin
+        2) Divide the results into chunks
+        3) Embed and persist the chunks
+
+        Returns:
+            The built context.
+        """
+        result = execute(prepared_source, plugin)
+
+        chunks = plugin.divide_context_into_chunks(result.context)
+        if not chunks:
+            logger.info("No chunks for %s — skipping.", prepared_source.path.name)
+            return result
+
+        self._chunk_embedding_service.embed_chunks(
+            chunks=chunks,
+            result=to_yaml_string(result.context),
+            full_type=prepared_source.datasource_type.full_type,
+            datasource_id=result.datasource_id,
+        )
+
+        return result

databao_context_engine/build_sources/build_wiring.py

@@ -0,0 +1,84 @@
+import logging
+from pathlib import Path
+
+from duckdb import DuckDBPyConnection
+
+from databao_context_engine.build_sources.build_runner import BuildContextResult, build
+from databao_context_engine.build_sources.build_service import BuildService
+from databao_context_engine.llm.descriptions.provider import DescriptionProvider
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.llm.factory import (
+    create_ollama_description_provider,
+    create_ollama_embedding_provider,
+    create_ollama_service,
+)
+from databao_context_engine.project.layout import ensure_project_dir
+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
+from databao_context_engine.storage.migrate import migrate
+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]:
+    """Build the context for all datasources in the project.
+
+    - Instantiates the build service
+    - Delegates the actual build logic to the build runner
+
+    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()}")
+
+    # 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.parent.mkdir(parents=True, exist_ok=True)
+    if db_path.exists():
+        db_path.unlink()
+
+    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 = (
+            create_ollama_description_provider(ollama_service)
+            if chunk_embedding_mode.should_generate_description()
+            else None
+        )
+        build_service = _create_build_service(
+            conn,
+            embedding_provider=embedding_provider,
+            description_provider=description_provider,
+            chunk_embedding_mode=chunk_embedding_mode,
+        )
+        return build(
+            project_dir=project_dir,
+            build_service=build_service,
+        )
+
+
+def _create_build_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+    description_provider: DescriptionProvider | None,
+    chunk_embedding_mode: ChunkEmbeddingMode,
+) -> BuildService:
+    chunk_embedding_service = create_chunk_embedding_service(
+        conn,
+        embedding_provider=embedding_provider,
+        description_provider=description_provider,
+        chunk_embedding_mode=chunk_embedding_mode,
+    )
+
+    return BuildService(
+        chunk_embedding_service=chunk_embedding_service,
+    )

databao_context_engine/build_sources/export_results.py

@@ -0,0 +1,41 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.build_sources.plugin_execution import BuiltDatasourceContext
+from databao_context_engine.datasources.datasource_context import get_context_header_for_datasource
+from databao_context_engine.datasources.types import DatasourceId
+from databao_context_engine.project.layout import ALL_RESULTS_FILE_NAME
+from databao_context_engine.serialization.yaml import write_yaml_to_stream
+
+logger = logging.getLogger(__name__)
+
+
+def export_build_result(output_dir: Path, result: BuiltDatasourceContext) -> Path:
+    datasource_id = DatasourceId.from_string_repr(result.datasource_id)
+    export_file_path = output_dir.joinpath(datasource_id.relative_path_to_context_file())
+
+    # Make sure the parent folder exists
+    export_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with export_file_path.open("w") as export_file:
+        write_yaml_to_stream(data=result, file_stream=export_file)
+
+    logger.info(f"Exported result to {export_file_path.resolve()}")
+
+    return export_file_path
+
+
+def append_result_to_all_results(output_dir: Path, result: BuiltDatasourceContext):
+    path = output_dir.joinpath(ALL_RESULTS_FILE_NAME)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a", encoding="utf-8") as export_file:
+        export_file.write(get_context_header_for_datasource(DatasourceId.from_string_repr(result.datasource_id)))
+        write_yaml_to_stream(data=result, file_stream=export_file)
+        export_file.write("\n")
+
+
+# Here temporarily because it needs to be reset between runs.
+# A subsequent PR will remove the existence of the all_results file
+def reset_all_results(output_dir: Path):
+    path = output_dir.joinpath(ALL_RESULTS_FILE_NAME)
+    path.unlink(missing_ok=True)

databao_context_engine/build_sources/{internal/plugin_execution.py → plugin_execution.py}

@@ -2,20 +2,18 @@ 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.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     BuildFilePlugin,
     BuildPlugin,
 )
 from databao_context_engine.pluginlib.plugin_utils import execute_datasource_plugin, execute_file_plugin
-from databao_context_engine.project.types import PreparedConfig, PreparedDatasource, DatasourceId
 
 
 @dataclass()
 class BuiltDatasourceContext:
-    """
-    Dataclass defining the result of building a datasource's context.
-    """
+    """Dataclass defining the result of building a datasource's context."""
 
     datasource_id: str
     """
@@ -53,9 +51,7 @@ def execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> Bui
 
 
 def _execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> Any:
-    """
-    Run a prepared source through the plugin
-    """
+    """Run a prepared source through the plugin."""
     if isinstance(prepared_datasource, PreparedConfig):
         ds_plugin = cast(BuildDatasourcePlugin, plugin)
 

databao_context_engine/cli/add_datasource_config.py

@@ -5,22 +5,27 @@ from typing import Any
 
 import click
 
-from databao_context_engine import DatabaoContextProjectManager, DatasourceId, DatasourceType
-from databao_context_engine.datasource_config.add_config import (
-    get_config_file_structure_for_datasource_type,
+from databao_context_engine import (
+    ConfigPropertyDefinition,
+    DatabaoContextPluginLoader,
+    DatabaoContextProjectManager,
+    DatasourceId,
+    DatasourceType,
 )
-from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
-from databao_context_engine.plugins.plugin_loader import get_all_available_plugin_types
+from databao_context_engine.pluginlib.config import ConfigUnionPropertyDefinition
 
 
 def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
     project_manager = DatabaoContextProjectManager(project_dir=project_dir)
+    plugin_loader = DatabaoContextPluginLoader()
 
     click.echo(
         f"We will guide you to add a new datasource in your Databao Context Engine project, at {project_dir.resolve()}"
     )
 
-    datasource_type = _ask_for_datasource_type()
+    datasource_type = _ask_for_datasource_type(
+        plugin_loader.get_all_supported_datasource_types(exclude_file_plugins=True)
+    )
     datasource_name = click.prompt("Datasource name?", type=str)
 
     is_datasource_existing = project_manager.datasource_config_exists(
@@ -33,7 +38,9 @@ def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
             default=False,
         )
 
-    config_content = _ask_for_config_details(datasource_type)
+    config_content = _ask_for_config_details(
+        plugin_loader.get_config_file_structure_for_datasource_type(datasource_type)
+    )
 
     config_file = project_manager.create_datasource_config(
         datasource_type, datasource_name, config_content, overwrite_existing=True
@@ -44,10 +51,8 @@ def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
     return config_file.datasource_id
 
 
-def _ask_for_datasource_type() -> DatasourceType:
-    supported_types_by_folder = _group_supported_types_by_folder(
-        get_all_available_plugin_types(exclude_file_plugins=True)
-    )
+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(
@@ -68,9 +73,7 @@ def _ask_for_datasource_type() -> DatasourceType:
     return DatasourceType.from_main_and_subtypes(config_folder, config_type)
 
 
-def _ask_for_config_details(datasource_type: DatasourceType) -> dict[str, Any]:
-    config_file_structure = get_config_file_structure_for_datasource_type(datasource_type)
-
+def _ask_for_config_details(config_file_structure: list[ConfigPropertyDefinition]) -> dict[str, Any]:
     # Adds a new line before asking for the config values specific to the plugin
     click.echo("")
 
@@ -84,13 +87,35 @@ def _ask_for_config_details(datasource_type: DatasourceType) -> dict[str, Any]:
 
 
 def _build_config_content_from_properties(
-    properties: list[ConfigPropertyDefinition], properties_prefix: str
+    properties: list[ConfigPropertyDefinition], properties_prefix: str, in_union: bool = False
 ) -> dict[str, Any]:
     config_content: dict[str, Any] = {}
     for config_file_property in properties:
         if config_file_property.property_key in ["type", "name"] and len(properties_prefix) == 0:
             # We ignore type and name properties as they've already been filled
             continue
+        if in_union and config_file_property.property_key == "type":
+            continue
+
+        if isinstance(config_file_property, ConfigUnionPropertyDefinition):
+            choices = {t.__name__: t for t in config_file_property.types}
+
+            chosen = click.prompt(
+                f"{properties_prefix}{config_file_property.property_key}.type?",
+                type=click.Choice(sorted(choices.keys())),
+            )
+
+            chosen_type = choices[chosen]
+
+            nested_props = config_file_property.type_properties[chosen_type]
+            nested_content = _build_config_content_from_properties(
+                nested_props, f"{properties_prefix}{config_file_property.property_key}.", in_union=True
+            )
+
+            config_content[config_file_property.property_key] = {
+                **nested_content,
+            }
+            continue
 
         if config_file_property.nested_properties is not None and len(config_file_property.nested_properties) > 0:
             nested_content = _build_config_content_from_properties(
@@ -114,6 +139,7 @@ def _build_config_content_from_properties(
                 type=str,
                 default=default_value,
                 show_default=default_value is not None and default_value != "",
+                hide_input=config_file_property.secret,
             )
 
             if property_value.strip():

databao_context_engine/cli/commands.py

@@ -19,7 +19,6 @@ 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
-from databao_context_engine.storage.migrate import migrate
 
 
 @click.group()
@@ -49,25 +48,18 @@ def dce(ctx: Context, verbose: bool, quiet: bool, project_dir: str | None) -> No
     ctx.obj["quiet"] = quiet
     ctx.obj["project_dir"] = project_path
 
-    migrate()
-
 
 @dce.command()
 @click.pass_context
 def info(ctx: Context) -> None:
-    """
-    Display system-wide information
-    """
-
+    """Display system-wide information."""
     echo_info(ctx.obj["project_dir"])
 
 
 @dce.command()
 @click.pass_context
 def init(ctx: Context) -> None:
-    """
-    Create an empty Databao Context Engine project
-    """
+    """Create an empty Databao Context Engine project."""
     project_dir = ctx.obj["project_dir"]
     try:
         init_dce_project(project_dir=project_dir)
@@ -97,17 +89,14 @@ def init(ctx: Context) -> None:
 
 @dce.group()
 def datasource() -> None:
-    """
-    Manage datasource configurations
-    """
+    """Manage datasource configurations."""
     pass
 
 
 @datasource.command(name="add")
 @click.pass_context
 def add_datasource_config(ctx: Context) -> None:
-    """
-    Add a new datasource configuration.
+    """Add a new datasource configuration.
 
     The command will ask all relevant information for that datasource and save it in your Databao Context Engine project.
     """
@@ -122,15 +111,13 @@ def add_datasource_config(ctx: Context) -> None:
 )
 @click.pass_context
 def check_datasource_config(ctx: Context, datasources_config_files: list[str] | None) -> None:
-    """
-    Check whether a datasource configuration is valid.
+    """Check whether a datasource configuration is valid.
 
     The configuration is considered as valid if a connection with the datasource can be established.
 
     By default, all datasources declared in the project will be checked.
     You can explicitely list which datasources to validate by using the [DATASOURCES_CONFIG_FILES] argument. Each argument must be the path to the file within the src folder (e.g: my-folder/my-config.yaml)
     """
-
     datasource_ids = (
         [DatasourceId.from_string_repr(datasource_config_file) for datasource_config_file in datasources_config_files]
         if datasources_config_files is not None
@@ -157,10 +144,9 @@ def build(
         "embeddable_text_only", "generated_description_only", "embeddable_text_and_generated_description"
     ],
 ) -> None:
-    """
-    Build context for all datasources
+    """Build context for all datasources.
 
-    The output of the build command will be saved in a "run" folder in the output directory.
+    The output of the build command will be saved in the output directory.
 
     Internally, this indexes the context to be used by the MCP server and the "retrieve" command.
     """
@@ -177,12 +163,6 @@ def build(
     nargs=-1,
     required=True,
 )
-@click.option(
-    "-r",
-    "--run-name",
-    type=click.STRING,
-    help="Build run to use (the run folder name). Defaults to the latest run in the project.",
-)
 @click.option(
     "-l",
     "--limit",
@@ -200,19 +180,16 @@ def build(
 def retrieve(
     ctx: Context,
     retrieve_text: tuple[str, ...],
-    run_name: str | None,
     limit: int | None,
     output_format: Literal["file", "streamed"],
 ) -> None:
-    """
-    Search the project's built context for the most relevant chunks.
-    """
+    """Search the project's built context for the most relevant chunks."""
     text = " ".join(retrieve_text)
 
     databao_engine = DatabaoContextEngine(project_dir=ctx.obj["project_dir"])
 
     retrieve_results = databao_engine.search_context(
-        retrieve_text=text, run_name=run_name, limit=limit, export_to_file=output_format == "file"
+        retrieve_text=text, limit=limit, export_to_file=output_format == "file"
     )
 
     if output_format == "streamed":
@@ -221,12 +198,6 @@ def retrieve(
 
 
 @dce.command()
-@click.option(
-    "-r",
-    "--run-name",
-    type=click.STRING,
-    help="Name of the build run you want to use (aka. the name of the run folder in your project's output). Defaults to the latest one in the project.",
-)
 @click.option(
     "-H",
     "--host",
@@ -247,10 +218,8 @@ def retrieve(
     help="Transport to use. Defaults to stdio",
 )
 @click.pass_context
-def mcp(ctx: Context, run_name: str | None, host: str | None, port: int | None, transport: McpTransport) -> None:
-    """
-    Run Databao Context Engine's MCP server
-    """
+def mcp(ctx: Context, host: str | None, port: int | None, transport: McpTransport) -> None:
+    """Run Databao Context Engine's MCP server."""
     if transport == "stdio":
         configure_logging(verbose=False, quiet=True, project_dir=ctx.obj["project_dir"])
-    run_mcp_server(project_dir=ctx.obj["project_dir"], run_name=run_name, transport=transport, host=host, port=port)
+    run_mcp_server(project_dir=ctx.obj["project_dir"], transport=transport, host=host, port=port)

databao_context_engine/cli/info.py

@@ -15,6 +15,7 @@ def _generate_info_string(command_info: DceInfo) -> str:
     info_lines = []
     info_lines.append(f"Databao context engine version: {command_info.version}")
     info_lines.append(f"Databao context engine storage dir: {command_info.dce_path}")
+    info_lines.append(f"Databao context engine plugins: {command_info.plugin_ids}")
 
     info_lines.append("")
 
@@ -23,10 +24,10 @@ def _generate_info_string(command_info: DceInfo) -> str:
 
     info_lines.append("")
 
-    if command_info.project_info.is_initialised:
+    if command_info.project_info.is_initialized:
         info_lines.append(f"Project dir: {command_info.project_info.project_path.resolve()}")
         info_lines.append(f"Project ID: {str(command_info.project_info.project_id)}")
     else:
-        info_lines.append(f"Project not initialised at {command_info.project_info.project_path.resolve()}")
+        info_lines.append(f"Project not initialized at {command_info.project_info.project_path.resolve()}")
 
     return os.linesep.join(info_lines)