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

databao_context_engine/cli/commands.py

@@ -0,0 +1,256 @@
+import sys
+from pathlib import Path
+from typing import Literal
+
+import click
+from click import Context
+
+from databao_context_engine import (
+    ChunkEmbeddingMode,
+    DatabaoContextEngine,
+    DatabaoContextProjectManager,
+    DatasourceId,
+    InitErrorReason,
+    InitProjectError,
+    init_dce_project,
+)
+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
+from databao_context_engine.storage.migrate import migrate
+
+
+@click.group()
+@click.option("-v", "--verbose", is_flag=True, help="Enable debug logging")
+@click.option("-q", "--quiet", is_flag=True, help="Disable all console logging")
+@click.option(
+    "-d",
+    "--project-dir",
+    type=click.STRING,
+    help="Location of your Databao Context Engine project",
+)
+@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)
+        exit(1)
+
+    if project_dir is None:
+        project_path = Path.cwd()
+    else:
+        project_path = Path(project_dir).expanduser()
+
+    configure_logging(verbose=verbose, quiet=quiet, project_dir=project_path)
+
+    ctx.ensure_object(dict)
+    ctx.obj["verbose"] = verbose
+    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
+    """
+
+    echo_info(ctx.obj["project_dir"])
+
+
+@dce.command()
+@click.pass_context
+def init(ctx: Context) -> None:
+    """
+    Create an empty Databao Context Engine project
+    """
+    project_dir = ctx.obj["project_dir"]
+    try:
+        init_dce_project(project_dir=project_dir)
+    except InitProjectError as e:
+        if e.reason == InitErrorReason.PROJECT_DIR_DOESNT_EXIST:
+            if click.confirm(
+                f"The directory {ctx.obj['project_dir'].resolve()} does not exist. Do you want to create it?",
+                default=True,
+            ):
+                project_dir.mkdir(parents=True, exist_ok=False)
+                init_dce_project(project_dir=project_dir)
+            else:
+                return
+        else:
+            raise e
+
+    click.echo(f"Project initialized successfully at {project_dir.resolve()}")
+
+    try:
+        resolve_ollama_bin()
+    except RuntimeError as e:
+        click.echo(str(e), err=True)
+
+    if click.confirm("\nDo you want to configure a datasource now?"):
+        add_datasource_config_cli(project_dir)
+
+
+@dce.group()
+def datasource() -> None:
+    """
+    Manage datasource configurations
+    """
+    pass
+
+
+@datasource.command(name="add")
+@click.pass_context
+def add_datasource_config(ctx: Context) -> None:
+    """
+    Add a new datasource configuration.
+
+    The command will ask all relevant information for that datasource and save it in your Databao Context Engine project.
+    """
+    add_datasource_config_cli(ctx.obj["project_dir"])
+
+
+@datasource.command(name="check")
+@click.argument(
+    "datasources-config-files",
+    type=click.STRING,
+    nargs=-1,
+)
+@click.pass_context
+def check_datasource_config(ctx: Context, datasources_config_files: list[str] | None) -> None:
+    """
+    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
+        else None
+    )
+
+    check_datasource_connection_cli(ctx.obj["project_dir"], datasource_ids=datasource_ids)
+
+
+@dce.command()
+@click.option(
+    "-m",
+    "--chunk-embedding-mode",
+    type=click.Choice(
+        ["embeddable_text_only", "generated_description_only", "embeddable_text_and_generated_description"]
+    ),
+    default="embeddable_text_only",
+    help="Choose how chunks will be embedded. If a mode with the generated_description is selected, a local LLM model will be downloaded and used.",
+)
+@click.pass_context
+def build(
+    ctx: Context,
+    chunk_embedding_mode: Literal[
+        "embeddable_text_only", "generated_description_only", "embeddable_text_and_generated_description"
+    ],
+) -> None:
+    """
+    Build context for all datasources
+
+    The output of the build command will be saved in a "run" folder in the output directory.
+
+    Internally, this indexes the context to be used by the MCP server and the "retrieve" command.
+    """
+    result = DatabaoContextProjectManager(project_dir=ctx.obj["project_dir"]).build_context(
+        datasource_ids=None, chunk_embedding_mode=ChunkEmbeddingMode(chunk_embedding_mode.upper())
+    )
+
+    click.echo(f"Build complete. Processed {len(result)} datasources.")
+
+
+@dce.command()
+@click.argument(
+    "retrieve-text",
+    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",
+    type=click.INT,
+    help="Maximum number of chunk matches to return.",
+)
+@click.option(
+    "-o",
+    "--output-format",
+    type=click.Choice(["file", "streamed"]),
+    default="file",
+    help="The output format [file (default), streamed]",
+)
+@click.pass_context
+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.
+    """
+    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"
+    )
+
+    if output_format == "streamed":
+        display_texts = [context_search_result.context_result for context_search_result in retrieve_results]
+        click.echo("\n".join(display_texts))
+
+
+@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",
+    type=click.STRING,
+    help="Host to bind to. Defaults to 127.0.0.1",
+)
+@click.option(
+    "-p",
+    "--port",
+    type=click.INT,
+    help="Port to bind to. Defaults to 8000",
+)
+@click.option(
+    "-t",
+    "--transport",
+    type=click.Choice(["stdio", "streamable-http"]),
+    default="stdio",
+    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
+    """
+    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)

