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

databao_context_engine/project/types.py

@@ -0,0 +1,134 @@
+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

@@ -0,0 +1,12 @@
+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

@@ -0,0 +1,34 @@
+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_service.py

@@ -0,0 +1,68 @@
+import logging
+from collections.abc import Sequence
+
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.project.runs import resolve_run_name_from_repo
+from databao_context_engine.services.embedding_shard_resolver import EmbeddingShardResolver
+from databao_context_engine.storage.repositories.run_repository import RunRepository
+from databao_context_engine.storage.repositories.vector_search_repository import (
+    VectorSearchRepository,
+    VectorSearchResult,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RetrieveService:
+    def __init__(
+        self,
+        *,
+        run_repo: RunRepository,
+        vector_search_repo: VectorSearchRepository,
+        shard_resolver: EmbeddingShardResolver,
+        provider: EmbeddingProvider,
+    ):
+        self._run_repo = run_repo
+        self._shard_resolver = shard_resolver
+        self._provider = provider
+        self._vector_search_repo = vector_search_repo
+
+    def retrieve(
+        self, *, project_id: str, text: str, run_name: str, limit: int | None = None
+    ) -> list[VectorSearchResult]:
+        if limit is None:
+            limit = 10
+
+        run = self._run_repo.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}'.")
+
+        table_name, dimension = self._shard_resolver.resolve(
+            embedder=self._provider.embedder, model_id=self._provider.model_id
+        )
+
+        retrieve_vec: Sequence[float] = self._provider.embed(text)
+
+        logger.debug(f"Retrieving display texts for run {run.run_id} in table {table_name}")
+
+        search_results = self._vector_search_repo.get_display_texts_by_similarity(
+            table_name=table_name,
+            run_id=run.run_id,
+            retrieve_vec=retrieve_vec,
+            dimension=dimension,
+            limit=limit,
+        )
+
+        logger.debug(f"Retrieved {len(search_results)} display texts for run {run.run_id} in table {table_name}")
+
+        if logger.isEnabledFor(logging.DEBUG):
+            closest_result = min(search_results, key=lambda result: result.cosine_distance)
+            logger.debug(f"Best result: ({closest_result.cosine_distance}, {closest_result.embeddable_text})")
+
+            farthest_result = max(search_results, key=lambda result: result.cosine_distance)
+            logger.debug(f"Worst result: ({farthest_result.cosine_distance}, {farthest_result.embeddable_text})")
+
+        return search_results
+
+    def resolve_run_name(self, *, project_id: str, run_name: str | None) -> str:
+        return resolve_run_name_from_repo(run_repository=self._run_repo, project_id=project_id, run_name=run_name)

databao_context_engine/retrieve_embeddings/internal/retrieve_wiring.py

@@ -0,0 +1,29 @@
+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

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

databao_context_engine/serialisation/yaml.py

@@ -0,0 +1,35 @@
+from typing import Any, Mapping, TextIO, cast
+
+import yaml
+from yaml import Node, SafeDumper
+
+
+def default_representer(dumper: SafeDumper, data: object) -> Node:
+    if isinstance(data, Mapping):
+        return dumper.represent_dict(data)
+    elif hasattr(data, "__dict__"):
+        # Doesn't serialise "private" attributes (that starts with an _)
+        data_public_attributes = {key: value for key, value in data.__dict__.items() if not key.startswith("_")}
+        if data_public_attributes:
+            return dumper.represent_dict(data_public_attributes)
+        else:
+            # If there is no public attributes, we default to the string representation
+            return dumper.represent_str(str(data))
+    else:
+        return dumper.represent_str(str(data))
+
+
+# Registers our default representer only once, when that file is imported
+yaml.add_multi_representer(object, default_representer, Dumper=SafeDumper)
+
+
+def write_yaml_to_stream(*, data: Any, file_stream: TextIO) -> None:
+    _to_yaml(data, file_stream)
+
+
+def to_yaml_string(data: Any) -> str:
+    return cast(str, _to_yaml(data, None))
+
+
+def _to_yaml(data: Any, stream: TextIO | None) -> str | None:
+    return yaml.safe_dump(data, stream, sort_keys=False, default_flow_style=False)

