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

databao_context_engine/services/table_name_policy.py

@@ -0,0 +1,15 @@
+import re
+
+_TABLE_NAME_RE = re.compile(r"^embedding_[a-z0-9_]+$")
+
+
+class TableNamePolicy:
+    @staticmethod
+    def build(*, embedder: str, model_id: str, dim: int) -> str:
+        safe_model = model_id.replace(":", "_").replace("-", "_").replace(" ", "_").replace(".", "_")
+        return f"embedding_{embedder}__{safe_model}__{dim}"
+
+    @staticmethod
+    def validate_table_name(*, table_name: str):
+        if not _TABLE_NAME_RE.fullmatch(table_name):
+            raise ValueError(f"invalid table_name {table_name!r}; expected pattern {_TABLE_NAME_RE.pattern}")

databao_context_engine/storage/connection.py

@@ -0,0 +1,32 @@
+import logging
+from contextlib import contextmanager
+from pathlib import Path
+
+import duckdb
+
+from databao_context_engine.system.properties import get_db_path
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def open_duckdb_connection(db_path: str | Path | None = None):
+    """
+    Open a DuckDB connection with vector search enabled and close on exist.
+    Loads the vss extension and enables HNSW experimental persistence.
+
+    Usage:
+        with open_duckdb_connection() as conn:
+    """
+    path = str(db_path or get_db_path())
+    conn = duckdb.connect(path)
+    logger.debug(f"Connected to DuckDB database at {path}")
+
+    try:
+        conn.execute("LOAD vss;")
+        conn.execute("SET hnsw_enable_experimental_persistence = true;")
+
+        logger.debug("Loaded Vector Similarity Search extension")
+        yield conn
+    finally:
+        conn.close()

databao_context_engine/storage/exceptions/exceptions.py

@@ -0,0 +1,6 @@
+class RepositoryError(Exception):
+    """Base exception for repository errors"""
+
+
+class IntegrityError(RepositoryError):
+    """Raised when a DB constraint is violated"""

databao_context_engine/storage/migrate.py

