databao-context-engine 0.1.2__py3-none-any.whl → 0.1.4.dev1__py3-none-any.whl

databao_context_engine/{project → datasources}/types.py

@@ -38,12 +38,16 @@ PreparedDatasource = PreparedConfig | PreparedFile
 
 @dataclass(kw_only=True, frozen=True, eq=True)
 class DatasourceId:
-    """
-    The ID of a datasource. The ID is the path to the datasource's config file relative to the src folder in the project.
+    """The ID of a datasource. The ID is the path to the datasource's config file relative to the src folder in the project.
 
     e.g: "databases/my_postgres_datasource.yaml"
 
     Use the provided factory methods `from_string_repr` and `from_datasource_config_file_path` to create a DatasourceId, rather than its constructor.
+
+    Attributes:
+        datasource_config_folder: The folder where the datasource's config file is located.
+        datasource_name: The name of the datasource.
+        config_file_suffix: The suffix of the config (or raw) file.
     """
 
     datasource_config_folder: str
@@ -78,18 +82,22 @@ class DatasourceId:
         return str(self.relative_path_to_config_file())
 
     def relative_path_to_config_file(self) -> Path:
-        """
-        Returns a path to the config file for this datasource.
+        """Return a path to the config file for this datasource.
 
         The returned path is relative to the src folder in the project.
+
+        Returns:
+            The path to the config file relative to the src folder in the project.
         """
         return Path(self.datasource_config_folder).joinpath(self.datasource_name + self.config_file_suffix)
 
     def relative_path_to_context_file(self) -> Path:
-        """
-        Returns a path to the config file for this datasource.
+        """Return a path to the config file for this datasource.
 
         The returned path is relative to an output run folder in the project.
+
+        Returns:
+            The path to the context file relative to the output folder in the project.
         """
         # Keep the suffix in the filename if this datasource is a raw file, to handle multiple files with the same name and different extensions
         suffix = ".yaml" if self.config_file_suffix == ".yaml" else (self.config_file_suffix + ".yaml")
@@ -97,13 +105,19 @@ class DatasourceId:
         return Path(self.datasource_config_folder).joinpath(self.datasource_name + suffix)
 
     @classmethod
-    def from_string_repr(cls, datasource_id_as_string: str):
-        """
-        Creates a DatasourceId from a string representation.
+    def from_string_repr(cls, datasource_id_as_string: str) -> "DatasourceId":
+        """Create a DatasourceId from a string representation.
+
+        Args:
+            datasource_id_as_string: The string representation of a DatasourceId.
+                This is the path to the datasource's config file relative to the src folder in the project.
+                (e.g. "databases/my_postgres_datasource.yaml")
 
-        The string representation of a DatasourceId is the path to the datasource's config file relative to the src folder in the project.
+        Returns:
+            The DatasourceId instance created from the string representation.
 
-        e.g: "databases/my_postgres_datasource.yaml"
+        Raises:
+            ValueError: If the string representation is invalid (e.g. too many parent folders).
         """
         config_file_path = Path(datasource_id_as_string)
 
@@ -115,11 +129,15 @@ class DatasourceId:
         return DatasourceId.from_datasource_config_file_path(config_file_path)
 
     @classmethod
-    def from_datasource_config_file_path(cls, datasource_config_file: Path):
-        """
-        Creates a DatasourceId from a config file path.
+    def from_datasource_config_file_path(cls, datasource_config_file: Path) -> "DatasourceId":
+        """Create a DatasourceId from a config file path.
 
