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

databao_context_engine/build_sources/public/api.py

@@ -1,4 +0,0 @@
-from databao_context_engine.build_sources.internal.build_runner import BuildContextResult
-from databao_context_engine.build_sources.internal.build_wiring import build_all_datasources
-
-__all__ = ["build_all_datasources", "BuildContextResult"]

databao_context_engine/databao_engine.py

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

databao_context_engine/datasource_config/add_config.py

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

databao_context_engine/datasource_config/datasource_context.py

@@ -1,60 +0,0 @@
-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/mcp/all_results_tool.py

@@ -1,5 +0,0 @@
-from databao_context_engine import DatabaoContextEngine
-
-
-def run_all_results_tool(databao_context_engine: DatabaoContextEngine, run_name: str | None) -> str:
-    return databao_context_engine.get_all_contexts_formatted(run_name=run_name)

databao_context_engine/mcp/retrieve_tool.py

@@ -1,22 +0,0 @@
-import datetime
-
-from databao_context_engine import DatabaoContextEngine
-
-
-def run_retrieve_tool(
-    *, databao_context_engine: DatabaoContextEngine, run_name: str | None, text: str, limit: int | None = None
-) -> str:
-    """
-    Execute the retrieve flow for MCP and return the matching display texts
-    Adds the current date to the end
-    """
-
-    retrieve_results = databao_context_engine.search_context(
-        retrieve_text=text, run_name=run_name, 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 {datetime.date.today()}")
-
-    return "\n".join(display_results)

databao_context_engine/plugins/databases/athena_introspector.py

@@ -1,101 +0,0 @@
-from __future__ import annotations
-
-from typing import Any, Mapping
-
-from pyathena import connect
-from pyathena.cursor import DictCursor
-from pydantic import Field
-
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
-from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
-from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
-from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
-
-
-class AthenaConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/athena")
-    connection: dict[str, Any] = Field(
-        description="Connection parameters for the Athena database. It can contain any of the keys supported by the Athena connection library"
-    )
-
-
-class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
-    _IGNORED_SCHEMAS = {
-        "information_schema",
-    }
-    supports_catalogs = True
-
-    def _connect(self, file_config: AthenaConfigFile):
-        connection = file_config.connection
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
-
-        return connect(**connection, cursor_class=DictCursor)
-
-    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
-        with connection.cursor() as cur:
-            cur.execute(sql, params or {})
-            return cur.fetchall()
-
-    def _get_catalogs(self, connection, file_config: AthenaConfigFile) -> list[str]:
-        catalog = file_config.connection.get("catalog", self._resolve_pseudo_catalog_name(file_config))
-        return [catalog]
-
-    def _connect_to_catalog(self, file_config: AthenaConfigFile, catalog: str):
-        self._connect(file_config)
-
-    def _sql_list_schemas(self, catalogs: list[str] | None) -> SQLQuery:
-        if not catalogs:
-            return SQLQuery("SELECT schema_name, catalog_name FROM information_schema.schemata", None)
-        catalog = catalogs[0]
-        sql = "SELECT schema_name, catalog_name FROM information_schema.schemata WHERE catalog_name = %(catalog)s"
-        return SQLQuery(sql, {"catalog": catalog})
-
-    # TODO: Incomplete plugin. Awaiting permission access to AWS to properly develop
-    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
-        if not schemas:
-            return []
-
-        comps = {"columns": self._sql_columns(catalog, schemas)}
-        results: dict[str, list[dict]] = {}
-
-        for name, q in comps.items():
-            results[name] = self._fetchall_dicts(connection, q.sql, q.params)
-
-        return IntrospectionModelBuilder.build_schemas_from_components(
-            schemas=schemas,
-            rels=results.get("relations", []),
-            cols=results.get("columns", []),
-            pk_cols=[],
-            uq_cols=[],
-            checks=[],
-            fk_cols=[],
-            idx_cols=[],
-        )
-
-    def _sql_columns(self, catalog: str, schemas: list[str]) -> SQLQuery:
-        sql = f"""
-        SELECT 
-            table_schema AS schema_name,
-            table_name, 
-            column_name, 
-            ordinal_position, 
-            data_type,
-            is_nullable
-        FROM 
-            {catalog}.information_schema.columns
-        WHERE 
-            table_schema IN ({schemas})
-        ORDER BY
-            table_schema,
-            table_name,
-            ordinal_position
-        """
-        return SQLQuery(sql, {"schema": schemas})
-
-    def _resolve_pseudo_catalog_name(self, file_config: AthenaConfigFile) -> str:
-        return "awsdatacatalog"
-
-    def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
-        sql = f'SELECT * FROM "{schema}"."{table}" LIMIT %(limit)s'
-        return SQLQuery(sql, {"limit": limit})

databao_context_engine/plugins/duckdb_db_plugin.py

@@ -1,12 +0,0 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.duckdb_introspector import DuckDBConfigFile, DuckDBIntrospector
-
-
-class DuckDbPlugin(BaseDatabasePlugin[DuckDBConfigFile]):
-    id = "jetbrains/duckdb"
-    name = "DuckDB Plugin"
-    supported = {"databases/duckdb"}
-    config_file_type = DuckDBConfigFile
-
-    def __init__(self):
-        super().__init__(DuckDBIntrospector())

databao_context_engine/plugins/postgresql_db_plugin.py

@@ -1,12 +0,0 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.postgresql_introspector import PostgresConfigFile, PostgresqlIntrospector
-
-
-class PostgresqlDbPlugin(BaseDatabasePlugin[PostgresConfigFile]):
-    id = "jetbrains/postgres"
-    name = "PostgreSQL DB Plugin"
-    supported = {"databases/postgres"}
-    config_file_type = PostgresConfigFile
-
-    def __init__(self):
-        super().__init__(PostgresqlIntrospector())

databao_context_engine/plugins/snowflake_db_plugin.py

@@ -1,12 +0,0 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.snowflake_introspector import SnowflakeConfigFile, SnowflakeIntrospector
-
-
-class SnowflakeDbPlugin(BaseDatabasePlugin[SnowflakeConfigFile]):
-    id = "jetbrains/snowflake"
-    name = "Snowflake DB Plugin"
-    supported = {"databases/snowflake"}
-    config_file_type = SnowflakeConfigFile
-
-    def __init__(self):
-        super().__init__(SnowflakeIntrospector())

databao_context_engine/project/datasource_discovery.py

@@ -1,141 +0,0 @@
-import logging
-from pathlib import Path
-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 (
-    DatasourceDescriptor,
-    DatasourceKind,
-    PreparedConfig,
-    PreparedDatasource,
-    PreparedFile,
-    DatasourceId,
-    Datasource,
-)
-from databao_context_engine.templating.renderer import render_template
-
-logger = logging.getLogger(__name__)
-
-
-def get_datasource_list(project_dir: Path) -> list[Datasource]:
-    result = []
-    for discovered_datasource in discover_datasources(project_dir=project_dir):
-        try:
-            prepared_source = prepare_source(discovered_datasource)
-        except Exception as e:
-            logger.debug(str(e), exc_info=True, stack_info=True)
-            logger.info(f"Invalid source at ({discovered_datasource.path}): {str(e)}")
-            continue
-
-        result.append(
-            Datasource(
-                id=DatasourceId.from_datasource_config_file_path(discovered_datasource.path),
-                type=prepared_source.datasource_type,
-            )
-        )
-
-    return result
-
-
-def discover_datasources(project_dir: Path) -> list[DatasourceDescriptor]:
-    """
-    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
-    """
-    src = get_source_dir(project_dir)
-    if not src.exists() or not src.is_dir():
-        raise ValueError(f"src directory does not exist in {project_dir}")
-
-    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)
-            if datasource is not None:
-                datasources.append(datasource)
-
-    return datasources
-
-
-def _is_datasource_file(p: Path) -> bool:
-    # ignore backup files
-    return p.is_file() and not p.suffix.endswith("~")
-
-
-def get_datasource_descriptors(project_dir: Path, datasource_ids: list[DatasourceId]):
-    src = get_source_dir(project_dir)
-    if not src.exists() or not src.is_dir():
-        raise ValueError(f"src directory does not exist in {project_dir}")
-
-    datasources: list[DatasourceDescriptor] = []
-    for datasource_id in datasource_ids:
-        config_file_path = src.joinpath(datasource_id.relative_path_to_config_file())
-        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)
-        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
-    """
-    if not path.is_file():
-        return None
-
-    parent_name = path.parent.name
-    extension = path.suffix.lower().lstrip(".")
-
-    if parent_name == "files":
-        return DatasourceDescriptor(path=path.resolve(), main_type=parent_name, kind=DatasourceKind.FILE)
-
-    if extension in {"yaml", "yml"}:
-        return DatasourceDescriptor(path=path.resolve(), main_type=parent_name, kind=DatasourceKind.CONFIG)
-
-    if extension:
-        return DatasourceDescriptor(path=path.resolve(), main_type=parent_name, kind=DatasourceKind.FILE)
-
-    logger.debug("Skipping file without extension: %s", path)
-    return None
-
-
-def prepare_source(datasource: DatasourceDescriptor) -> PreparedDatasource:
-    """
-    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(
-            datasource_type=DatasourceType.from_main_and_subtypes(main_type=datasource.main_type, subtype=file_subtype),
-            path=datasource.path,
-        )
-
-    else:
-        config = _parse_config_file(datasource.path)
-
-        subtype = config.get("type")
-        if not subtype or not isinstance(subtype, str):
-            raise ValueError("Config missing 'type' at %s - skipping", datasource.path)
-
-        return PreparedConfig(
-            datasource_type=DatasourceType.from_main_and_subtypes(main_type=datasource.main_type, subtype=subtype),
-            path=datasource.path,
-            config=config,
-            datasource_name=datasource.path.stem,
-        )
-
-
-def _parse_config_file(file_path: Path) -> dict[Any, Any]:
-    rendered_file = render_template(file_path.read_text())
-
-    return yaml.safe_load(rendered_file) or {}