databao_context_engine/cli/datasources.py

@@ -0,0 +1,64 @@
+import os
+from pathlib import Path
+
+import click
+
+from databao_context_engine import (
+    CheckDatasourceConnectionResult,
+    DatabaoContextProjectManager,
+    DatasourceConnectionStatus,
+    DatasourceId,
+)
+from databao_context_engine.cli.add_datasource_config import add_datasource_config_interactive
+
+
+def add_datasource_config_cli(project_dir: Path) -> None:
+    datasource_id = add_datasource_config_interactive(project_dir)
+
+    if click.confirm("\nDo you want to check the connection to this new datasource?"):
+        check_datasource_connection_cli(project_dir, datasource_ids=[datasource_id])
+
+
+def check_datasource_connection_cli(project_dir: Path, *, datasource_ids: list[DatasourceId] | None) -> None:
+    results = DatabaoContextProjectManager(project_dir=project_dir).check_datasource_connection(
+        datasource_ids=datasource_ids
+    )
+
+    _print_check_datasource_connection_results(results)
+
+
+def _print_check_datasource_connection_results(results: list[CheckDatasourceConnectionResult]) -> None:
+    if len(results) > 0:
+        valid_datasources = [
+            result for result in results if result.connection_status == DatasourceConnectionStatus.VALID
+        ]
+        invalid_datasources = [
+            result for result in results if result.connection_status == DatasourceConnectionStatus.INVALID
+        ]
+        unknown_datasources = [
+            result for result in results if result.connection_status == DatasourceConnectionStatus.UNKNOWN
+        ]
+
+        # Print all errors
+        for check_result in invalid_datasources:
+            click.echo(
+                f"Error for datasource {str(check_result.datasource_id)}:{os.linesep}{check_result.full_message}{os.linesep}"
+            )
+
+        results_summary = (
+            os.linesep.join(
+                [
+                    f"{str(check_result.datasource_id)}: {check_result.format(show_summary_only=True)}"
+                    for check_result in results
+                ]
+            )
+            if results
+            else "No datasource found"
+        )
+
+        click.echo(
+            f"Validation completed with {len(valid_datasources)} valid datasource(s) and {len(invalid_datasources) + len(unknown_datasources)} invalid (or unknown status) datasource(s)"
+            f"{os.linesep}{results_summary}"
+        )
+    else:
+        click.echo("No datasource found")

databao_context_engine/cli/info.py

@@ -0,0 +1,32 @@
+import os
+import sys
+from pathlib import Path
+
+import click
+
+from databao_context_engine import DceInfo, get_databao_context_engine_info
+
+
+def echo_info(project_dir: Path) -> None:
+    click.echo(_generate_info_string(get_databao_context_engine_info(project_dir=project_dir)))
+
+
+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("")
+
+    info_lines.append(f"OS name: {sys.platform}")
+    info_lines.append(f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}")
+
+    info_lines.append("")
+
+    if command_info.project_info.is_initialised:
+        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()}")
+
+    return os.linesep.join(info_lines)

databao_context_engine/config/log_config.yaml

@@ -0,0 +1,16 @@
+version: 1
+formatters:
+  main:
+    format: '%(asctime)s %(levelname)s %(name)s: %(message)s'
+  console:
+    format: '%(asctime)s %(levelname)s: %(message)s'
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: console
+    stream: ext://sys.stdout
+loggers:
+  databao_context_engine:
+    level: INFO
+    propagate: false
+    handlers: [ console ]

databao_context_engine/config/logging.py

