databao-context-engine 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

databao_context_engine/databao_context_engine.py

@@ -0,0 +1,137 @@
+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()
+
+        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,
+        limit: int | None = None,
+        export_to_file: bool = False,
+        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.
+            export_to_file: Whether the results should be exported to a file as a side-effect. If True, the results will be exported in a file in the run directory.
+            datasource_ids: Not Implemented yet: providing this argument changes nothing to the search
+
+        Returns:
+            A list of the results found for the search, sorted by distance.
+        """
+        # TODO: Filter with datasource_ids
+        # TODO: When no run_name is required, we can extract the "export_to_file" side-effect and let the caller (the CLI) do it themselves
+
+        results = retrieve_embeddings(
+            project_layout=self._project_layout,
+            retrieve_text=retrieve_text,
+            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]:
+        """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,43 +2,100 @@ 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 (
+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.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.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.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
 
     def __init__(self, project_dir: Path) -> None:
+        """Initialize the DatabaoContextProjectManager.
+
+        Args:
+            project_dir: The root directory of the Databao Context Project.
+        """
         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_dir)
+
     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)
 
     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(),
             key=lambda result: str(result.datasource_id),
@@ -51,6 +108,18 @@ 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,
@@ -77,6 +146,19 @@ class DatabaoContextProjectManager:
         datasource_name: str | None = None,
         datasource_id: DatasourceId | None = None,
     ) -> bool:
+        """Check if a datasource configuration file already exists in the project.
+
+        Args:
+            datasource_type: The type of the datasource. This must be used in conjunction with datasource_name.
+            datasource_name: The name of the datasource. This must be used in conjunction with datasource_type.
+            datasource_id: The id of the datasource. If provided, datasource_type and datasource_name will be ignored.
+
+        Returns:
+            True if there is already a datasource configuration file for this datasource, False otherwise.
+
+        Raises:
+            ValueError: If the wrong set of arguments is provided.
+        """
         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:
@@ -90,3 +172,11 @@ class DatabaoContextProjectManager:
             return False
         except ValueError:
             return True
+
+    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)

databao_context_engine/datasources/add_config.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.datasources.types import DatasourceId
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+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
+
+
+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",
+    )

databao_context_engine/{datasource_config → datasources}/check_config.py

@@ -6,21 +6,23 @@ from pathlib import Path
 
 from pydantic import ValidationError
 