databao_context_engine/services/chunk_embedding_service.py

@@ -0,0 +1,104 @@
+import logging
+from enum import Enum
+from typing import cast
+
+from databao_context_engine.llm.descriptions.provider import DescriptionProvider
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.serialisation.yaml import to_yaml_string
+from databao_context_engine.services.embedding_shard_resolver import EmbeddingShardResolver
+from databao_context_engine.services.models import ChunkEmbedding
+from databao_context_engine.services.persistence_service import PersistenceService
+
+logger = logging.getLogger(__name__)
+
+
+class ChunkEmbeddingMode(Enum):
+    EMBEDDABLE_TEXT_ONLY = "EMBEDDABLE_TEXT_ONLY"
+    GENERATED_DESCRIPTION_ONLY = "GENERATED_DESCRIPTION_ONLY"
+    EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION = "EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION"
+
+    def should_generate_description(self) -> bool:
+        return self in (
+            ChunkEmbeddingMode.GENERATED_DESCRIPTION_ONLY,
+            ChunkEmbeddingMode.EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION,
+        )
+
+
+class ChunkEmbeddingService:
+    def __init__(
+        self,
+        *,
+        persistence_service: PersistenceService,
+        embedding_provider: EmbeddingProvider,
+        description_provider: DescriptionProvider | None,
+        shard_resolver: EmbeddingShardResolver,
+        chunk_embedding_mode: ChunkEmbeddingMode = ChunkEmbeddingMode.EMBEDDABLE_TEXT_ONLY,
+    ):
+        self._persistence_service = persistence_service
+        self._embedding_provider = embedding_provider
+        self._description_provider = description_provider
+        self._shard_resolver = shard_resolver
+        self._chunk_embedding_mode = chunk_embedding_mode
+
+        if self._chunk_embedding_mode.should_generate_description() and description_provider is None:
+            raise ValueError("A DescriptionProvider must be provided when generating descriptions")
+
+    def embed_chunks(self, *, datasource_run_id: int, chunks: list[EmbeddableChunk], result: str) -> None:
+        """
+        Turn plugin chunks into persisted chunks and embeddings
+
+        Flow:
+        1) Embed each chunk into an embedded vector
+        2) Get or create embedding table for the appropriate model and embedding dimensions
+        3) Persist chunks and embeddings vectors in a single transaction
+        """
+
+        if not chunks:
+            return
+
+        logger.debug(
+            f"Embedding {len(chunks)} chunks for datasource run {datasource_run_id}, with chunk_embedding_mode={self._chunk_embedding_mode}"
+        )
+
+        enriched_embeddings: list[ChunkEmbedding] = []
+        for chunk in chunks:
+            chunk_display_text = to_yaml_string(chunk.content)
+
+            generated_description = ""
+            match self._chunk_embedding_mode:
+                case ChunkEmbeddingMode.EMBEDDABLE_TEXT_ONLY:
+                    embedding_text = chunk.embeddable_text
+                case ChunkEmbeddingMode.GENERATED_DESCRIPTION_ONLY:
+                    generated_description = cast(DescriptionProvider, self._description_provider).describe(
+                        text=chunk_display_text, context=result
+                    )
+                    embedding_text = generated_description
+                case ChunkEmbeddingMode.EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION:
+                    generated_description = cast(DescriptionProvider, self._description_provider).describe(
+                        text=chunk_display_text, context=result
+                    )
+                    embedding_text = generated_description + "\n" + chunk.embeddable_text
+
+            vec = self._embedding_provider.embed(embedding_text)
+
+            enriched_embeddings.append(
+                ChunkEmbedding(
+                    chunk=chunk,
+                    vec=vec,
+                    display_text=chunk_display_text,
+                    generated_description=generated_description,
+                )
+            )
+
+        table_name = self._shard_resolver.resolve_or_create(
+            embedder=self._embedding_provider.embedder,
+            model_id=self._embedding_provider.model_id,
+            dim=self._embedding_provider.dim,
+        )
+
+        self._persistence_service.write_chunks_and_embeddings(
+            datasource_run_id=datasource_run_id,
+            chunk_embeddings=enriched_embeddings,
+            table_name=table_name,
+        )

