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

databao_context_engine/cli/commands.py

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

databao_context_engine/databao_context_engine.py

@@ -0,0 +1,127 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.datasources.datasource_context import (
+    DatasourceContext,
+    get_all_contexts,
+    get_context_header_for_datasource,
+    get_datasource_context,
+    get_introspected_datasource_list,
+)
+from databao_context_engine.datasources.types import Datasource, DatasourceId
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.layout import ProjectLayout, ensure_project_dir
+from databao_context_engine.retrieve_embeddings import retrieve_embeddings
+
+
+@dataclass
+class ContextSearchResult:
+    """The result of a search in the project's contexts.
+
+    Attributes:
+        datasource_id: The ID of the datasource that generated the result.
+        datasource_type: The type of the datasource that generated the result.
+        distance: The distance between the search text and the result.
+        context_result: The actual content of the result that was found as a YAML string.
+            This content will be a subpart of the full context of the datasource.
+            In some cases, its content won't contain the exact same attributes as what can be
+            found directly in the full context.
+    """
+
+    datasource_id: DatasourceId
+    datasource_type: DatasourceType
+    distance: float
+    context_result: str
+
+
+class DatabaoContextEngine:
+    """Engine for reading and using the contexts generated in a Databao Context Project.
+
+    The Databao Context Project should already have datasources configured and built (see DatabaoContextProjectManager), so that they can be used in the Engine.
+
+    Attributes:
+        project_dir: The root directory of the Databao Context Project.
+    """
+
+    project_dir: Path
+    _project_layout: ProjectLayout
+
+    def __init__(self, project_dir: Path) -> None:
+        """Initialize the DatabaoContextEngine.
+
+        Args:
+            project_dir: The root directory of the Databao Context Project.
+                There must be a valid DatabaoContextProject in this directory.
+        """
+        self._project_layout = ensure_project_dir(project_dir=project_dir)
+        self.project_dir = project_dir
+
+    def get_introspected_datasource_list(self) -> list[Datasource]:
+        """Return the list of datasources for which a context is available.
+
+        Returns:
+            A list of the datasources for which a context is available.
+        """
+        return get_introspected_datasource_list(self._project_layout)
+
+    def get_datasource_context(self, datasource_id: DatasourceId) -> DatasourceContext:
+        """Return the context available for a given datasource.
+
+        Args:
+            datasource_id: The ID of the datasource.
+
+        Returns:
+            The context for this datasource.
+        """
+        return get_datasource_context(project_layout=self._project_layout, datasource_id=datasource_id)
+
+    def get_all_contexts(self) -> list[DatasourceContext]:
+        """Return all contexts generated in the project.
+
+        Returns:
+             A list of all contexts generated in the project.
+        """
+        return get_all_contexts(project_layout=self._project_layout)
+
+    def get_all_contexts_formatted(self) -> str:
+        all_contexts = self.get_all_contexts()
+
+        return os.linesep.join(
+            [f"{get_context_header_for_datasource(context.datasource_id)}{context.context}" for context in all_contexts]
+        )
+
+    def search_context(
+        self,
+        retrieve_text: str,
+        limit: int | None = None,
+        datasource_ids: list[DatasourceId] | None = None,
+    ) -> list[ContextSearchResult]:
+        """Search in the avaialable context for the closest matches to the given text.
+
+        Args:
+            retrieve_text: The text to search for in the contexts.
+            limit: The maximum number of results to return. If None is provided, a default limit of 10 will be used.
+            datasource_ids: If provided, the search results will only come from the datasources with these IDs.
+
+        Returns:
+            A list of the results found for the search, sorted by distance.
+        """
+        results = retrieve_embeddings(
+            project_layout=self._project_layout, retrieve_text=retrieve_text, limit=limit, datasource_ids=datasource_ids
+        )
+
+        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]:
+        """Not Implemented yet. This will allow to run a SQL query against a datasource (if the datasource supports it)."""
+        raise NotImplementedError("Running SQL is not supported yet")

databao_context_engine/databao_context_project_manager.py