@@ -0,0 +1,43 @@
+from datetime import datetime
+from logging.config import dictConfig
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from databao_context_engine.project.layout import get_logs_dir, is_project_dir_valid
+
+
+def configure_logging(verbose: bool, quiet: bool, project_dir: Path) -> None:
+    with Path(__file__).parent.joinpath("log_config.yaml").open(mode="r") as log_config_file:
+        log_config = yaml.safe_load(log_config_file)
+
+        if is_project_dir_valid(project_dir):
+            logs_dir_path = get_logs_dir(project_dir)
+            logs_dir_path.mkdir(exist_ok=True)
+
+            file_handler_name = "logFile"
+            log_config["handlers"][file_handler_name] = _get_logging_file_handler(logs_dir_path)
+            log_config["loggers"]["databao_context_engine"]["handlers"].append(file_handler_name)
+
+        if quiet:
+            log_config["loggers"]["databao_context_engine"]["handlers"].remove("console")
+        if verbose:
+            log_config["loggers"]["databao_context_engine"]["level"] = "DEBUG"
+
+        dictConfig(log_config)
+
+
+def _get_logging_file_handler(logs_dir_path: Path) -> dict[str, Any]:
+    return {
+        "filename": str(logs_dir_path.joinpath(_get_current_log_filename())),
+        "class": "logging.handlers.RotatingFileHandler",
+        "formatter": "main",
+        "maxBytes": 100000000,  # 100MB
+        "backupCount": 12,
+    }
+
+
+def _get_current_log_filename() -> str:
+    # Creates a new log file every month
+    return datetime.now().strftime("log-%Y-%m.txt")

databao_context_engine/databao_context_project_manager.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, overload
+
+from databao_context_engine.build_sources.public.api import BuildContextResult, build_all_datasources
+from databao_context_engine.datasource_config.add_config import (
+    create_datasource_config_file,
+    get_datasource_id_for_config_file,
+)
+from databao_context_engine.datasource_config.check_config import (
+    CheckDatasourceConnectionResult,
+    check_datasource_connection as check_datasource_connection_internal,
+)
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.datasource_discovery import DatasourceId
+from databao_context_engine.project.layout import ensure_datasource_config_file_doesnt_exist, ensure_project_dir
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
+
+
+@dataclass
+class DatasourceConfigFile:
+    datasource_id: DatasourceId
+    config_file_path: Path
+
+
+class DatabaoContextProjectManager:
+    project_dir: Path
+
+    def __init__(self, project_dir: Path) -> None:
+        ensure_project_dir(project_dir=project_dir)
+        self.project_dir = project_dir
+
+    def build_context(
+        self, datasource_ids: list[DatasourceId] | None, chunk_embedding_mode: ChunkEmbeddingMode
+    ) -> list[BuildContextResult]:
+        # TODO: Filter which datasources to build by datasource_ids
+        return build_all_datasources(project_dir=self.project_dir, chunk_embedding_mode=chunk_embedding_mode)
+
+    def check_datasource_connection(
+        self, datasource_ids: list[DatasourceId] | None
+    ) -> list[CheckDatasourceConnectionResult]:
+        return sorted(
+            check_datasource_connection_internal(project_dir=self.project_dir, datasource_ids=datasource_ids).values(),
+            key=lambda result: str(result.datasource_id),
+        )
+
+    def create_datasource_config(
+        self,
+        datasource_type: DatasourceType,
+        datasource_name: str,
+        config_content: dict[str, Any],
+        overwrite_existing: bool = False,
+    ) -> DatasourceConfigFile:
+        # 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,
+            datasource_type=datasource_type,
+            datasource_name=datasource_name,
+            config_content=config_content,
+            overwrite_existing=overwrite_existing,
+        )
+
+        return DatasourceConfigFile(
+            datasource_id=DatasourceId.from_datasource_config_file_path(config_file_path),
+            config_file_path=config_file_path,
+        )
+
+    @overload
+    def datasource_config_exists(self, *, datasource_type: DatasourceType, datasource_name: str) -> bool: ...
+    @overload
+    def datasource_config_exists(self, *, datasource_id: DatasourceId) -> bool: ...
+
+    def datasource_config_exists(
+        self,
+        *,
+        datasource_type: DatasourceType | None = None,
+        datasource_name: str | None = None,
+        datasource_id: DatasourceId | None = None,
+    ) -> bool:
+        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:
+            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
+        except ValueError:
+            return True

databao_context_engine/databao_engine.py

