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

databao_context_engine/datasource_config/check_config.py

@@ -0,0 +1,131 @@
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+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 (
+    discover_datasources,
+    get_datasource_descriptors,
+    prepare_source,
+)
+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):
+    VALID = "Valid"
+    INVALID = "Invalid"
+    UNKNOWN = "Unknown"
+
+
+@dataclass(kw_only=True)
+class CheckDatasourceConnectionResult:
+    datasource_id: DatasourceId
+    connection_status: DatasourceConnectionStatus
+    summary: str | None
+    full_message: str | None = None
+
+    def format(self, show_summary_only: bool = True) -> str:
+        formatted_string = str(self.connection_status.value)
+        if self.summary:
+            formatted_string += f" - {self.summary}"
+        if not show_summary_only and self.full_message:
+            formatted_string += f"{os.linesep}{self.full_message}"
+
+        return formatted_string
+
+
+def check_datasource_connection(
+    project_dir: Path, *, datasource_ids: list[DatasourceId] | None = None
+) -> dict[DatasourceId, CheckDatasourceConnectionResult]:
+    ensure_project_dir(project_dir)
+
+    if datasource_ids:
+        logger.info(f"Validating datasource(s): {datasource_ids}")
+        datasources_to_traverse = get_datasource_descriptors(project_dir, datasource_ids)
+    else:
+        datasources_to_traverse = discover_datasources(project_dir)
+
+    plugins = load_plugins(exclude_file_plugins=True)
+
+    result = {}
+    for discovered_datasource in datasources_to_traverse:
+        result_key = DatasourceId.from_datasource_config_file_path(discovered_datasource.path)
+
+        try:
+            prepared_source = prepare_source(discovered_datasource)
+        except Exception as e:
+            result[result_key] = CheckDatasourceConnectionResult(
+                datasource_id=result_key,
+                connection_status=DatasourceConnectionStatus.INVALID,
+                summary="Failed to prepare source",
+                full_message=str(e),
+            )
+            continue
+
+        plugin = plugins.get(prepared_source.datasource_type)
+        if plugin is None:
+            logger.debug(
+                "No plugin for '%s' (datasource=%s) — skipping.",
+                prepared_source.datasource_type.full_type,
+                prepared_source.path,
+            )
+            result[result_key] = CheckDatasourceConnectionResult(
+                datasource_id=result_key,
+                connection_status=DatasourceConnectionStatus.INVALID,
+                summary="No compatible plugin found",
+            )
+            continue
+
+        if isinstance(prepared_source, PreparedConfig) and isinstance(plugin, BuildDatasourcePlugin):
+            try:
+                check_connection_for_datasource(
+                    plugin=plugin,
+                    datasource_type=prepared_source.datasource_type,
+                    config=prepared_source.config,
+                    datasource_name=prepared_source.datasource_name,
+                )
+
+                result[result_key] = CheckDatasourceConnectionResult(
+                    datasource_id=result_key, connection_status=DatasourceConnectionStatus.VALID, summary=None
+                )
+            except Exception as e:
+                logger.debug(
+                    f"Connection failed for {prepared_source.datasource_name} with error: {str(e)}",
+                    exc_info=True,
+                    stack_info=True,
+                )
+                result[result_key] = get_validation_result_from_error(result_key, e)
+
+    return result
+
+
+def get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
+    if isinstance(e, ValidationError):
+        return CheckDatasourceConnectionResult(
+            datasource_id=datasource_id,
+            connection_status=DatasourceConnectionStatus.INVALID,
+            summary="Config file is invalid",
+            full_message=str(e),
+        )
+    elif isinstance(e, NotImplementedError | NotSupportedError):
+        return CheckDatasourceConnectionResult(
+            datasource_id=datasource_id,
+            connection_status=DatasourceConnectionStatus.UNKNOWN,
+            summary="Plugin doesn't support validating its config",
+        )
+    else:
+        return CheckDatasourceConnectionResult(
+            datasource_id=datasource_id,
+            connection_status=DatasourceConnectionStatus.INVALID,
+            summary="Connection with the datasource can not be established",
+            full_message=str(e),
+        )

databao_context_engine/datasource_config/datasource_context.py

