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

databao_context_engine/build_sources/internal/build_service.py

@@ -1,77 +0,0 @@
-from __future__ import annotations
-
-import logging
-from datetime import datetime
-
-from databao_context_engine.build_sources.internal.plugin_execution import BuiltDatasourceContext, execute
-from databao_context_engine.pluginlib.build_plugin import (
-    BuildPlugin,
-)
-from databao_context_engine.project.types import PreparedDatasource
-from databao_context_engine.serialisation.yaml import to_yaml_string
-from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingService
-from databao_context_engine.storage.models import RunDTO
-from databao_context_engine.storage.repositories.datasource_run_repository import DatasourceRunRepository
-from databao_context_engine.storage.repositories.run_repository import RunRepository
-
-logger = logging.getLogger(__name__)
-
-
-class BuildService:
-    def __init__(
-        self,
-        *,
-        run_repo: RunRepository,
-        datasource_run_repo: DatasourceRunRepository,
-        chunk_embedding_service: ChunkEmbeddingService,
-    ) -> None:
-        self._run_repo = run_repo
-        self._datasource_run_repo = datasource_run_repo
-        self._chunk_embedding_service = chunk_embedding_service
-
-    def start_run(self, *, project_id: str, dce_version: str) -> RunDTO:
-        """
-        Create a new run row and return (run_id, started_at).
-        """
-        return self._run_repo.create(project_id=project_id, dce_version=dce_version)
-
-    def finalize_run(self, *, run_id: int):
-        """
-        Mark the run as complete (sets ended_at).
-        """
-        self._run_repo.update(run_id=run_id, ended_at=datetime.now())
-
-    def process_prepared_source(
-        self,
-        *,
-        run_id: int,
-        prepared_source: PreparedDatasource,
-        plugin: BuildPlugin,
-    ) -> BuiltDatasourceContext:
-        """
-        Process a single source.
-
-        1) Execute the plugin
-        2) Divide the results into chunks
-        3) Embed and persist the chunks
-        """
-        result = execute(prepared_source, plugin)
-
-        chunks = plugin.divide_context_into_chunks(result.context)
-        if not chunks:
-            logger.info("No chunks for %s — skipping.", prepared_source.path.name)
-            return result
-
-        datasource_run = self._datasource_run_repo.create(
-            run_id=run_id,
-            plugin=plugin.name,
-            full_type=prepared_source.datasource_type.full_type,
-            source_id=result.datasource_id,
-            storage_directory=str(prepared_source.path.parent),
-        )
-
-        self._chunk_embedding_service.embed_chunks(
-            datasource_run_id=datasource_run.datasource_run_id, chunks=chunks, result=to_yaml_string(result.context)
-        )
-
-        return result

databao_context_engine/build_sources/internal/build_wiring.py

@@ -1,52 +0,0 @@
-import logging
-from pathlib import Path
-
-from databao_context_engine.build_sources.internal.build_runner import BuildContextResult, build
-from databao_context_engine.llm.factory import (
-    create_ollama_description_provider,
-    create_ollama_embedding_provider,
-    create_ollama_service,
-)
-from databao_context_engine.project.info import get_dce_version
-from databao_context_engine.project.layout import ensure_project_dir
-from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
-from databao_context_engine.services.factories import (
-    create_build_service,
-)
-from databao_context_engine.storage.connection import open_duckdb_connection
-from databao_context_engine.system.properties import get_db_path
-
-logger = logging.getLogger(__name__)
-
-
-def build_all_datasources(project_dir: Path, chunk_embedding_mode: ChunkEmbeddingMode) -> list[BuildContextResult]:
-    """
-    Public build entrypoint
-    - Instantiates the build service
-    - Delegates the actual build logic to the build runner
-    """
-    project_layout = ensure_project_dir(project_dir)
-
-    logger.debug(f"Starting to build datasources in project {project_dir.resolve()}")
-
-    with open_duckdb_connection(get_db_path()) as conn:
-        ollama_service = create_ollama_service()
-        embedding_provider = create_ollama_embedding_provider(ollama_service)
-        description_provider = (
-            create_ollama_description_provider(ollama_service)
-            if chunk_embedding_mode.should_generate_description()
-            else None
-        )
-        build_service = create_build_service(
-            conn,
-            embedding_provider=embedding_provider,
-            description_provider=description_provider,
-            chunk_embedding_mode=chunk_embedding_mode,
-        )
-        dce_config = project_layout.read_config_file()
-        return build(
-            project_dir=project_dir,
-            build_service=build_service,
-            project_id=str(dce_config.project_id),
-            dce_version=get_dce_version(),
-        )