@@ -0,0 +1,127 @@
+import hashlib
+import logging
+import re
+from dataclasses import dataclass
+from importlib.resources import files
+from pathlib import Path
+from typing import LiteralString
+
+import duckdb
+
+from databao_context_engine.system.properties import get_db_path
+
+logger = logging.getLogger(__name__)
+
+
+def migrate(db_path: str | Path | None = None, migration_files: list[Path] | None = None) -> None:
+    if migration_files is None:
+        migration_files = [
+            migration
+            for migration in files("databao_context_engine.storage.migrations").iterdir()
+            if isinstance(migration, Path) and ".sql" == migration.suffix
+        ]
+
+    db = Path(db_path or get_db_path()).expanduser().resolve()
+    db.parent.mkdir(parents=True, exist_ok=True)
+    logger.debug("Running migrations on database: %s", db)
+
+    migration_manager = _MigrationManager(db, migration_files)
+    migration_manager.migrate()
+    logger.debug("Migration complete")
+
+
+@dataclass(frozen=True)
+class MigrationDTO:
+    name: str
+    version: int
+    checksum: str
+
+
+class MigrationError(Exception):
+    """Base class for migration errors."""
+
+
+def load_migrations(conn) -> list[MigrationDTO]:
+    with conn.cursor() as cur:
+        cur.execute(
+            "SELECT name, version, checksum, applied_at FROM migration_history",
+        )
+        rows = cur.fetchall()
+        return [
+            MigrationDTO(name=name, version=version, checksum=checksum)
+            for (name, version, checksum, applied_at) in rows
+        ]
+
+
+def _extract_version_from_name(name: str) -> int:
+    version_groups = re.findall(r"(\d+)__", name)
+    if not version_groups:
+        raise ValueError(f"Invalid migration name: {name}")
+    return int(version_groups[0])
+
+
+@dataclass(frozen=True)
+class _Migration:
+    name: str
+    version: int
+    checksum: str
+    query: str
+
+
+def _create_migration(file: Path) -> _Migration:
+    query_bytes = file.read_bytes()
+    query = query_bytes.decode("utf-8")
+    checksum = hashlib.md5(query_bytes).hexdigest()
+    version = _extract_version_from_name(file.name)
+    return _Migration(name=file.name, version=version, checksum=checksum, query=query)
+
+
+class _MigrationManager:
+    _init_migration_table_sql: LiteralString = """
+        CREATE SEQUENCE IF NOT EXISTS migration_history_id_seq START 1;
+        
+        CREATE TABLE IF NOT EXISTS migration_history (
+            id              BIGINT PRIMARY KEY DEFAULT nextval('migration_history_id_seq'),
+            name            TEXT NOT NULL,
+            version         INTEGER NOT NULL,
+            checksum        TEXT NOT NULL,
+            applied_at      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            UNIQUE (version)
+        );
+    """
+
+    _insert_migration_sql: LiteralString = "INSERT INTO migration_history (name, version, checksum) VALUES (?, ?, ?)"
+
+    def __init__(self, db_path: Path, migration_files: list[Path]):
+        self._migration_files = migration_files
+        self._db_path = db_path
+        self._requested_migrations = [_create_migration(file) for file in migration_files]
+
+    def migrate(self) -> None:
+        applied_migrations: list[MigrationDTO] = self.init_db_and_load_applied_migrations()
+        applied_checksums = [m.checksum for m in applied_migrations]
+        applied_versions = [m.version for m in applied_migrations]
+        migrations_to_apply = [m for m in self._requested_migrations if m.checksum not in applied_checksums]
+        duplicated_versions = [
+            migration.version for migration in migrations_to_apply if migration.version in applied_versions
+        ]
+        if any(duplicated_versions):
+            raise MigrationError(f"Migrations with versions {duplicated_versions} already exist")
+        with duckdb.connect(self._db_path) as conn:
+            for migration in migrations_to_apply:
+                logger.debug("Applying migration %s", migration.name)
+                with conn.cursor() as cur:
+                    cur.execute("START TRANSACTION;")
+                    try:
+                        cur.execute(migration.query)
+                        cur.execute(self._insert_migration_sql, [migration.name, migration.version, migration.checksum])
+                        cur.commit()
+                    except Exception:
+                        cur.rollback()
+                        raise MigrationError(f"Failed to apply migration {migration.name}. Aborting migration process.")
+
+    def init_db_and_load_applied_migrations(self) -> list[MigrationDTO]:
+        with duckdb.connect(str(self._db_path)) as conn:
+            conn.execute(self._init_migration_table_sql)
+            conn.commit()
+            return load_migrations(conn)

databao_context_engine/storage/migrations/V01__init.sql

