databao-context-engine 0.1.1__py3-none-any.whl → 0.1.5__py3-none-any.whl

databao_context_engine/__init__.py

@@ -1,16 +1,30 @@
-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.project.info import DceInfo, get_databao_context_engine_info
+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,
+    BuildFilePlugin,
+    BuildPlugin,
+    DatasourceType,
+)
+from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
+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.project.types import Datasource, DatasourceId
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
 __all__ = [
@@ -28,8 +42,19 @@ __all__ = [
     "DatasourceType",
     "get_databao_context_engine_info",
     "DceInfo",
+    "DceProjectInfo",
     "init_dce_project",
     "init_or_get_dce_project",
     "InitErrorReason",
     "InitProjectError",
+    "DatabaoContextPluginLoader",
+    "ConfigPropertyDefinition",
+    "BuildPlugin",
+    "BuildDatasourcePlugin",
+    "BuildFilePlugin",
+    "install_ollama_if_needed",
+    "download_ollama_models_if_needed",
+    "OllamaError",
+    "OllamaTransientError",
+    "OllamaPermanentError",
 ]

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 ProjectLayout
 
 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
@@ -26,33 +36,34 @@ class BuildContextResult:
 
 
 def build(
-    project_dir: Path,
+    project_layout: ProjectLayout,
     *,
     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()
 
-    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 []
 
-    run = None
-    run_dir = None
-
     number_of_failed_builds = 0
     build_result = []
+    reset_all_results(project_layout.output_dir)
     for discovered_datasource in datasources:
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -71,23 +82,19 @@ 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 = 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,
@@ -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,82 @@
+import logging
+
+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 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
+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_layout: ProjectLayout, 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.
+    """
+    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_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:
+        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_layout=project_layout,
+            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 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
     """
@@ -41,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,
@@ -53,9 +48,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)
 
@@ -65,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,39 +1,43 @@
 import os
-from collections import defaultdict
 from pathlib import Path
 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(
-        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,
         )
 
-    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,33 +48,19 @@ 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)
-    )
-
-    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])
+def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType:
+    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(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 +74,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,17 +126,10 @@ 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():
                 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