databao_context_engine/build_sources/internal/export_results.py

@@ -1,43 +0,0 @@
-import logging
-from pathlib import Path
-
-from databao_context_engine.build_sources.internal.plugin_execution import BuiltDatasourceContext
-from databao_context_engine.datasource_config.datasource_context import get_context_header_for_datasource
-from databao_context_engine.project.layout import ALL_RESULTS_FILE_NAME, get_output_dir
-from databao_context_engine.project.types import DatasourceId
-from databao_context_engine.serialisation.yaml import write_yaml_to_stream
-
-logger = logging.getLogger(__name__)
-
-
-def create_run_dir(project_dir: Path, run_name: str) -> Path:
-    output_dir = get_output_dir(project_dir)
-
-    run_dir = output_dir.joinpath(run_name)
-    run_dir.mkdir(parents=True, exist_ok=False)
-
-    return run_dir
-
-
-def export_build_result(run_dir: Path, result: BuiltDatasourceContext) -> Path:
-    datasource_id = DatasourceId.from_string_repr(result.datasource_id)
-    export_file_path = run_dir.joinpath(datasource_id.relative_path_to_context_file())
-
-    # Make sure the parent folder exists
-    export_file_path.parent.mkdir(exist_ok=True)
-
-    with export_file_path.open("w") as export_file:
-        write_yaml_to_stream(data=result, file_stream=export_file)
-
-    logger.info(f"Exported result to {export_file_path.resolve()}")
-
-    return export_file_path
-
-
-def append_result_to_all_results(run_dir: Path, result: BuiltDatasourceContext):
-    path = run_dir.joinpath(ALL_RESULTS_FILE_NAME)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("a", encoding="utf-8") as export_file:
-        export_file.write(get_context_header_for_datasource(DatasourceId.from_string_repr(result.datasource_id)))
-        write_yaml_to_stream(data=result, file_stream=export_file)
-        export_file.write("\n")

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/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

databao_context_engine/retrieve_embeddings/internal/retrieve_wiring.py

@@ -1,29 +0,0 @@
-from databao_context_engine.llm.factory import create_ollama_embedding_provider, create_ollama_service
-from databao_context_engine.project.layout import ProjectLayout
-from databao_context_engine.retrieve_embeddings.internal.retrieve_runner import retrieve
-from databao_context_engine.services.factories import create_retrieve_service
-from databao_context_engine.storage.connection import open_duckdb_connection
-from databao_context_engine.storage.repositories.vector_search_repository import VectorSearchResult
-from databao_context_engine.system.properties import get_db_path
-
-
-def retrieve_embeddings(
-    project_layout: ProjectLayout,
-    retrieve_text: str,
-    run_name: str | None,
-    limit: int | None,
-    export_to_file: bool,
-) -> list[VectorSearchResult]:
-    with open_duckdb_connection(get_db_path()) as conn:
-        ollama_service = create_ollama_service()
-        embedding_provider = create_ollama_embedding_provider(ollama_service)
-        retrieve_service = create_retrieve_service(conn, embedding_provider=embedding_provider)
-        return retrieve(
-            project_dir=project_layout.project_dir,
-            retrieve_service=retrieve_service,
-            project_id=str(project_layout.read_config_file().project_id),
-            text=retrieve_text,
-            run_name=run_name,
-            limit=limit,
-            export_to_file=export_to_file,
-        )

databao_context_engine/retrieve_embeddings/public/api.py

@@ -1,3 +0,0 @@
-from databao_context_engine.retrieve_embeddings.internal.retrieve_wiring import retrieve_embeddings
-
-__all__ = ["retrieve_embeddings"]

databao_context_engine/services/run_name_policy.py