@@ -0,0 +1,63 @@
+INSTALL vss;
+LOAD vss;
+SET hnsw_enable_experimental_persistence = true;
+
+CREATE SEQUENCE IF NOT EXISTS run_id_seq START 1;
+CREATE SEQUENCE IF NOT EXISTS datasource_run_id_seq START 1;
+CREATE SEQUENCE IF NOT EXISTS chunk_id_seq START 1;
+
+CREATE TABLE IF NOT EXISTS run (
+    run_id          BIGINT PRIMARY KEY DEFAULT nextval('run_id_seq'),
+    project_id      TEXT NOT NULL,
+    started_at      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    ended_at        TIMESTAMP,
+    nemory_version TEXT,
+    run_name TEXT NOT NULL,
+);
+
+CREATE TABLE IF NOT EXISTS datasource_run (
+    datasource_run_id   BIGINT PRIMARY KEY DEFAULT nextval('datasource_run_id_seq'),
+    run_id              BIGINT NOT NULL REFERENCES run(run_id),
+    plugin              TEXT NOT NULL,
+    source_id           TEXT NOT NULL,
+    storage_directory   TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    full_type TEXT NOT NULL,
+);
+
+CREATE TABLE IF NOT EXISTS chunk (
+    chunk_id       BIGINT PRIMARY KEY DEFAULT nextval('chunk_id_seq'),
+    datasource_run_id        BIGINT NOT NULL REFERENCES datasource_run(datasource_run_id),
+    embeddable_text  TEXT NOT NULL,
+    display_text     TEXT,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    generated_description TEXT,
+);
+
+CREATE TABLE IF NOT EXISTS embedding_model_registry (
+    embedder    TEXT NOT NULL,
+    model_id    TEXT NOT NULL,
+    dim         INTEGER NOT NULL,
+    table_name  TEXT NOT NULL,
+    created_at  TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (embedder, model_id),
+    UNIQUE (table_name)
+);
+
+CREATE TABLE IF NOT EXISTS embedding_ollama__nomic_embed_text_v1_5__768 (
+    chunk_id   BIGINT NOT NULL REFERENCES chunk(chunk_id),
+    vec          FLOAT[768] NOT NULL,
+    created_at   TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (chunk_id)
+);
+CREATE INDEX IF NOT EXISTS emb_hnsw_embedding_ollama__nomic_embed_text_v1_5__768 ON embedding_ollama__nomic_embed_text_v1_5__768 USING HNSW (vec) WITH (metric = 'cosine');
+
+INSERT
+OR IGNORE INTO
+    embedding_model_registry(embedder, model_id, dim, table_name)
+VALUES
+    ('ollama', 'nomic-embed-text:v1.5', 768, 'embedding_ollama__nomic_embed_text_v1_5__768');
+
+CREATE INDEX IF NOT EXISTS idx_datasource_run_run ON datasource_run(run_id);
+CREATE INDEX IF NOT EXISTS idx_datasource_run_plugin_run ON datasource_run(plugin, run_id);
+CREATE INDEX IF NOT EXISTS idx_chunk_datasource_run ON chunk(datasource_run_id);

databao_context_engine/storage/models.py

@@ -0,0 +1,51 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class RunDTO:
+    run_id: int
+    run_name: str
+    project_id: str
+    started_at: datetime
+    ended_at: Optional[datetime]
+    nemory_version: str
+
+
+@dataclass(frozen=True)
+class DatasourceRunDTO:
+    datasource_run_id: int
+    run_id: int
+    plugin: str
+    full_type: str
+    source_id: str
+    storage_directory: str
+    created_at: datetime
+
+
+@dataclass(frozen=True)
+class ChunkDTO:
+    chunk_id: int
+    datasource_run_id: int
+    embeddable_text: str
+    display_text: Optional[str]
+    generated_description: str
+    created_at: datetime
+
+
+@dataclass(frozen=True)
+class EmbeddingModelRegistryDTO:
+    embedder: str
+    model_id: str
+    dim: int
+    table_name: str
+    created_at: datetime
+
+
+@dataclass(frozen=True)
+class EmbeddingDTO:
+    chunk_id: int
+    vec: Sequence[float]
+    created_at: datetime

databao_context_engine/storage/repositories/chunk_repository.py