-        The `datasource_config_file` path provided can either be the config file path relative to the src folder or the full path to the config file.
+        Args:
+            datasource_config_file: The path to the datasource config file.
+                This path can either be the config file path relative to the src folder or the full path to the config file.
+
+        Returns:
+            The DatasourceId instance created from the config file path.
         """
         return DatasourceId(
             datasource_config_folder=datasource_config_file.parent.name,
@@ -127,8 +145,39 @@ class DatasourceId:
             config_file_suffix=datasource_config_file.suffix,
         )
 
+    @classmethod
+    def from_datasource_context_file_path(cls, datasource_context_file: Path) -> "DatasourceId":
+        """Create a DatasourceId from a context file path.
+
+        This factory handles the case where the context was generated from a raw file rather than from a config.
+        In that case, the context file name will look like "<my_datasource_name>.<raw_file_extension>.yaml"
+
+        Args:
+            datasource_context_file: The path to the datasource context file.
+
+        Returns:
+            The DatasourceId instance created from the context file path.
+        """
+        if len(datasource_context_file.suffixes) > 1 and datasource_context_file.suffix == ".yaml":
+            # If there is more than 1 suffix, we remove the latest suffix (.yaml) to keep only the actual datasource file suffix
+            context_file_name_without_yaml_extension = datasource_context_file.name[: -len(".yaml")]
+            datasource_context_file = datasource_context_file.with_name(context_file_name_without_yaml_extension)
+
+        return DatasourceId(
+            datasource_config_folder=datasource_context_file.parent.name,
+            datasource_name=datasource_context_file.stem,
+            config_file_suffix=datasource_context_file.suffix,
+        )
+
 
 @dataclass
 class Datasource:
+    """A datasource contained in the project.
+
+    Attributes:
+        id: The unique identifier of the datasource.
+        type: The type of the datasource.
+    """
+
     id: DatasourceId
     type: DatasourceType

databao_context_engine/init_project.py

@@ -1,17 +1,39 @@
 from pathlib import Path
 
 from databao_context_engine.databao_context_project_manager import DatabaoContextProjectManager
-from databao_context_engine.project.init_project import init_project_dir
+from databao_context_engine.project.init_project import InitProjectError, init_project_dir
 from databao_context_engine.project.layout import is_project_dir_valid
 
 
 def init_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
-    initialised_project_dir = init_project_dir(project_dir=project_dir)
+    """Initialize a Databao Context project in the given directory.
 