databao_context_engine/project/runs.py

@@ -1,39 +0,0 @@
-from pathlib import Path
-
-from databao_context_engine.project.layout import ProjectLayout, get_output_dir
-from databao_context_engine.storage.connection import open_duckdb_connection
-from databao_context_engine.storage.repositories.factories import create_run_repository
-from databao_context_engine.storage.repositories.run_repository import RunRepository
-from databao_context_engine.system.properties import get_db_path
-
-
-def resolve_run_name(*, project_layout: ProjectLayout, run_name: str | None) -> str:
-    project_id = str(project_layout.read_config_file().project_id)
-
-    with open_duckdb_connection(get_db_path()) as conn:
-        run_repository = create_run_repository(conn)
-
-        return resolve_run_name_from_repo(run_repository=run_repository, project_id=project_id, run_name=run_name)
-
-
-def resolve_run_name_from_repo(*, run_repository: RunRepository, project_id: str, run_name: str | None) -> str:
-    if run_name is None:
-        latest = run_repository.get_latest_run_for_project(project_id=project_id)
-        if latest is None:
-            raise LookupError(f"No runs found for project '{project_id}'. Run a build first.")
-        return latest.run_name
-    else:
-        run = run_repository.get_by_run_name(project_id=project_id, run_name=run_name)
-        if run is None:
-            raise LookupError(f"Run '{run_name}' not found for project '{project_id}'.")
-        return run.run_name
-
-
-def get_run_dir(project_dir: Path, run_name: str) -> Path:
-    run_dir = get_output_dir(project_dir).joinpath(run_name)
-    if not run_dir.is_dir():
-        raise ValueError(
-            f"The run with name {run_name} doesn't exist in the project. [project_dir: {project_dir.resolve()}]"
-        )
-
-    return run_dir