@@ -0,0 +1,130 @@
+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 ChunkDTO
+
+
+class ChunkRepository:
+    def __init__(self, conn: duckdb.DuckDBPyConnection):
+        self._conn = conn
+
+    def create(
+        self,
+        *,
+        datasource_run_id: int,
+        embeddable_text: str,
+        display_text: Optional[str],
+        generated_description: str,
+    ) -> ChunkDTO:
+        try:
+            row = self._conn.execute(
+                """
+            INSERT INTO
+                chunk(datasource_run_id, embeddable_text, display_text, generated_description)
+            VALUES
+                (?, ?, ?, ?)
+            RETURNING
+                *
+            """,
+                [datasource_run_id, embeddable_text, display_text, generated_description],
+            ).fetchone()
+            if row is None:
+                raise RuntimeError("chunk creation returned no object")
+            return self._row_to_dto(row)
+        except ConstraintException as e:
+            raise IntegrityError from e
+
+    def get(self, chunk_id: int) -> Optional[ChunkDTO]:
+        row = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                chunk
+            WHERE
+                chunk_id = ?
+        """,
+            [chunk_id],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def update(
+        self,
+        chunk_id: int,
+        *,
+        embeddable_text: Optional[str] = None,
+        display_text: Optional[str] = None,
+        generated_description: Optional[str] = None,
+    ) -> Optional[ChunkDTO]:
+        sets: list[Any] = []
+        params: list[Any] = []
+
+        if embeddable_text is not None:
+            sets.append("embeddable_text = ?")
+            params.append(embeddable_text)
+        if display_text is not None:
+            sets.append("display_text = ?")
+            params.append(display_text)
+        if generated_description is not None:
+            sets.append("generated_description = ?")
+            params.append(generated_description)
+
+        if not sets:
+            return self.get(chunk_id)
+
+        params.append(chunk_id)
+        self._conn.execute(
+            f"""
+            UPDATE
+                chunk
+            SET
+                {", ".join(sets)}
+            WHERE
+                chunk_id = ?
+        """,
+            params,
+        )
+
+        return self.get(chunk_id)
+
+    def delete(self, chunk_id: int) -> int:
+        row = self._conn.execute(
+            """
+            DELETE FROM
+                chunk
+            WHERE
+                chunk_id = ?
+            RETURNING
+                chunk_id
+            """,
+            [chunk_id],
+        )
+        return 1 if row else 0
+
+    def list(self) -> list[ChunkDTO]:
+        rows = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                chunk
+            ORDER BY
+                chunk_id DESC
+            """
+        ).fetchall()
+        return [self._row_to_dto(r) for r in rows]
+
+    @staticmethod
+    def _row_to_dto(row: Tuple) -> ChunkDTO:
+        chunk_id, datasource_run_id, embeddable_text, display_text, created_at, generated_description = row
+        return ChunkDTO(
+            chunk_id=int(chunk_id),
+            datasource_run_id=int(datasource_run_id),
+            embeddable_text=str(embeddable_text),
+            display_text=display_text,
+            generated_description=str(generated_description),
+            created_at=created_at,
+        )

databao_context_engine/storage/repositories/datasource_run_repository.py

@@ -0,0 +1,136 @@
+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/embedding_model_registry_repository.py

@@ -0,0 +1,87 @@
+from datetime import datetime
+from typing import Optional
+
+import duckdb
+
+from databao_context_engine.services.table_name_policy import TableNamePolicy
+from databao_context_engine.storage.models import EmbeddingModelRegistryDTO
+
+
+class EmbeddingModelRegistryRepository:
+    def __init__(self, conn: duckdb.DuckDBPyConnection):
+        self._conn = conn
+
+    def create(
+        self,
+        *,
+        embedder: str,
+        model_id: str,
+        dim: int,
+        table_name: str,
+    ) -> EmbeddingModelRegistryDTO:
+        TableNamePolicy.validate_table_name(table_name=table_name)
+        row = self._conn.execute(
+            """
+            INSERT INTO
+                embedding_model_registry(embedder, model_id, dim, table_name)
+            VALUES
+                (?, ?, ?, ?)
+            RETURNING
+                *
+            """,
+            [embedder, model_id, dim, table_name],
+        ).fetchone()
+        if row is None:
+            raise RuntimeError("Embedding_model_registry creatuib returned no object")
+        return self._row_to_dto(row)
+
+    def get(
+        self,
+        *,
+        embedder: str,
+        model_id: str,
+    ) -> Optional[EmbeddingModelRegistryDTO]:
+        row = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                embedding_model_registry
+            WHERE
+                embedder = ?
+                AND model_id = ?
+            """,
+            [embedder, model_id],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def delete(
+        self,
+        *,
+        embedder: str,
+        model_id: str,
+    ) -> int:
+        row = self._conn.execute(
+            """
+            DELETE FROM
+                embedding_model_registry
+            WHERE
+                embedder = ?
+                AND model_id = ?
+            RETURNING
+                model_id
+            """,
+            [embedder, model_id],
+        ).fetchone()
+        return 1 if row else 0
+
+    @staticmethod
+    def _row_to_dto(row: tuple) -> EmbeddingModelRegistryDTO:
+        embedder, model_id, dim, table_name, created_at = row
+        return EmbeddingModelRegistryDTO(
+            embedder=str(embedder),
+            model_id=str(model_id),
+            dim=int(dim),
+            table_name=str(table_name),
+            created_at=created_at if isinstance(created_at, datetime) else created_at,
+        )