-    return DatabaoContextProjectManager(project_dir=initialised_project_dir)
+    Args:
+        project_dir: The directory where the project should be initialized.
+
+    Returns:
+        A DatabaoContextProjectManager for the initialized project.
+
+    Raises:
+        InitProjectError: If the project could not be initialized.
+    """
+    try:
+        initialized_project_dir = init_project_dir(project_dir=project_dir)
+
+        return DatabaoContextProjectManager(project_dir=initialized_project_dir)
+    except InitProjectError as e:
+        raise e
 
 
 def init_or_get_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
+    """Initialize a Databao Context project in the given directory or get a DatabaoContextProjectManager for the project if it already exists.
+
+    Args:
+        project_dir: The directory where the project should be initialized if it doesn't exist yet..
+
+    Returns:
+        A DatabaoContextProjectManager for the project in project_dir.
+    """
     if is_project_dir_valid(project_dir):
         return DatabaoContextProjectManager(project_dir=project_dir)
 

databao_context_engine/introspection/property_extract.py

@@ -1,11 +1,16 @@
 import types
 from dataclasses import MISSING, fields, is_dataclass
-from typing import Annotated, Any, ForwardRef, Union, get_origin, get_type_hints
+from typing import Annotated, Any, ForwardRef, Union, get_args, get_origin, get_type_hints
 
-from pydantic import BaseModel, _internal
+from pydantic import BaseModel
 from pydantic_core import PydanticUndefinedType
 
-from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation, ConfigPropertyDefinition
+from databao_context_engine.pluginlib.config import (
+    ConfigPropertyAnnotation,
+    ConfigPropertyDefinition,
+    ConfigSinglePropertyDefinition,
+    ConfigUnionPropertyDefinition,
+)
 
 
 def get_property_list_from_type(root_type: type) -> list[ConfigPropertyDefinition]:
@@ -54,18 +59,14 @@ def _get_property_list_from_dataclass(parent_type: type) -> list[ConfigPropertyD
         raise ValueError(f"{parent_type} is not a dataclass")
 
     dataclass_fields = fields(parent_type)
+    type_hints = get_type_hints(parent_type, include_extras=True)
 
     result = []
     for field in dataclass_fields:
         has_field_default = field.default is not None and field.default != MISSING
 
-        if isinstance(field.type, str):
-            try:
-                property_type = _evaluate_type_string(field.type)
-            except Exception:
-                continue
-        else:
-            property_type = field.type
+        # Use the type hints if the field type wasn't resolved (aka. if it is a ForwardRef or a str)
+        property_type = type_hints[field.name] if isinstance(field.type, ForwardRef | str) else field.type
 
         property_for_field = _create_property(
             property_type=property_type,
@@ -84,6 +85,10 @@ def _get_property_list_from_pydantic_base_model(parent_type: type):
     if not issubclass(parent_type, BaseModel):
         raise ValueError(f"{parent_type} is not a Pydantic BaseModel")
 
+    if any(isinstance(field.annotation, ForwardRef) for field in parent_type.model_fields.values()):
+        # If any field's future type wasn't resolved yet, we rebuild the model to resolve them
+        parent_type.model_rebuild(force=True)
+
     pydantic_fields = parent_type.model_fields
     result = []
 
@@ -123,26 +128,55 @@ def _create_property(
     if annotation is not None and annotation.ignored_for_config_wizard:
         return None
 
-    actual_property_type = _read_actual_property_type(property_type)
+    actual_property_types = _read_actual_property_type(property_type)
 
-    try:
-        nested_properties = _get_property_list_from_type(parent_type=actual_property_type, is_root_type=False)
-    except TypeError:
-        return None
+    required = annotation.required if annotation else is_property_required
+    secret = annotation.secret if annotation else False
 
-    default_value = compute_default_value(
-        annotation=annotation,
-        property_default=property_default,
-        has_nested_properties=nested_properties is not None and len(nested_properties) > 0,
-    )
+    if len(actual_property_types) > 1:
+        type_properties: dict[type, list[ConfigPropertyDefinition]] = {}
 
-    return ConfigPropertyDefinition(
-        property_key=property_name,
-        property_type=actual_property_type if not nested_properties else None,
-        required=annotation.required if annotation else is_property_required,
-        default_value=default_value,
-        nested_properties=nested_properties if nested_properties else None,
-    )
+        for union_type in actual_property_types:
+            try:
+                nested_props = _get_property_list_from_type(
+                    parent_type=union_type,
+                    is_root_type=False,
+                )
+            except TypeError:
+                nested_props = []
+
+            type_properties[union_type] = nested_props
+
+        return ConfigUnionPropertyDefinition(
+            property_key=property_name,
+            types=actual_property_types,
+            type_properties=type_properties,
+        )
+    else:
+        actual_property_type = actual_property_types[0]
+        try:
+            nested_properties = _get_property_list_from_type(
+                parent_type=actual_property_type,
+                is_root_type=False,
+            )
+        except TypeError:
+            return None
+
+        resolved_type = actual_property_type if not nested_properties else None
+        default_value = compute_default_value(
+            annotation=annotation,
+            property_default=property_default,
+            has_nested_properties=nested_properties is not None and len(nested_properties) > 0,
+        )
+
+        return ConfigSinglePropertyDefinition(
+            property_key=property_name,
+            property_type=resolved_type,
+            required=required,
+            default_value=default_value,
+            nested_properties=nested_properties or None,
+            secret=secret,
+        )
 
 
 def _get_config_property_annotation(property_type) -> ConfigPropertyAnnotation | None:
@@ -155,21 +189,16 @@ def _get_config_property_annotation(property_type) -> ConfigPropertyAnnotation |
     return None
 
 
-def _read_actual_property_type(property_type: type) -> type:
+def _read_actual_property_type(property_type: type) -> tuple[type, ...]:
     property_type_origin = get_origin(property_type)
 
     if property_type_origin is Annotated:
-        return property_type.__origin__  # type: ignore[attr-defined]
-    elif property_type_origin is Union or property_type_origin is types.UnionType:
-        type_args = property_type.__args__  # type: ignore[attr-defined]
-        if len(type_args) == 2 and type(None) in type_args:
-            # Uses the actual type T when the Union is "T | None" (or "None | T")
-            return next(arg for arg in type_args if arg is not None)
-        else:
-            # Ignoring Union types when it is not used as type | None as we wouldn't which type to pick
-            return type(None)
+        return _read_actual_property_type(property_type.__origin__)  # type: ignore[attr-defined]
+    elif property_type_origin in (Union, types.UnionType):
+        args = tuple(arg for arg in get_args(property_type) if arg is not type(None))
+        return args
 
-    return property_type
+    return (property_type,)
 
 
 def compute_default_value(
@@ -185,18 +214,3 @@ def compute_default_value(
         return str(property_default)
 
     return None
-
-
-def _evaluate_type_string(property_type: str) -> type:
-    try:
-        # Using a pydantic internal function for this, to avoid having to implement type evaluation manually...
-        return _internal._typing_extra.eval_type(property_type)
-    except Exception as initial_error:
-        try:
-            # Try to convert it ourselves if Pydantic didn't work
-            return ForwardRef(property_type)._evaluate(  # type: ignore[return-value]
-                globalns=globals(), localns=locals(), recursive_guard=frozenset()
-            )
-        except Exception as e:
-            # Ignore if we didn't manage to convert the str to a type
-            raise e from initial_error

databao_context_engine/llm/errors.py

@@ -3,14 +3,8 @@ class OllamaError(Exception):
 
 
 class OllamaTransientError(OllamaError):
-    """
-    Errors that are likely temporary (network issues, timeouts, 5xx, etc.).
-    Typically worth retrying.
-    """
+    """Errors that are likely temporary (network issues, timeouts, 5xx, etc.), typically worth retrying."""
 
 
 class OllamaPermanentError(OllamaError):
-    """
-    Errors that are unlikely to succeed on retry without changing inputs
-    or configuration (4xx, bad response schema, etc.).
-    """
+    """Errors that are unlikely to succeed on retry without changing inputs or configuration (4xx, bad response schema, etc.)."""

databao_context_engine/llm/install.py

@@ -49,14 +49,15 @@ ARTIFACTS: dict[str, ArtifactInfo] = {
 
 
 def resolve_ollama_bin() -> str:
-    """
-    Decide which `ollama` binary to use, in this order:
+    """Decide which `ollama` binary to use.
 
+    Here is the priority order:
     1. DCE_OLLAMA_BIN env var, if set and exists
     2. `ollama` found on PATH
     3. Managed installation under MANAGED_OLLAMA_BIN
 
-    Returns the full path to the binary
+    Returns:
+        The full path to the binary
     """
     override = os.environ.get("DCE_OLLAMA_BIN")
     if override:
@@ -76,9 +77,7 @@ def resolve_ollama_bin() -> str:
 
 
 def _detect_platform() -> str:
-    """
-    Return one of: 'darwin', 'linux-amd64', 'linux-arm64', 'windows-amd64', 'windows-arm64'.
-    """
+    """Return one of: 'darwin', 'linux-amd64', 'linux-arm64', 'windows-amd64', 'windows-arm64'."""
     os_name = sys.platform.lower()
     arch = (os.uname().machine if hasattr(os, "uname") else "").lower()
 
@@ -97,9 +96,7 @@ def _detect_platform() -> str:
 
 
 def _download_to_temp(url: str) -> Path:
-    """
-    Download to a temporary file and return its path.
-    """
+    """Download to a temporary file and return its path."""
     import urllib.request
 
     tmp_dir = Path(tempfile.mkdtemp(prefix="ollama-download-"))
@@ -114,9 +111,7 @@ def _download_to_temp(url: str) -> Path:
 
 
 def _verify_sha256(path: Path, expected_hex: str) -> None:
-    """
-    Verify SHA-256 of path matches expected_hex
-    """
+    """Verify SHA-256 of path matches expected_hex."""
     h = hashlib.sha256()
     with path.open("rb") as f:
         for chunk in iter(lambda: f.read(8192), b""):