databao_context_engine/services/embedding_shard_resolver.py

@@ -0,0 +1,64 @@
+import duckdb
+
+from databao_context_engine.services.table_name_policy import TableNamePolicy
+from databao_context_engine.storage.repositories.embedding_model_registry_repository import (
+    EmbeddingModelRegistryRepository,
+)
+
+
+class EmbeddingShardResolver:
+    def __init__(
+        self,
+        *,
+        conn: duckdb.DuckDBPyConnection,
+        registry_repo: EmbeddingModelRegistryRepository,
+        table_name_policy: TableNamePolicy | None = None,
+    ):
+        self._conn = conn
+        self._registry = registry_repo
+        self._policy = table_name_policy or TableNamePolicy()
+
+    def resolve(self, *, embedder: str, model_id: str) -> tuple[str, int]:
+        row = self._registry.get(embedder=embedder, model_id=model_id)
+        if not row:
+            raise ValueError(f"Model not registered: {embedder}:{model_id}")
+        return row.table_name, row.dim
+
+    def resolve_or_create(self, *, embedder: str, model_id: str, dim: int) -> str:
+        row = self._registry.get(embedder=embedder, model_id=model_id)
+        if row:
+            if row.dim != dim:
+                raise ValueError(f"Model already registered with dim={row.dim}, requested dim={dim}")
+            return row.table_name
+
+        table_name = self._policy.build(embedder=embedder, model_id=model_id, dim=dim)
+        self._create_table_and_index(table_name, dim)
+
+        self._registry.create(
+            embedder=embedder,
+            model_id=model_id,
+            dim=dim,
+            table_name=table_name,
+        )
+
+        return table_name
+
+    def _create_table_and_index(self, table_name: str, dim: int) -> None:
+        self._conn.execute("LOAD vss;")
+        self._conn.execute("SET hnsw_enable_experimental_persistence = true;")
+
+        self._conn.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {table_name} (
+                chunk_id BIGINT NOT NULL REFERENCES chunk(chunk_id),
+                vec FLOAT[{dim}] NOT NULL,
+                created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                PRIMARY KEY (chunk_id)
+            )
+            """
+        )
+        self._conn.execute(
+            f"""
+            CREATE INDEX IF NOT EXISTS emb_hnsw_{table_name} ON {table_name} USING HNSW (vec) WITH (metric='cosine');
+            """
+        )

databao_context_engine/services/factories.py

@@ -0,0 +1,88 @@
+from _duckdb import DuckDBPyConnection
+
+from databao_context_engine.build_sources.internal.build_service import BuildService
+from databao_context_engine.llm.descriptions.provider import DescriptionProvider
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.retrieve_embeddings.internal.retrieve_service import RetrieveService
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode, ChunkEmbeddingService
+from databao_context_engine.services.embedding_shard_resolver import EmbeddingShardResolver
+from databao_context_engine.services.persistence_service import PersistenceService
+from databao_context_engine.services.table_name_policy import TableNamePolicy
+from databao_context_engine.storage.repositories.factories import (
+    create_chunk_repository,
+    create_datasource_run_repository,
+    create_embedding_repository,
+    create_registry_repository,
+    create_run_repository,
+    create_vector_search_repository,
+)
+
+
+def create_shard_resolver(conn: DuckDBPyConnection, policy: TableNamePolicy | None = None) -> EmbeddingShardResolver:
+    return EmbeddingShardResolver(
+        conn=conn, registry_repo=create_registry_repository(conn), table_name_policy=policy or TableNamePolicy()
+    )
+
+
+def create_persistence_service(conn: DuckDBPyConnection) -> PersistenceService:
+    return PersistenceService(
+        conn=conn, chunk_repo=create_chunk_repository(conn), embedding_repo=create_embedding_repository(conn)
+    )
+
+
+def create_chunk_embedding_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+    description_provider: DescriptionProvider | None,
+    chunk_embedding_mode: ChunkEmbeddingMode,
+) -> ChunkEmbeddingService:
+    resolver = create_shard_resolver(conn)
+    persistence = create_persistence_service(conn)
+    return ChunkEmbeddingService(
+        persistence_service=persistence,
+        embedding_provider=embedding_provider,
+        shard_resolver=resolver,
+        description_provider=description_provider,
+        chunk_embedding_mode=chunk_embedding_mode,
+    )
+
+
+def create_build_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+    description_provider: DescriptionProvider | None,
+    chunk_embedding_mode: ChunkEmbeddingMode,
+) -> BuildService:
+    run_repo = create_run_repository(conn)
+    datasource_run_repo = create_datasource_run_repository(conn)
+    chunk_embedding_service = create_chunk_embedding_service(
+        conn,
+        embedding_provider=embedding_provider,
+        description_provider=description_provider,
+        chunk_embedding_mode=chunk_embedding_mode,
+    )
+
+    return BuildService(
+        run_repo=run_repo,
+        datasource_run_repo=datasource_run_repo,
+        chunk_embedding_service=chunk_embedding_service,
+    )
+
+
+def create_retrieve_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+) -> RetrieveService:
+    run_repo = create_run_repository(conn)
+    vector_search_repo = create_vector_search_repository(conn)
+    shard_resolver = create_shard_resolver(conn)
+
+    return RetrieveService(
+        run_repo=run_repo,
+        vector_search_repo=vector_search_repo,
+        shard_resolver=shard_resolver,
+        provider=embedding_provider,
+    )

databao_context_engine/services/models.py

@@ -0,0 +1,12 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+
+
+@dataclass(frozen=True)
+class ChunkEmbedding:
+    chunk: EmbeddableChunk
+    vec: Sequence[float]
+    display_text: str
+    generated_description: str

databao_context_engine/services/persistence_service.py

@@ -0,0 +1,61 @@
+from collections.abc import Sequence
+
+import duckdb
+
+from databao_context_engine.services.models import ChunkEmbedding
+from databao_context_engine.storage.models import ChunkDTO
+from databao_context_engine.storage.repositories.chunk_repository import ChunkRepository
+from databao_context_engine.storage.repositories.embedding_repository import EmbeddingRepository
+from databao_context_engine.storage.transaction import transaction
+
+
+class PersistenceService:
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection,
+        chunk_repo: ChunkRepository,
+        embedding_repo: EmbeddingRepository,
+        *,
+        dim: int = 768,
+    ):
+        self._conn = conn
+        self._chunk_repo = chunk_repo
+        self._embedding_repo = embedding_repo
+        self._dim = dim
+
+    def write_chunks_and_embeddings(
+        self, *, datasource_run_id: int, chunk_embeddings: list[ChunkEmbedding], table_name: str
+    ):
+        """
+        Atomically persist chunks and their vectors.
+        Returns the number of embeddings written.
+        """
+        if not chunk_embeddings:
+            raise ValueError("chunk_embeddings must be a non-empty list")
+
+        with transaction(self._conn):
+            for chunk_embedding in chunk_embeddings:
+                chunk_dto = self.create_chunk(
+                    datasource_run_id=datasource_run_id,
+                    embeddable_text=chunk_embedding.chunk.embeddable_text,
+                    display_text=chunk_embedding.display_text,
+                    generated_description=chunk_embedding.generated_description,
+                )
+                self.create_embedding(table_name=table_name, chunk_id=chunk_dto.chunk_id, vec=chunk_embedding.vec)
+
+    def create_chunk(
+        self, *, datasource_run_id: int, embeddable_text: str, display_text: str, generated_description: str
+    ) -> ChunkDTO:
+        return self._chunk_repo.create(
+            datasource_run_id=datasource_run_id,
+            embeddable_text=embeddable_text,
+            display_text=display_text,
+            generated_description=generated_description,
+        )
+
+    def create_embedding(self, *, table_name: str, chunk_id: int, vec: Sequence[float]):
+        self._embedding_repo.create(
+            table_name=table_name,
+            chunk_id=chunk_id,
+            vec=vec,
+        )

databao_context_engine/services/run_name_policy.py

@@ -0,0 +1,8 @@
+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')}"