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

databao_context_engine/project/types.py

@@ -1,134 +0,0 @@
-import os
-from dataclasses import dataclass
-from enum import StrEnum
-from pathlib import Path
-from typing import Any
-
-from databao_context_engine.pluginlib.build_plugin import DatasourceType
-
-
-class DatasourceKind(StrEnum):
-    CONFIG = "config"
-    FILE = "file"
-
-
-@dataclass(frozen=True)
-class DatasourceDescriptor:
-    path: Path
-    kind: DatasourceKind
-    main_type: str
-
-
-@dataclass(frozen=True)
-class PreparedConfig:
-    datasource_type: DatasourceType
-    path: Path
-    config: dict[Any, Any]
-    datasource_name: str
-
-
-@dataclass(frozen=True)
-class PreparedFile:
-    datasource_type: DatasourceType
-    path: Path
-
-
-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.
-
-    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.
-    """
-
-    datasource_config_folder: str
-    datasource_name: str
-    config_file_suffix: str
-
-    def __post_init__(self):
-        if not self.datasource_config_folder.strip():
-            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_config_folder must not be empty")
-        if not self.datasource_name.strip():
-            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_name must not be empty")
-        if not self.config_file_suffix.strip():
-            raise ValueError(f"Invalid DatasourceId ({str(self)}): config_file_suffix must not be empty")
-
-        if os.sep in self.datasource_config_folder:
-            raise ValueError(
-                f"Invalid DatasourceId ({str(self)}): datasource_config_folder must not contain a path separator"
-            )
-
-        if os.sep in self.datasource_name:
-            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_name must not contain a path separator")
-
-        if not self.config_file_suffix.startswith("."):
-            raise ValueError(
-                f'Invalid DatasourceId ({str(self)}): config_file_suffix must start with a dot "." (e.g.: .yaml)'
-            )
-
-        if self.datasource_name.endswith(self.config_file_suffix):
-            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_name must not contain the file suffix")
-
-    def __str__(self):
-        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.
-
-        The returned path is 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.
-
-        The returned path is relative to an output run 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")
-
-        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.
-
-        The string representation of a DatasourceId is the path to the datasource's config file relative to the src folder in the project.
-
-        e.g: "databases/my_postgres_datasource.yaml"
-        """
-        config_file_path = Path(datasource_id_as_string)
-
-        if len(config_file_path.parents) > 2:
-            raise ValueError(
-                f"Invalid string representation of a DatasourceId: too many parent folders defined in {datasource_id_as_string}"
-            )
-
-        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.
-
-        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.
-        """
-        return DatasourceId(
-            datasource_config_folder=datasource_config_file.parent.name,
-            datasource_name=datasource_config_file.stem,
-            config_file_suffix=datasource_config_file.suffix,
-        )
-
-
-@dataclass
-class Datasource:
-    id: DatasourceId
-    type: DatasourceType

databao_context_engine/retrieve_embeddings/internal/export_results.py

@@ -1,12 +0,0 @@
-from pathlib import Path
-
-
-def export_retrieve_results(run_dir: Path, retrieve_results: list[str]) -> Path:
-    path = run_dir.joinpath("context_duckdb.yaml")
-
-    with path.open("w") as export_file:
-        for result in retrieve_results:
-            export_file.write(result)
-            export_file.write("\n")
-
-    return path

databao_context_engine/retrieve_embeddings/internal/retrieve_runner.py

@@ -1,34 +0,0 @@
-import logging
-from pathlib import Path
-
-from databao_context_engine.project.runs import get_run_dir
-from databao_context_engine.retrieve_embeddings.internal.export_results import export_retrieve_results
-from databao_context_engine.retrieve_embeddings.internal.retrieve_service import RetrieveService
-from databao_context_engine.storage.repositories.vector_search_repository import VectorSearchResult
-
-logger = logging.getLogger(__name__)
-
-
-def retrieve(
-    project_dir: Path,
-    *,
-    retrieve_service: RetrieveService,
-    project_id: str,
-    text: str,
-    run_name: str | None,
-    limit: int | None,
-    export_to_file: bool,
-) -> list[VectorSearchResult]:
-    resolved_run_name = retrieve_service.resolve_run_name(project_id=project_id, run_name=run_name)
-    retrieve_results = retrieve_service.retrieve(
-        project_id=project_id, text=text, run_name=resolved_run_name, limit=limit
-    )
-
-    if export_to_file:
-        export_directory = get_run_dir(project_dir=project_dir, run_name=resolved_run_name)
-
-        display_texts = [result.display_text for result in retrieve_results]
-        export_file = export_retrieve_results(export_directory, display_texts)
-        logger.info(f"Exported results to {export_file}")
-
-    return retrieve_results

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,
-        )

databao_context_engine/storage/repositories/run_repository.py

@@ -1,157 +0,0 @@
-from datetime import datetime
-from typing import Any, Optional
-
-import duckdb
-
-from databao_context_engine.services.run_name_policy import RunNamePolicy
-from databao_context_engine.storage.models import RunDTO
-
-
-class RunRepository:
-    def __init__(self, conn: duckdb.DuckDBPyConnection, run_name_policy: RunNamePolicy):
-        self._conn = conn
-        self._run_name_policy = run_name_policy
-
-    def create(
-        self, *, project_id: str, dce_version: Optional[str] = None, started_at: datetime | None = None
-    ) -> RunDTO:
-        if started_at is None:
-            started_at = datetime.now()
-        run_name = self._run_name_policy.build(run_started_at=started_at)
-
-        row = self._conn.execute(
-            """
-            INSERT INTO 
-                run (project_id, nemory_version, started_at, run_name)
-            VALUES 
-                (?, ?, ?, ?)
-            RETURNING 
-                *
-            """,
-            [project_id, dce_version, started_at, run_name],
-        ).fetchone()
-        if row is None:
-            raise RuntimeError("Run creation returned no object")
-        return self._row_to_dto(row)
-
-    def get(self, run_id: int) -> Optional[RunDTO]:
-        row = self._conn.execute(
-            """
-            SELECT 
-                *
-            FROM
-                run
-            WHERE
-                run_id = ?
-            """,
-            [run_id],
-        ).fetchone()
-        return self._row_to_dto(row) if row else None
-
-    def get_by_run_name(self, *, project_id: str, run_name: str) -> RunDTO | None:
-        row = self._conn.execute(
-            """
-            SELECT
-                *
-            FROM
-                run
-            WHERE
-                run.project_id = ? AND run_name = ?
-            """,
-            [project_id, run_name],
-        ).fetchone()
-        return self._row_to_dto(row) if row else None
-
-    def get_latest_run_for_project(self, project_id: str) -> RunDTO | None:
-        row = self._conn.execute(
-            """
-            SELECT
-                *
-            FROM
-                run
-            WHERE
-                run.project_id = ?
-            ORDER BY run.started_at DESC
-            LIMIT 1
-            """,
-            [project_id],
-        ).fetchone()
-        return self._row_to_dto(row) if row else None
-
-    def update(
-        self,
-        run_id: int,
-        *,
-        project_id: Optional[str] = None,
-        ended_at: Optional[datetime] = None,
-        dce_version: Optional[str] = None,
-    ) -> Optional[RunDTO]:
-        sets: list[Any] = []
-        params: list[Any] = []
-
-        if project_id is not None:
-            sets.append("project_id = ?")
-            params.append(project_id)
-        if ended_at is not None:
-            sets.append("ended_at = ?")
-            params.append(ended_at)
-        if dce_version is not None:
-            sets.append("nemory_version = ?")
-            params.append(dce_version)
-
-        if not sets:
-            return self.get(run_id)
-
-        params.append(run_id)
-        self._conn.execute(
-            f"""
-                    UPDATE
-                        run
-                    SET 
-                        {", ".join(sets)}
-                    WHERE
-                        run_id = ?
-                    """,
-            params,
-        )
-
-        return self.get(run_id)
-
-    def delete(self, run_id: int) -> int:
-        row = self._conn.execute(
-            """
-            DELETE FROM
-                run
-            WHERE
-                run_id = ?
-            RETURNING
-                run_id
-            """,
-            [run_id],
-        ).fetchone()
-        return 1 if row else 0
-
-    def list(self) -> list[RunDTO]:
-        rows = self._conn.execute(
-            """
-            SELECT
-                *
-            FROM
-                run
-            ORDER BY
-                run_id DESC
-            """
-        ).fetchall()
-        return [self._row_to_dto(r) for r in rows]
-
-    @staticmethod
-    def _row_to_dto(row: tuple) -> RunDTO:
-        run_id, project_id, started_at, ended_at, dce_version, run_name = row
-        return RunDTO(
-            run_id=int(run_id),
-            run_name=run_name,
-            project_id=str(project_id),
-            started_at=started_at,
-            ended_at=ended_at,
-            nemory_version=dce_version,
-        )