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

databao_context_engine/__init__.py

@@ -0,0 +1,35 @@
+from databao_context_engine.build_sources.internal.build_runner import BuildContextResult
+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 (
+    CheckDatasourceConnectionResult,
+    DatasourceConnectionStatus,
+)
+from databao_context_engine.datasource_config.datasource_context import DatasourceContext
+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.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__ = [
+    "DatabaoContextEngine",
+    "Datasource",
+    "DatasourceId",
+    "DatasourceContext",
+    "ContextSearchResult",
+    "DatabaoContextProjectManager",
+    "ChunkEmbeddingMode",
+    "BuildContextResult",
+    "CheckDatasourceConnectionResult",
+    "DatasourceConnectionStatus",
+    "DatasourceConfigFile",
+    "DatasourceType",
+    "get_databao_context_engine_info",
+    "DceInfo",
+    "init_dce_project",
+    "init_or_get_dce_project",
+    "InitErrorReason",
+    "InitProjectError",
+]

databao_context_engine/build_sources/internal/build_runner.py

@@ -0,0 +1,111 @@
+import logging
+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 (
+    append_result_to_all_results,
+    create_run_dir,
+    export_build_result,
+)
+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
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BuildContextResult:
+    datasource_id: DatasourceId
+    datasource_type: DatasourceType
+    context_built_at: datetime
+    context_file_path: Path
+
+
+def build(
+    project_dir: Path,
+    *,
+    build_service: BuildService,
+    project_id: str,
+    dce_version: str,
+) -> list[BuildContextResult]:
+    """
+    Build entrypoint.
+
+    1) Load available plugins
+    2) Discover sources
+    3) Create a run
+    4) For each source, call process_source
+    """
+    plugins = load_plugins()
+
+    datasources = discover_datasources(project_dir)
+
+    if not datasources:
+        logger.info("No sources discovered under %s", project_dir)
+        return []
+
+    run = None
+    run_dir = None
+
+    number_of_failed_builds = 0
+    build_result = []
+    for discovered_datasource in datasources:
+        try:
+            prepared_source = prepare_source(discovered_datasource)
+
+            logger.info(
+                f'Found datasource of type "{prepared_source.datasource_type.full_type}" with name {prepared_source.path.stem}'
+            )
+
+            plugin = plugins.get(prepared_source.datasource_type)
+            if plugin is None:
+                logger.warning(
+                    "No plugin for '%s' (datasource=%s) — skipping.",
+                    prepared_source.datasource_type.full_type,
+                    prepared_source.path,
+                )
+                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)
+
+            build_result.append(
+                BuildContextResult(
+                    datasource_id=DatasourceId.from_string_repr(result.datasource_id),
+                    datasource_type=DatasourceType(full_type=result.datasource_type),
+                    context_built_at=result.context_built_at,
+                    context_file_path=context_file_path,
+                )
+            )
+        except Exception as e:
+            logger.debug(str(e), exc_info=True, stack_info=True)
+            logger.info(f"Failed to build source at ({discovered_datasource.path}): {str(e)}")
+
+            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),
+        f"Failed to build {number_of_failed_builds}." if number_of_failed_builds > 0 else "",
+    )
+
+    return build_result

databao_context_engine/build_sources/internal/build_service.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from databao_context_engine.build_sources.internal.plugin_execution import BuiltDatasourceContext, execute
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildPlugin,
+)
+from databao_context_engine.project.types import PreparedDatasource
+from databao_context_engine.serialisation.yaml import to_yaml_string
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingService
+from databao_context_engine.storage.models import RunDTO
+from databao_context_engine.storage.repositories.datasource_run_repository import DatasourceRunRepository
+from databao_context_engine.storage.repositories.run_repository import RunRepository
+
+logger = logging.getLogger(__name__)
+
+
+class BuildService:
+    def __init__(
+        self,
+        *,
+        run_repo: RunRepository,
+        datasource_run_repo: DatasourceRunRepository,
+        chunk_embedding_service: ChunkEmbeddingService,
+    ) -> None:
+        self._run_repo = run_repo
+        self._datasource_run_repo = datasource_run_repo
+        self._chunk_embedding_service = chunk_embedding_service
+
+    def start_run(self, *, project_id: str, dce_version: str) -> RunDTO:
+        """
+        Create a new run row and return (run_id, started_at).
+        """
+        return self._run_repo.create(project_id=project_id, dce_version=dce_version)
+
+    def finalize_run(self, *, run_id: int):
+        """
+        Mark the run as complete (sets ended_at).
+        """
+        self._run_repo.update(run_id=run_id, ended_at=datetime.now())
+
+    def process_prepared_source(
+        self,
+        *,
+        run_id: int,
+        prepared_source: PreparedDatasource,
+        plugin: BuildPlugin,
+    ) -> BuiltDatasourceContext:
+        """
+        Process a single source.
+
+        1) Execute the plugin
+        2) Divide the results into chunks
+        3) Embed and persist the chunks
+        """
+        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
+
+        datasource_run = self._datasource_run_repo.create(
+            run_id=run_id,
+            plugin=plugin.name,
+            full_type=prepared_source.datasource_type.full_type,
+            source_id=result.datasource_id,
+            storage_directory=str(prepared_source.path.parent),
+        )
+
+        self._chunk_embedding_service.embed_chunks(
+            datasource_run_id=datasource_run.datasource_run_id, chunks=chunks, result=to_yaml_string(result.context)
+        )
+
+        return result

databao_context_engine/build_sources/internal/build_wiring.py