@@ -0,0 +1,85 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.datasource_config.datasource_context import (
+    DatasourceContext,
+    get_all_contexts,
+    get_context_header_for_datasource,
+    get_datasource_context,
+)
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.datasource_discovery import get_datasource_list
+from databao_context_engine.project.layout import ensure_project_dir
+from databao_context_engine.project.types import Datasource, DatasourceId
+from databao_context_engine.retrieve_embeddings.public.api import retrieve_embeddings
+
+
+@dataclass
+class ContextSearchResult:
+    datasource_id: DatasourceId
+    datasource_type: DatasourceType
+    distance: float
+    context_result: str
+
+
+class DatabaoContextEngine:
+    project_dir: Path
+
+    def __init__(self, project_dir: Path) -> None:
+        self.project_layout = ensure_project_dir(project_dir=project_dir)
+        self.project_dir = project_dir
+
+    def get_datasource_list(self) -> list[Datasource]:
+        # TODO: Should this return the list of built datasources rather than the list of datasources within the src folder?
+        return get_datasource_list(self.project_dir)
+
+    def get_datasource_context(self, datasource_id: DatasourceId, run_name: str | None = None) -> DatasourceContext:
+        return get_datasource_context(
+            project_layout=self.project_layout, datasource_id=datasource_id, run_name=run_name
+        )
+
+    def get_all_contexts(self, run_name: str | None = None) -> list[DatasourceContext]:
+        return get_all_contexts(project_layout=self.project_layout, run_name=run_name)
+
+    def get_all_contexts_formatted(self, run_name: str | None = None) -> str:
+        all_contexts = self.get_all_contexts(run_name=run_name)
+
+        all_results = 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,
+        run_name: str | None,
+        limit: int | None,
+        export_to_file: bool,
+        datasource_ids: list[DatasourceId] | None = None,
+    ) -> list[ContextSearchResult]:
+        # TODO: Filter with datasource_ids
+        # TODO: Remove the need for a run_name
+
+        results = retrieve_embeddings(
+            project_layout=self.project_layout,
+            retrieve_text=retrieve_text,
+            run_name=run_name,
+            limit=limit,
+            export_to_file=export_to_file,
+        )
+
+        return [
+            ContextSearchResult(
+                datasource_id=result.datasource_id,
+                datasource_type=result.datasource_type,
+                distance=result.cosine_distance,
+                context_result=result.display_text,
+            )
+            for result in results
+        ]
+
+    def run_sql(self, datasource_id: DatasourceId, sql: str, params: list[str]) -> dict[str, Any]:
+        raise NotImplementedError("Running SQL is not supported yet")

databao_context_engine/datasource_config/add_config.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.introspection.property_extract import get_property_list_from_type
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, DatasourceType
+from databao_context_engine.pluginlib.config import ConfigPropertyDefinition, CustomiseConfigProperties
+from databao_context_engine.plugins.plugin_loader import get_plugin_for_type
+from databao_context_engine.project.layout import (
+    create_datasource_config_file as create_datasource_config_file_internal,
+)
+from databao_context_engine.project.types import DatasourceId
+from databao_context_engine.serialisation.yaml import to_yaml_string
+
+
+def get_config_file_structure_for_datasource_type(datasource_type: DatasourceType) -> list[ConfigPropertyDefinition]:
+    plugin = get_plugin_for_type(datasource_type)
+
+    if isinstance(plugin, CustomiseConfigProperties):
+        return plugin.get_config_file_properties()
+    elif isinstance(plugin, BuildDatasourcePlugin):
+        return get_property_list_from_type(plugin.config_file_type)
+    else:
+        raise ValueError(
+            f"Impossible to create a config for type {datasource_type.full_type}. The plugin for this type is not a BuildDatasourcePlugin or CustomiseConfigProperties"
+        )
+
+
+def create_datasource_config_file(
+    project_dir: Path,
+    datasource_type: DatasourceType,
+    datasource_name: str,
+    config_content: dict[str, Any],
+    overwrite_existing: bool,
+) -> Path:
+    basic_config = {"type": datasource_type.subtype, "name": datasource_name}
+
+    return create_datasource_config_file_internal(
+        project_dir,
+        get_datasource_id_for_config_file(datasource_type, datasource_name),
+        to_yaml_string(basic_config | config_content),
+        overwrite_existing=overwrite_existing,
+    )
+
+
+def get_datasource_id_for_config_file(datasource_type: DatasourceType, datasource_name: str) -> DatasourceId:
+    return DatasourceId(
+        datasource_config_folder=datasource_type.config_folder,
+        datasource_name=datasource_name,
+        config_file_suffix=".yaml",
+    )