@@ -0,0 +1,60 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from databao_context_engine.project.layout import ProjectLayout
+from databao_context_engine.project.runs import get_run_dir, resolve_run_name
+from databao_context_engine.project.types import DatasourceId
+
+
+@dataclass(eq=True, frozen=True)
+class DatasourceContext:
+    datasource_id: DatasourceId
+    # TODO: Read the context as a BuildExecutionResult instead of a Yaml string?
+    context: str
+
+
+def get_datasource_context(
+    project_layout: ProjectLayout, datasource_id: DatasourceId, run_name: str | None = None
+) -> DatasourceContext:
+    run_dir = _resolve_run_dir(project_layout, run_name)
+
+    context_path = run_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)} in run {run_dir.name}")
+
+    context = context_path.read_text()
+    return DatasourceContext(datasource_id=datasource_id, context=context)
+
+
+def get_all_contexts(project_layout: ProjectLayout, run_name: str | None = None) -> list[DatasourceContext]:
+    run_dir = _resolve_run_dir(project_layout, run_name)
+
+    result = []
+    for main_type_dir in sorted((p for p in run_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(
+                    # FIXME: The extension will always be yaml here even if the datasource is a file with a different extension
+                    datasource_id=DatasourceId.from_datasource_config_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}"
+
+
+def _resolve_run_dir(project_layout: ProjectLayout, run_name: str | None) -> Path:
+    resolved_run_name = resolve_run_name(project_layout=project_layout, run_name=run_name)
+
+    run_dir = get_run_dir(project_dir=project_layout.project_dir, run_name=resolved_run_name)
+    if not run_dir.is_dir():
+        raise ValueError(f"Run {resolved_run_name} does not exist at {run_dir.resolve()}")
+
+    return run_dir

databao_context_engine/event_journal/writer.py

@@ -0,0 +1,29 @@
+import datetime
+import json
+from pathlib import Path
+from uuid import UUID, uuid1
+
+from databao_context_engine.system.properties import get_dce_path
+
+
+def get_journal_file(dce_path: Path) -> Path:
+    return dce_path / "event-journal" / "journal.txt"
+
+
+def log_event(*, project_id: UUID, dce_version: str, event_type: str, event_id: UUID = uuid1(), **kwargs):
+    current_timestamp = datetime.datetime.now().isoformat()
+    journal_file = get_journal_file(get_dce_path())
+    journal_file.parent.mkdir(parents=True, exist_ok=True)
+    with journal_file.open("a+", encoding="utf-8") as journal_handle:
+        event_json = json.dumps(
+            {
+                "id": str(event_id),
+                "project_id": str(project_id),
+                "dce_version": dce_version,
+                "timestamp": current_timestamp,
+                "type": event_type,
+                **kwargs,
+            }
+        )
+        journal_handle.write(event_json)
+        journal_handle.write("\n")

databao_context_engine/generate_configs_schemas.py

@@ -0,0 +1,92 @@
+import logging
+import os
+from typing import cast
+
+import click
+
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin
+from databao_context_engine.pluginlib.plugin_utils import format_json_schema_for_output, generate_json_schema
+from databao_context_engine.plugins.plugin_loader import load_plugins
+
+logger = logging.getLogger(__name__)
+
+
+@click.command()
+@click.option(
+    "-i",
+    "--include-plugins",
+    multiple=True,
+    type=str,
+    help="""
+    Plugin ID for the plugins to include in the json schema generation (-i can be specified multiple times to include multiple plugins). 
+    When set, only the plugins in this list are included in the json schema generation.
+    This can't be used in conjunction with --exclude-plugins  
+    """,
+)
+@click.option(
+    "-e",
+    "--exclude-plugins",
+    multiple=True,
+    help="""
+    Plugin ID for the plugins to exclude from the json schema generation (-e can be specified multiple times to exclude multiple plugins). 
+    This can't be used in conjunction with --include-plugins  
+    """,
+)
+def generate_configs_schemas(
+    include_plugins: tuple[str, ...] | None = None, exclude_plugins: tuple[str, ...] | None = None
+):
+    try:
+        schema_list = _generate_json_schema_output_for_plugins(include_plugins, exclude_plugins)
+
+        print(f"{os.linesep}{os.linesep}".join(schema_list))
+    except Exception as e:
+        print(str(e))
+
+
+def _generate_json_schema_output_for_plugins(
+    include_plugins: tuple[str, ...] | None = None, exclude_plugins: tuple[str, ...] | None = None
+) -> list[str]:
+    if include_plugins and exclude_plugins:
+        raise ValueError("Can't use --include-plugins and --exclude-plugins together")
+
+    filtered_plugins = _get_plugins_for_schema_generation(include_plugins, exclude_plugins)
+
+    if len(filtered_plugins) == 0:
+        if include_plugins:
+            raise ValueError(f"No plugin found with id in {include_plugins}")
+        elif exclude_plugins:
+            raise ValueError(f"No plugin found when excluding {exclude_plugins}")
+        else:
+            raise ValueError("No plugin found")
+
+    results = []
+    for plugin in filtered_plugins:
+        if not isinstance(plugin, BuildDatasourcePlugin):
+            continue
+
+        json_schema = generate_json_schema(plugin)
+        if json_schema is not None:
+            results.append(format_json_schema_for_output(plugin, json_schema))
+
+    return results
+
+
+def _get_plugins_for_schema_generation(
+    include_plugins: tuple[str, ...] | None = None, exclude_plugins: tuple[str, ...] | None = None
+) -> list[BuildDatasourcePlugin]:
+    all_plugins = load_plugins(exclude_file_plugins=True).values()
+
+    return [
+        cast(BuildDatasourcePlugin, plugin)
+        for plugin in all_plugins
+        if (plugin.id in include_plugins if include_plugins else True)
+        and (plugin.id not in exclude_plugins if exclude_plugins else True)
+    ]
+
+
+def main():
+    generate_configs_schemas()
+
+
+if __name__ == "__main__":
+    main()

databao_context_engine/init_project.py

@@ -0,0 +1,18 @@
+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.layout import is_project_dir_valid
+
+
+def init_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
+    initialised_project_dir = init_project_dir(project_dir=project_dir)
+
+    return DatabaoContextProjectManager(project_dir=initialised_project_dir)
+
+
+def init_or_get_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
+    if is_project_dir_valid(project_dir):
+        return DatabaoContextProjectManager(project_dir=project_dir)
+
+    return init_dce_project(project_dir=project_dir)

databao_context_engine/introspection/property_extract.py

@@ -0,0 +1,202 @@
+import types
+from dataclasses import MISSING, fields, is_dataclass
+from typing import Annotated, Any, ForwardRef, Union, get_origin, get_type_hints
+
+from pydantic import BaseModel, _internal
+from pydantic_core import PydanticUndefinedType
+
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation, ConfigPropertyDefinition
+
+
+def get_property_list_from_type(root_type: type) -> list[ConfigPropertyDefinition]:
+    return _get_property_list_from_type(parent_type=root_type, is_root_type=True)
+
+
+def _get_property_list_from_type(*, parent_type: type, is_root_type: bool) -> list[ConfigPropertyDefinition]:
+    if is_dataclass(parent_type):
+        return _get_property_list_from_dataclass(parent_type=parent_type)
+
+    try:
+        if issubclass(parent_type, BaseModel):
+            return _get_property_list_from_pydantic_base_model(parent_type=parent_type)
+    except TypeError:
+        # when trying to compare ABC Metadata classes to BaseModel, e.g: issubclass(Mapping[str, str], BaseModel)
+        # issubclass is raising a TypeError: issubclass() arg 1 must be a class
+        pass
+
+    return _get_property_list_from_type_hints(parent_type=parent_type, is_root_type=is_root_type)
+
+
+def _get_property_list_from_type_hints(*, parent_type: type, is_root_type: bool) -> list[ConfigPropertyDefinition]:
+    try:
+        type_hints = get_type_hints(parent_type, include_extras=True)
+    except TypeError as e:
+        if is_root_type:
+            # Ignore root types that don't have type hints like dict or list
+            return []
+        else:
+            # If we're evaluating a nested property, we want to propagate the exception
+            # to let the parent property know that this type should be ignored
+            raise e
+
+    result = []
+    for property_key, property_type in type_hints.items():
+        config_property = _create_property(property_type=property_type, property_name=property_key)
+
+        if config_property is not None:
+            result.append(config_property)
+
+    return result
+
+
+def _get_property_list_from_dataclass(parent_type: type) -> list[ConfigPropertyDefinition]:
+    if not is_dataclass(parent_type):
+        raise ValueError(f"{parent_type} is not a dataclass")
+
+    dataclass_fields = fields(parent_type)
+
+    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
+
+        property_for_field = _create_property(
+            property_type=property_type,
+            property_name=field.name,
+            property_default=field.default if has_field_default else None,
+            is_property_required=not has_field_default,
+        )
+
+        if property_for_field is not None:
+            result.append(property_for_field)
+
+    return result
+
+
+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")
+
+    pydantic_fields = parent_type.model_fields
+    result = []
+
+    for field_name, field_info in pydantic_fields.items():
+        has_field_default = type(field_info.default) is not PydanticUndefinedType
+
+        if field_info.annotation is None:
+            # No type: ignore the field
+            continue
+
+        property_for_field = _create_property(
+            property_type=field_info.annotation,
+            property_name=field_name,
+            property_default=field_info.default if has_field_default else None,
+            is_property_required=not has_field_default,
+            annotation=next(
+                (metadata for metadata in field_info.metadata if isinstance(metadata, ConfigPropertyAnnotation)), None
+            ),
+        )
+
+        if property_for_field is not None:
+            result.append(property_for_field)
+
+    return result
+
+
+def _create_property(
+    *,
+    property_type: type,
+    property_name: str,
+    property_default: Any | None = None,
+    is_property_required: bool = False,
+    annotation: ConfigPropertyAnnotation | None = None,
+) -> ConfigPropertyDefinition | None:
+    annotation = annotation or _get_config_property_annotation(property_type)
+
+    if annotation is not None and annotation.ignored_for_config_wizard:
+        return None
+
+    actual_property_type = _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
+
+    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 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,
+    )
+
+
+def _get_config_property_annotation(property_type) -> ConfigPropertyAnnotation | None:
+    if get_origin(property_type) is Annotated:
+        return next(
+            (metadata for metadata in property_type.__metadata__ if isinstance(metadata, ConfigPropertyAnnotation)),
+            None,
+        )
+
+    return None
+
+
+def _read_actual_property_type(property_type: type) -> 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 property_type
+
+
+def compute_default_value(
+    *, annotation: ConfigPropertyAnnotation | None, property_default: Any | None = None, has_nested_properties: bool
+) -> str | None:
+    if has_nested_properties:
+        return None
+
+    if annotation is not None and annotation.default_value is not None:
+        return str(annotation.default_value)
+
+    if property_default is not None:
+        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/config.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class OllamaConfig:
+    host: str = "127.0.0.1"
+    port: int = 11434
+
+    timeout: float = 30.0
+    headers: Optional[dict[str, str]] = None
+
+    bin_path: str = "ollama"
+    work_dir: Optional[Path] = None
+    extra_env: Optional[dict[str, str]] = None
+
+    @property
+    def base_url(self) -> str:
+        return f"http://{self.host}:{self.port}"

databao_context_engine/llm/descriptions/ollama.py

@@ -0,0 +1,21 @@
+from databao_context_engine.llm.descriptions.provider import DescriptionProvider
+from databao_context_engine.llm.service import OllamaService
+
+
+class OllamaDescriptionProvider(DescriptionProvider):
+    def __init__(self, *, service: OllamaService, model_id: str):
+        self._service = service
+        self._model_id = model_id
+
+    @property
+    def describer(self) -> str:
+        return "ollama"
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    def describe(self, text: str, context: str) -> str:
+        description = self._service.describe(model=self._model_id, text=text, context=context)
+
+        return description

databao_context_engine/llm/descriptions/provider.py

@@ -0,0 +1,10 @@
+from typing import Protocol
+
+
+class DescriptionProvider(Protocol):
+    @property
+    def describer(self) -> str: ...
+    @property
+    def model_id(self) -> str: ...
+
+    def describe(self, text: str, context: str) -> str: ...

databao_context_engine/llm/embeddings/ollama.py

@@ -0,0 +1,37 @@
+from collections.abc import Sequence
+
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.llm.service import OllamaService
+
+
+class OllamaEmbeddingProvider(EmbeddingProvider):
+    def __init__(
+        self,
+        *,
+        service: OllamaService,
+        model_id: str,
+        dim: int = 768,
+    ):
+        self._service = service
+        self._model_id = model_id
+        self._dim: int = dim
+
+    @property
+    def embedder(self) -> str:
+        return "ollama"
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    def embed(self, text: str) -> Sequence[float]:
+        vec = self._service.embed(model=self._model_id, text=text)
+
+        if len(vec) != self._dim:
+            raise ValueError(f"provider returned dim={len(vec)} but expected {self._dim}")
+
+        return [float(x) for x in vec]

databao_context_engine/llm/embeddings/provider.py

@@ -0,0 +1,13 @@
+from collections.abc import Sequence
+from typing import Protocol
+
+
+class EmbeddingProvider(Protocol):
+    @property
+    def embedder(self) -> str: ...
+    @property
+    def model_id(self) -> str: ...
+    @property
+    def dim(self) -> int: ...
+
+    def embed(self, text: str) -> Sequence[float]: ...

databao_context_engine/llm/errors.py

@@ -0,0 +1,16 @@
+class OllamaError(Exception):
+    """Base class for Ollama-related errors."""
+
+
+class OllamaTransientError(OllamaError):
+    """
+    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.).
+    """