@@ -0,0 +1,52 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.build_sources.internal.build_runner import BuildContextResult, build
+from databao_context_engine.llm.factory import (
+    create_ollama_description_provider,
+    create_ollama_embedding_provider,
+    create_ollama_service,
+)
+from databao_context_engine.project.info import get_dce_version
+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_build_service,
+)
+from databao_context_engine.storage.connection import open_duckdb_connection
+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]:
+    """
+    Public build entrypoint
+    - Instantiates the build service
+    - Delegates the actual build logic to the build runner
+    """
+    project_layout = ensure_project_dir(project_dir)
+
+    logger.debug(f"Starting to build datasources in project {project_dir.resolve()}")
+
+    with open_duckdb_connection(get_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,
+        )
+        dce_config = project_layout.read_config_file()
+        return build(
+            project_dir=project_dir,
+            build_service=build_service,
+            project_id=str(dce_config.project_id),
+            dce_version=get_dce_version(),
+        )

databao_context_engine/build_sources/internal/export_results.py

@@ -0,0 +1,43 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.build_sources.internal.plugin_execution import BuiltDatasourceContext
+from databao_context_engine.datasource_config.datasource_context import get_context_header_for_datasource
+from databao_context_engine.project.layout import ALL_RESULTS_FILE_NAME, get_output_dir
+from databao_context_engine.project.types import DatasourceId
+from databao_context_engine.serialisation.yaml import write_yaml_to_stream
+
+logger = logging.getLogger(__name__)
+
+
+def create_run_dir(project_dir: Path, run_name: str) -> Path:
+    output_dir = get_output_dir(project_dir)
+
+    run_dir = output_dir.joinpath(run_name)
+    run_dir.mkdir(parents=True, exist_ok=False)
+
+    return run_dir
+
+
+def export_build_result(run_dir: Path, result: BuiltDatasourceContext) -> Path:
+    datasource_id = DatasourceId.from_string_repr(result.datasource_id)
+    export_file_path = run_dir.joinpath(datasource_id.relative_path_to_context_file())
+
+    # Make sure the parent folder exists
+    export_file_path.parent.mkdir(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(run_dir: Path, result: BuiltDatasourceContext):
+    path = run_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")

databao_context_engine/build_sources/internal/plugin_execution.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, cast
+
+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.
+    """
+
+    datasource_id: str
+    """
+    The ID of the built data source
+    """
+
+    datasource_type: str
+    """
+    The type of the built data source
+    """
+
+    context_built_at: datetime
+    """
+    The time at which the context was built
+    """
+
+    context: Any
+    """
+    A dictionary containing the actual context generated for the data source.
+    This dictionary should be serializable in YAML format.
+    """
+
+
+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_type=prepared_datasource.datasource_type.full_type,
+        context_built_at=datetime.now(),
+        context=built_context,
+    )
+
+
+def _execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> Any:
+    """
+    Run a prepared source through the plugin
+    """
+    if isinstance(prepared_datasource, PreparedConfig):
+        ds_plugin = cast(BuildDatasourcePlugin, plugin)
+
+        return execute_datasource_plugin(
+            plugin=ds_plugin,
+            datasource_type=prepared_datasource.datasource_type,
+            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,
+        )

databao_context_engine/build_sources/public/api.py

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

databao_context_engine/cli/add_datasource_config.py

@@ -0,0 +1,130 @@
+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.pluginlib.config import ConfigPropertyDefinition
+from databao_context_engine.plugins.plugin_loader import get_all_available_plugin_types
+
+
+def add_datasource_config_interactive(project_dir: Path) -> DatasourceId:
+    project_manager = DatabaoContextProjectManager(project_dir=project_dir)
+
+    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_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:
+        click.confirm(
+            f"A config file already exists for this datasource ({datasource_type.config_folder}/{datasource_name}). Do you want to overwrite it?",
+            abort=True,
+            default=False,
+        )
+
+    config_content = _ask_for_config_details(datasource_type)
+
+    config_file = project_manager.create_datasource_config(
+        datasource_type, datasource_name, config_content, overwrite_existing=True
+    )
+
+    click.echo(f"{os.linesep}We've created a new config file for your datasource at: {config_file.config_file_path}")
+
+    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])
+    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,
+    )
+    click.echo(f"Selected subtype: {config_type}")
+
+    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)
+
+    # Adds a new line before asking for the config values specific to the plugin
+    click.echo("")
+
+    if len(config_file_structure) == 0:
+        click.echo(
+            "The plugin for this datasource doesn't declare its configuration. Please check the documentation and fill the config file manually."
+        )
+        return {}
+
+    return _build_config_content_from_properties(config_file_structure, properties_prefix="")
+
+
+def _build_config_content_from_properties(
+    properties: list[ConfigPropertyDefinition], properties_prefix: str
+) -> 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 config_file_property.nested_properties is not None and len(config_file_property.nested_properties) > 0:
+            nested_content = _build_config_content_from_properties(
+                config_file_property.nested_properties,
+                properties_prefix=f"{properties_prefix}.{config_file_property.property_key}."
+                if properties_prefix
+                else f"{config_file_property.property_key}.",
+            )
+            if len(nested_content.keys()) > 0:
+                config_content[config_file_property.property_key] = nested_content
+        else:
+            default_value: str | None
+            if config_file_property.default_value:
+                default_value = config_file_property.default_value
+            else:
+                # We need to add an empty string default value for non-required fields
+                default_value = None if config_file_property.required else ""
+
+            property_value = click.prompt(
+                f"{properties_prefix}{config_file_property.property_key}? {'(Optional)' if not config_file_property.required else ''}",
+                type=str,
+                default=default_value,
+                show_default=default_value is not None and default_value != "",
+            )
+
+            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