@@ -127,9 +122,7 @@ def _verify_sha256(path: Path, expected_hex: str) -> None:
 
 
 def _extract_archive(archive: Path, target_dir: Path) -> None:
-    """
-    Extract archive into target_dir.
-    """
+    """Extract archive into target_dir."""
     name = archive.name.lower()
     target_dir.mkdir(parents=True, exist_ok=True)
 
@@ -144,9 +137,7 @@ def _extract_archive(archive: Path, target_dir: Path) -> None:
 
 
 def _ensure_executable(path: Path) -> None:
-    """
-    Mark path as executable
-    """
+    """Mark path as executable."""
     try:
         mode = path.stat().st_mode
         path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
@@ -155,8 +146,7 @@ def _ensure_executable(path: Path) -> None:
 
 
 def install_ollama_to(target: Path) -> None:
-    """
-    Ensure an Ollama binary exists.
+    """Ensure an Ollama binary exist.
 
     If it doesn't exist, this will:
     - detect OS
@@ -164,6 +154,9 @@ def install_ollama_to(target: Path) -> None:
     - verify its SHA-256 checksum
     - extract into the installation directory
     - make the binary executable
+
+    Raises:
+        RuntimeError: If the user's platform is not supported
     """
     target = target.expanduser()
     if target.parent.name == "bin":

databao_context_engine/llm/service.py

@@ -36,9 +36,7 @@ class OllamaService:
         return [float(x) for x in vec]
 
     def describe(self, *, model: str, text: str, context: str) -> str:
-        """
-        Ask Ollama to generate a short description for `text`
-        """
+        """Ask Ollama to generate a short description for `text`."""
         prompt = self._build_description_prompt(text=text, context=context)
 
         payload: dict[str, Any] = {"model": model, "prompt": prompt, "stream": False, "options": {"temperature": 0.1}}

databao_context_engine/mcp/mcp_runner.py

@@ -2,15 +2,17 @@ import logging
 from pathlib import Path
 
 from databao_context_engine.mcp.mcp_server import McpServer, McpTransport
+from databao_context_engine.project.layout import ensure_project_dir
 
 logger = logging.getLogger(__name__)
 
 
 def run_mcp_server(
     project_dir: Path,
-    run_name: str | None,
     transport: McpTransport,
     host: str | None = None,
     port: int | None = None,
 ) -> None:
-    McpServer(project_dir, run_name, host, port).run(transport)
+    ensure_project_dir(project_dir=project_dir)
+
+    McpServer(project_dir, host, port).run(transport)

databao_context_engine/mcp/mcp_server.py

@@ -1,5 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
+from datetime import date
 from pathlib import Path
 from typing import Literal
 
@@ -7,8 +8,6 @@ from mcp.server import FastMCP
 from mcp.types import ToolAnnotations
 
 from databao_context_engine import DatabaoContextEngine
-from databao_context_engine.mcp.all_results_tool import run_all_results_tool
-from databao_context_engine.mcp.retrieve_tool import run_retrieve_tool
 
 logger = logging.getLogger(__name__)
 
@@ -26,12 +25,10 @@ class McpServer:
     def __init__(
         self,
         project_dir: Path,
-        run_name: str | None,
         host: str | None = None,
         port: int | None = None,
     ):
         self._databao_context_engine = DatabaoContextEngine(project_dir)
-        self._run_name = run_name
 
         self._mcp_server = self._create_mcp_server(host, port)
 
@@ -43,20 +40,23 @@ class McpServer:
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def all_results_tool():
-            return run_all_results_tool(self._databao_context_engine, self._run_name)
+            return self._databao_context_engine.get_all_contexts_formatted()
 
         @mcp.tool(
             description="Retrieve the context built from various resources, including databases, dbt tools, plain and structured files, to retrieve relevant information",
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def retrieve_tool(text: str, limit: int | None):
-            return run_retrieve_tool(
-                databao_context_engine=self._databao_context_engine,
-                run_name=self._run_name,
-                text=text,
-                limit=limit or 50,
+            retrieve_results = self._databao_context_engine.search_context(
+                retrieve_text=text, limit=limit, export_to_file=False
             )
 
+            display_results = [context_search_result.context_result for context_search_result in retrieve_results]
+
+            display_results.append(f"\nToday's date is {date.today()}")
+
+            return "\n".join(display_results)
+
         return mcp
 
     def run(self, transport: McpTransport):

databao_context_engine/plugin_loader.py

@@ -0,0 +1,111 @@
+from databao_context_engine.introspection.property_extract import get_property_list_from_type
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    BuildFilePlugin,
+    BuildPlugin,
+    DatasourceType,
+)
+from databao_context_engine.pluginlib.config import ConfigPropertyDefinition, CustomiseConfigProperties
+from databao_context_engine.plugins.plugin_loader import (
+    load_plugins,
+)
+
+
+class DatabaoContextPluginLoader:
+    """Loader for plugins installed in the current environment."""
+
+    def __init__(self, plugins_by_type: dict[DatasourceType, BuildPlugin] | None = None):
+        """Initialize the DatabaoContextEngine.
+
+        Args:
+            plugins_by_type: Override the list of plugins loaded from the environment.
+                Typical usage should not provide this argument and leave it as None.
+        """
+        self._all_plugins_by_type = load_plugins() if plugins_by_type is None else plugins_by_type
+
+    def get_all_supported_datasource_types(self, exclude_file_plugins: bool = False) -> set[DatasourceType]:
+        """Return the list of all supported datasource types.
+
+        Args:
+            exclude_file_plugins: If True, do not return datasource types from plugins that deal with raw files.
+
+        Returns:
+            A set of all DatasourceType supported in the current installation environment.
+        """
+        if exclude_file_plugins:
+            return {
+                datasource_type
+                for (datasource_type, plugin) in self._all_plugins_by_type.items()
+                if not isinstance(plugin, BuildFilePlugin)
+            }
+        else:
+            return set(self._all_plugins_by_type.keys())
+
+    def get_plugin_for_datasource_type(self, datasource_type: DatasourceType) -> BuildPlugin:
+        """Return the plugin able to build a context for the given datasource type.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the plugin.
+
+        Returns:
+            The plugin able to build a context for the given datasource type.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+        """
+        if datasource_type not in self._all_plugins_by_type:
+            raise ValueError(f"No plugin found for type '{datasource_type.full_type}'")
+
+        return self._all_plugins_by_type[datasource_type]
+
+    def get_config_file_type_for_datasource_type(self, datasource_type: DatasourceType) -> type:
+        """Return the type of the config file for the given datasource type.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the config file type.
+
+        Returns:
+            The type of the config file for the given datasource type.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+            ValueError: If the plugin does not support config files.
+        """
+        plugin = self.get_plugin_for_datasource_type(datasource_type)
+
+        if isinstance(plugin, BuildDatasourcePlugin):
+            return plugin.config_file_type
+
+        raise ValueError(
+            f'Impossible to get a config file type for datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin'
+        )
+
+    def get_config_file_structure_for_datasource_type(
+        self, datasource_type: DatasourceType
+    ) -> list[ConfigPropertyDefinition]:
+        """Return the property structure of the config file for the given datasource type.
+
+        This can be used to generate a form for the user to fill in the config file.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the config file structure.
+
+        Returns:
+            The structure of the config file for the given datasource type.
+                This structure is a list of ConfigPropertyDefinition objects.
+                Each object in the list describes a property of the config file and its potential nested properties.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+            ValueError: If the plugin does not support config files.
+        """
+        plugin = self.get_plugin_for_datasource_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 datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin or CustomiseConfigProperties'
+            )