-from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, NotSupportedError
-from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource
-from databao_context_engine.plugins.plugin_loader import load_plugins
-from databao_context_engine.project.datasource_discovery import (
+from databao_context_engine.datasources.datasource_discovery import (
     discover_datasources,
     get_datasource_descriptors,
     prepare_source,
 )
+from databao_context_engine.datasources.types import DatasourceId, PreparedConfig
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, NotSupportedError
+from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource
+from databao_context_engine.plugins.plugin_loader import load_plugins
 from databao_context_engine.project.layout import ensure_project_dir
-from databao_context_engine.project.types import DatasourceId, PreparedConfig
 
 logger = logging.getLogger(__name__)
 
 
 class DatasourceConnectionStatus(Enum):
+    """Status of the connection to a datasource."""
+
     VALID = "Valid"
     INVALID = "Invalid"
     UNKNOWN = "Unknown"
@@ -28,6 +30,15 @@ class DatasourceConnectionStatus(Enum):
 
 @dataclass(kw_only=True)
 class CheckDatasourceConnectionResult:
+    """Result of checking the connection status of a datasource.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        connection_status: The connection status of the datasource.
+        summary: A summary of the connection status' error, or None if the connection is valid.
+        full_message: A detailed message about the connection status' error, or None if the connection is valid.
+    """
+
     datasource_id: DatasourceId
     connection_status: DatasourceConnectionStatus
     summary: str | None
@@ -103,12 +114,12 @@ def check_datasource_connection(
                     exc_info=True,
                     stack_info=True,
                 )
-                result[result_key] = get_validation_result_from_error(result_key, e)
+                result[result_key] = _get_validation_result_from_error(result_key, e)
 
     return result
 
 
-def get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
+def _get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
     if isinstance(e, ValidationError):
         return CheckDatasourceConnectionResult(
             datasource_id=datasource_id,

databao_context_engine/datasources/datasource_context.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+from databao_context_engine.datasources.types import Datasource, DatasourceId, DatasourceType
+from databao_context_engine.project.layout import ProjectLayout, get_output_dir
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(eq=True, frozen=True)
+class DatasourceContext:
+    """A generated Context for a Datasource.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        context: The context generated for the datasource.
+    """
+
+    datasource_id: DatasourceId
+    # TODO: Read the context as a BuildExecutionResult instead of a Yaml string?
+    context: str
+
+
+def _read_datasource_type_from_context_file(context_path: Path) -> DatasourceType:
+    with context_path.open("r") as context_file:
+        type_key = "datasource_type"
+        for line in context_file:
+            if line.startswith(f"{type_key}: "):
+                datasource_type = yaml.safe_load(line)[type_key]
+                return DatasourceType(full_type=datasource_type)
+
+    raise ValueError(f"Could not find type in context file {context_path}")
+
+
+def get_introspected_datasource_list(project_layout: ProjectLayout) -> list[Datasource]:
+    output_dir = get_output_dir(project_dir=project_layout.project_dir)
+
+    result = []
+    for main_type_dir in sorted((p for p in output_dir.iterdir() if p.is_dir()), key=lambda p: p.name.lower()):
+        for context_path in sorted(
+            (p for p in main_type_dir.iterdir() if p.suffix in [".yaml", ".yml"]), key=lambda p: p.name.lower()
+        ):
+            try:
+                result.append(
+                    Datasource(
+                        id=DatasourceId.from_datasource_config_file_path(context_path),
+                        type=_read_datasource_type_from_context_file(context_path),
+                    )
+                )
+            except ValueError as e:
+                logger.debug(str(e), exc_info=True, stack_info=True)
+                logger.warning(
+                    f"Ignoring introspected datasource: Failed to read datasource_type from context file at {context_path.resolve()}"
+                )
+
+    return result
+
+
+def get_datasource_context(project_layout: ProjectLayout, datasource_id: DatasourceId) -> DatasourceContext:
+    output_dir = get_output_dir(project_dir=project_layout.project_dir)
+
+    context_path = output_dir.joinpath(datasource_id.relative_path_to_context_file())
+    if not context_path.is_file():
+        raise ValueError(f"Context file not found for datasource {str(datasource_id)}")
+
+    context = context_path.read_text()
+    return DatasourceContext(datasource_id=datasource_id, context=context)
+
+
+def get_all_contexts(project_layout: ProjectLayout) -> list[DatasourceContext]:
+    output_dir = get_output_dir(project_dir=project_layout.project_dir)
+
+    result = []
+    for main_type_dir in sorted((p for p in output_dir.iterdir() if p.is_dir()), key=lambda p: p.name.lower()):
+        for context_path in sorted(
+            (p for p in main_type_dir.iterdir() if p.suffix in [".yaml", ".yml"]), key=lambda p: p.name.lower()
+        ):
+            result.append(
+                DatasourceContext(
+                    datasource_id=DatasourceId.from_datasource_context_file_path(context_path),
+                    context=context_path.read_text(),
+                )
+            )
+
+    return result
+
+
+def get_context_header_for_datasource(datasource_id: DatasourceId) -> str:
+    return f"# ===== {str(datasource_id)} ====={os.linesep}"

databao_context_engine/{project → datasources}/datasource_discovery.py

@@ -4,17 +4,17 @@ from typing import Any
 
 import yaml
 
-from databao_context_engine.pluginlib.build_plugin import DatasourceType
-from databao_context_engine.project.layout import get_source_dir
-from databao_context_engine.project.types import (
+from databao_context_engine.datasources.types import (
+    Datasource,
     DatasourceDescriptor,
+    DatasourceId,
     DatasourceKind,
     PreparedConfig,
     PreparedDatasource,
     PreparedFile,
-    DatasourceId,
-    Datasource,
 )
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.layout import get_source_dir
 from databao_context_engine.templating.renderer import render_template
 
 logger = logging.getLogger(__name__)
@@ -41,13 +41,18 @@ def get_datasource_list(project_dir: Path) -> list[Datasource]:
 
 
 def discover_datasources(project_dir: Path) -> list[DatasourceDescriptor]:
-    """
-    Scan the project's src/ directory and return all discovered sources.
+    """Scan the project's src/ directory and return all discovered sources.
 
     Rules:
         - Each first-level directory under src/ is treated as a main_type
         - Unsupported or unreadable entries are skipped.
         - The returned list is sorted by directory and then filename
+
+    Returns:
+        A list of DatasourceDescriptor instances representing the discovered datasources.
+
+    Raises:
+        ValueError: If the src directory does not exist in the project directory.
     """
     src = get_source_dir(project_dir)
     if not src.exists() or not src.is_dir():
@@ -56,7 +61,7 @@ def discover_datasources(project_dir: Path) -> list[DatasourceDescriptor]:
     datasources: list[DatasourceDescriptor] = []
     for main_dir in sorted((p for p in src.iterdir() if p.is_dir()), key=lambda p: p.name.lower()):
         for path in sorted((p for p in main_dir.iterdir() if _is_datasource_file(p)), key=lambda p: p.name.lower()):
-            datasource = load_datasource_descriptor(path)
+            datasource = _load_datasource_descriptor(path)
             if datasource is not None:
                 datasources.append(datasource)
 
@@ -79,17 +84,15 @@ def get_datasource_descriptors(project_dir: Path, datasource_ids: list[Datasourc
         if not config_file_path.is_file():
             raise ValueError(f"Datasource config file not found: {config_file_path}")
 
-        datasource = load_datasource_descriptor(config_file_path)
+        datasource = _load_datasource_descriptor(config_file_path)
         if datasource is not None:
             datasources.append(datasource)
 
     return datasources
 
 
-def load_datasource_descriptor(path: Path) -> DatasourceDescriptor | None:
-    """
-    Load a single file with src/<parent_name>/ into a DatasourceDescriptor
-    """
+def _load_datasource_descriptor(path: Path) -> DatasourceDescriptor | None:
+    """Load a single file with src/<parent_name>/ into a DatasourceDescriptor."""
     if not path.is_file():
         return None
 
@@ -110,9 +113,7 @@ def load_datasource_descriptor(path: Path) -> DatasourceDescriptor | None:
 
 
 def prepare_source(datasource: DatasourceDescriptor) -> PreparedDatasource:
-    """
-    Convert a discovered datasource into a prepared datasource ready for plugin execution
-    """
+    """Convert a discovered datasource into a prepared datasource ready for plugin execution."""
     if datasource.kind is DatasourceKind.FILE:
         file_subtype = datasource.path.suffix.lower().lstrip(".")
         return PreparedFile(