@@ -2,45 +2,108 @@ 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 (
+from databao_context_engine.build_sources import BuildContextResult, build_all_datasources
+from databao_context_engine.databao_context_engine import DatabaoContextEngine
+from databao_context_engine.datasources.check_config import (
     CheckDatasourceConnectionResult,
+)
+from databao_context_engine.datasources.check_config import (
     check_datasource_connection as check_datasource_connection_internal,
 )
+from databao_context_engine.datasources.datasource_discovery import get_datasource_list
+from databao_context_engine.datasources.types import Datasource, DatasourceId
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
-from databao_context_engine.project.datasource_discovery import DatasourceId
-from databao_context_engine.project.layout import ensure_datasource_config_file_doesnt_exist, ensure_project_dir
+from databao_context_engine.project.layout import (
+    ProjectLayout,
+    ensure_project_dir,
+)
+from databao_context_engine.project.layout import (
+    create_datasource_config_file as create_datasource_config_file_internal,
+)
+from databao_context_engine.serialization.yaml import to_yaml_string
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
 
 @dataclass
 class DatasourceConfigFile:
+    """A datasource config file that was created by the DatabaoContextProjectManager.
+
+    Attributes:
+        datasource_id: The unique identifier for the datasource.
+        config_file_path: The path to the datasource configuration file.
+    """
+
     datasource_id: DatasourceId
     config_file_path: Path
 
 
 class DatabaoContextProjectManager:
+    """Project Manager for Databao Context Projects.
+
+    This project manager is responsible for configuring and building a Databao Context Project.
+    The project_dir should already have been initialized before a Project manager can be used.
+
+    Attributes:
+        project_dir: The root directory of the Databao Context Project.
+    """
+
     project_dir: Path
+    _project_layout: ProjectLayout
 
     def __init__(self, project_dir: Path) -> None:
-        ensure_project_dir(project_dir=project_dir)
+        """Initialize the DatabaoContextProjectManager.
+
+        Args:
+            project_dir: The root directory of the Databao Context Project.
+        """
+        self._project_layout = ensure_project_dir(project_dir=project_dir)
         self.project_dir = project_dir
 
+    def get_configured_datasource_list(self) -> list[Datasource]:
+        """Return the list of datasources configured in the project.
+
+        This method returns all datasources configured in the src folder of the project,
+        no matter whether the datasource configuration is valid or not.
+
+        Returns:
+            The list of datasources configured in the project.
+        """
+        return get_datasource_list(self._project_layout)
+
     def build_context(
-        self, datasource_ids: list[DatasourceId] | None, chunk_embedding_mode: ChunkEmbeddingMode
+        self,
+        datasource_ids: list[DatasourceId] | None = None,
+        chunk_embedding_mode: ChunkEmbeddingMode = ChunkEmbeddingMode.EMBEDDABLE_TEXT_ONLY,
     ) -> list[BuildContextResult]:
+        """Build the context for datasources in the project.
+
+        Any datasource with an invalid configuration will be skipped.
+
+        Args:
+            datasource_ids: The list of datasource ids to build. If None, all datasources will be built.
+            chunk_embedding_mode: The mode to use for chunk embedding.
+
+        Returns:
+            The list of all built results.
+        """
         # TODO: Filter which datasources to build by datasource_ids
-        return build_all_datasources(project_dir=self.project_dir, chunk_embedding_mode=chunk_embedding_mode)
+        return build_all_datasources(project_layout=self._project_layout, chunk_embedding_mode=chunk_embedding_mode)
 
     def check_datasource_connection(
-        self, datasource_ids: list[DatasourceId] | None
+        self, datasource_ids: list[DatasourceId] | None = None
     ) -> list[CheckDatasourceConnectionResult]:
+        """Check the connection for datasources in the project.
+
+        Args:
+            datasource_ids: The list of datasource ids to check. If None, all datasources will be checked.
+
+        Returns:
+            The list of all connection check results, sorted by datasource id.
+        """
         return sorted(
-            check_datasource_connection_internal(project_dir=self.project_dir, datasource_ids=datasource_ids).values(),
+            check_datasource_connection_internal(
+                project_layout=self._project_layout, datasource_ids=datasource_ids
+            ).values(),
             key=lambda result: str(result.datasource_id),
         )
 
@@ -51,9 +114,21 @@ class DatabaoContextProjectManager:
         config_content: dict[str, Any],
         overwrite_existing: bool = False,
     ) -> DatasourceConfigFile:
+        """Create a new datasource configuration file in the project.
+
+        Args:
+            datasource_type: The type of the datasource to create.
+            datasource_name: The name of the datasource to create.
+            config_content: The content of the datasource configuration. This is a dictionary that will be written as-is in a yaml file.
+                The actual content of the configuration is not checked and might not be valid for the requested type..
+            overwrite_existing: Whether to overwrite an existing datasource configuration file if it already exists.
+
+        Returns:
+            The path to the created datasource configuration file.
+        """
         # TODO: Before creating the datasource, validate the config content based on which plugin will be used
-        config_file_path = create_datasource_config_file(
-            project_dir=self.project_dir,
+        config_file = _create_datasource_config_file(
+            project_layout=self._project_layout,
             datasource_type=datasource_type,
             datasource_name=datasource_name,
             config_content=config_content,
@@ -61,32 +136,74 @@ class DatabaoContextProjectManager:
         )
 
         return DatasourceConfigFile(
-            datasource_id=DatasourceId.from_datasource_config_file_path(config_file_path),
-            config_file_path=config_file_path,
+            datasource_id=DatasourceId.from_datasource_config_file_path(self._project_layout, config_file),
+            config_file_path=config_file,
         )
 
     @overload
-    def datasource_config_exists(self, *, datasource_type: DatasourceType, datasource_name: str) -> bool: ...
+    def datasource_config_exists(self, *, datasource_name: str) -> DatasourceId | None: ...
     @overload
-    def datasource_config_exists(self, *, datasource_id: DatasourceId) -> bool: ...
+    def datasource_config_exists(self, *, datasource_id: DatasourceId) -> DatasourceId | None: ...
 
     def datasource_config_exists(
         self,
         *,
-        datasource_type: DatasourceType | None = None,
         datasource_name: str | None = None,
         datasource_id: DatasourceId | None = None,
-    ) -> bool:
-        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:
+    ) -> DatasourceId | None:
+        """Check if a datasource configuration file already exists in the project.
+
+        Args:
+            datasource_name: The name of the datasource.
+            datasource_id: The id of the datasource. If provided, datasource_type and datasource_name will be ignored.
+
+        Returns:
+            datasource_id if there is already a datasource configuration file for this datasource, None otherwise.
+
+        Raises:
+            ValueError: If the wrong set of arguments is provided.
+        """
+        if datasource_name is not None:
+            datasource_id = DatasourceId.from_string_repr(f"{datasource_name}.yaml")
+            config_file = self._project_layout.src_dir / datasource_id.relative_path_to_config_file()
+            if config_file.is_file():
+                return datasource_id
+            return None
+
+        if datasource_id is None:
             raise ValueError("Either datasource_id or both datasource_type and datasource_name must be provided")
 
         try:
-            ensure_datasource_config_file_doesnt_exist(
-                project_dir=self.project_dir,
-                datasource_id=datasource_id,
-            )
-            return False
+            config_file = self._project_layout.src_dir.joinpath(datasource_id.relative_path_to_config_file())
+            if config_file.is_file():
+                raise ValueError(f"A config file already exists for {str(datasource_id)}")
+
+            return datasource_id
         except ValueError:
-            return True
+            return None
+
+    def get_engine_for_project(self) -> DatabaoContextEngine:
+        """Instantiate a DatabaoContextEngine for the project.
+
+        Returns:
+            A DatabaoContextEngine instance for the project.
+        """
+        return DatabaoContextEngine(project_dir=self.project_dir)
+
+
+def _create_datasource_config_file(
+    project_layout: ProjectLayout,
+    datasource_type: DatasourceType,
+    datasource_name: str,
+    config_content: dict[str, Any],
+    overwrite_existing: bool,
+) -> Path:
+    last_datasource_name = datasource_name.split("/")[-1]
+    basic_config = {"type": datasource_type.full_type, "name": last_datasource_name}
+
+    return create_datasource_config_file_internal(
+        project_layout,
+        f"{datasource_name}.yaml",
+        to_yaml_string(basic_config | config_content),
+        overwrite_existing=overwrite_existing,
+    )