@@ -1,8 +0,0 @@
-from datetime import datetime
-
-
-class RunNamePolicy:
-    _RUN_DIR_PREFIX = "run-"
-
-    def build(self, *, run_started_at: datetime):
-        return f"{RunNamePolicy._RUN_DIR_PREFIX}{run_started_at.isoformat(timespec='seconds')}"

databao_context_engine/storage/repositories/datasource_run_repository.py

@@ -1,136 +0,0 @@
-from typing import Any, Optional, Tuple
-
-import duckdb
-from _duckdb import ConstraintException
-
-from databao_context_engine.storage.exceptions.exceptions import IntegrityError
-from databao_context_engine.storage.models import DatasourceRunDTO
-
-
-class DatasourceRunRepository:
-    def __init__(self, conn: duckdb.DuckDBPyConnection):
-        self._conn = conn
-
-    def create(
-        self,
-        *,
-        run_id: int,
-        plugin: str,
-        full_type: str,
-        source_id: str,
-        storage_directory: str,
-    ) -> DatasourceRunDTO:
-        try:
-            row = self._conn.execute(
-                """
-            INSERT INTO
-                datasource_run(run_id, plugin, full_type, source_id, storage_directory)
-            VALUES
-                (?, ?, ?, ?, ?)
-            RETURNING
-                *
-            """,
-                [run_id, plugin, full_type, source_id, storage_directory],
-            ).fetchone()
-            if row is None:
-                raise RuntimeError("datasource_run creation returned no object")
-            return self._row_to_dto(row)
-        except ConstraintException as e:
-            raise IntegrityError from e
-
-    def get(self, datasource_run_id: int) -> Optional[DatasourceRunDTO]:
-        row = self._conn.execute(
-            """
-        SELECT
-            *
-        FROM
-            datasource_run 
-        WHERE
-            datasource_run_id = ?
-        """,
-            [datasource_run_id],
-        ).fetchone()
-        return self._row_to_dto(row) if row else None
-
-    def update(
-        self,
-        datasource_run_id: int,
-        *,
-        plugin: Optional[str] = None,
-        full_type: Optional[str] = None,
-        source_id: Optional[str] = None,
-        storage_directory: Optional[str] = None,
-    ) -> Optional[DatasourceRunDTO]:
-        sets: list[Any] = []
-        params: list[Any] = []
-
-        if plugin is not None:
-            sets.append("plugin = ?")
-            params.append(plugin)
-        if full_type is not None:
-            sets.append("full_type = ?")
-            params.append(full_type)
-        if source_id is not None:
-            sets.append("source_id = ?")
-            params.append(source_id)
-        if storage_directory is not None:
-            sets.append("storage_directory = ?")
-            params.append(storage_directory)
-
-        if not sets:
-            return self.get(datasource_run_id)
-
-        params.append(datasource_run_id)
-        self._conn.execute(
-            f"""
-            UPDATE 
-                datasource_run
-            SET 
-                {", ".join(sets)}
-            WHERE 
-                datasource_run_id = ?
-            """,
-            params,
-        )
-
-        return self.get(datasource_run_id)
-
-    def delete(self, datasource_run_id: int) -> int:
-        row = self._conn.execute(
-            """
-            DELETE FROM 
-                datasource_run 
-            WHERE 
-                datasource_run_id = ? 
-            RETURNING 
-                datasource_run_id
-            """,
-            [datasource_run_id],
-        ).fetchone()
-        return 1 if row else 0
-
-    def list(self) -> list[DatasourceRunDTO]:
-        rows = self._conn.execute(
-            """
-            SELECT
-              *
-            FROM 
-                datasource_run
-            ORDER BY
-                datasource_run_id DESC
-            """
-        ).fetchall()
-        return [self._row_to_dto(r) for r in rows]
-
-    @staticmethod
-    def _row_to_dto(row: Tuple) -> DatasourceRunDTO:
-        datasource_run_id, run_id, plugin, source_id, storage_directory, created_at, full_type = row
-        return DatasourceRunDTO(
-            datasource_run_id=int(datasource_run_id),
-            run_id=int(run_id),
-            plugin=str(plugin),
-            full_type=str(full_type),
-            source_id=str(source_id),
-            storage_directory=str(storage_directory),
-            created_at=created_at,
-        )