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

databao_context_engine/plugins/plugin_loader.py

@@ -14,29 +14,11 @@ class DuplicatePluginTypeError(RuntimeError):
     """Raised when two plugins register the same <main>/<sub> plugin key."""
 
 
-PluginList = dict[DatasourceType, BuildPlugin]
-
-
-def get_all_available_plugin_types(exclude_file_plugins: bool = False) -> set[DatasourceType]:
-    return set(load_plugins(exclude_file_plugins=exclude_file_plugins).keys())
-
-
-def get_plugin_for_type(datasource_type: DatasourceType) -> BuildPlugin:
-    all_plugins = load_plugins()
-
-    if datasource_type not in all_plugins:
-        raise ValueError(f"No plugin found for type '{datasource_type.full_type}'")
-
-    return load_plugins()[datasource_type]
-
-
-def load_plugins(exclude_file_plugins: bool = False) -> PluginList:
-    """
-    Loads both builtin and external plugins and merges them into one list
-    """
+def load_plugins(exclude_file_plugins: bool = False) -> dict[DatasourceType, BuildPlugin]:
+    """Load both builtin and external plugins and merges them into one list."""
     builtin_plugins = _load_builtin_plugins(exclude_file_plugins)
     external_plugins = _load_external_plugins(exclude_file_plugins)
-    plugins = merge_plugins(builtin_plugins, external_plugins)
+    plugins = _merge_plugins(builtin_plugins, external_plugins)
 
     return plugins
 
@@ -61,43 +43,70 @@ def _load_builtin_file_plugins() -> list[BuildFilePlugin]:
 
 
 def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
-    """
-    Statically register built-in plugins
-    """
-    from databao_context_engine.plugins.athena_db_plugin import AthenaDbPlugin
-    from databao_context_engine.plugins.clickhouse_db_plugin import ClickhouseDbPlugin
+    """Statically register built-in plugins."""
     from databao_context_engine.plugins.duckdb_db_plugin import DuckDbPlugin
-    from databao_context_engine.plugins.mssql_db_plugin import MSSQLDbPlugin
-    from databao_context_engine.plugins.mysql_db_plugin import MySQLDbPlugin
     from databao_context_engine.plugins.parquet_plugin import ParquetPlugin
-    from databao_context_engine.plugins.postgresql_db_plugin import PostgresqlDbPlugin
-    from databao_context_engine.plugins.snowflake_db_plugin import SnowflakeDbPlugin
 
-    return [
-        AthenaDbPlugin(),
-        ClickhouseDbPlugin(),
+    # optional plugins are added to the python environment via extras
+    optional_plugins: list[BuildDatasourcePlugin] = []
+    try:
+        from databao_context_engine.plugins.mssql_db_plugin import MSSQLDbPlugin
+
+        optional_plugins = [MSSQLDbPlugin()]
+    except ImportError:
+        pass
+
+    try:
+        from databao_context_engine.plugins.clickhouse_db_plugin import ClickhouseDbPlugin
+
+        optional_plugins.append(ClickhouseDbPlugin())
+    except ImportError:
+        pass
+
+    try:
+        from databao_context_engine.plugins.athena_db_plugin import AthenaDbPlugin
+
+        optional_plugins.append(AthenaDbPlugin())
+    except ImportError:
+        pass
+
+    try:
+        from databao_context_engine.plugins.snowflake_db_plugin import SnowflakeDbPlugin
+
+        optional_plugins.append(SnowflakeDbPlugin())
+    except ImportError:
+        pass
+
+    try:
+        from databao_context_engine.plugins.mysql_db_plugin import MySQLDbPlugin
+
+        optional_plugins.append(MySQLDbPlugin())
+    except ImportError:
+        pass
+
+    try:
+        from databao_context_engine.plugins.postgresql_db_plugin import PostgresqlDbPlugin
+
+        optional_plugins.append(PostgresqlDbPlugin())
+    except ImportError:
+        pass
+
+    required_plugins: list[BuildDatasourcePlugin] = [
         DuckDbPlugin(),
-        MSSQLDbPlugin(),
-        MySQLDbPlugin(),
-        PostgresqlDbPlugin(),
-        SnowflakeDbPlugin(),
         ParquetPlugin(),
     ]
+    return required_plugins + optional_plugins
 
 
 def _load_external_plugins(exclude_file_plugins: bool = False) -> list[BuildPlugin]:
-    """
-    Discover external plugins via entry points
-    """
+    """Discover external plugins via entry points."""
     # TODO: implement external plugin loading
     return []
 
 
-def merge_plugins(*plugin_lists: list[BuildPlugin]) -> PluginList:
-    """
-    Merge multiple plugin maps
-    """
-    registry: PluginList = {}
+def _merge_plugins(*plugin_lists: list[BuildPlugin]) -> dict[DatasourceType, BuildPlugin]:
+    """Merge multiple plugin maps."""
+    registry: dict[DatasourceType, BuildPlugin] = {}
     for plugins in plugin_lists:
         for plugin in plugins:
             for full_type in plugin.supported_types():

databao_context_engine/plugins/resources/parquet_introspector.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass, replace
 from urllib.parse import urlparse
 
 import duckdb
-from _duckdb import DuckDBPyConnection
+from duckdb import DuckDBPyConnection
 from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import DuckDBSecret
@@ -18,9 +18,8 @@ logger = logging.getLogger(__name__)
 
 class ParquetConfigFile(BaseModel):
     name: str | None = Field(default=None)
-    type: str = Field(default=parquet_type)
+    type: str = Field(default="parquet")
     url: str = Field(
-        default=type,
         description="Parquet resource location. Should be a valid URL or a path to a local file. "
         "Examples: s3://your_bucket/file.parquet, s3://your-bucket/*.parquet, https://some.url/some_file.parquet, ~/path_to/file.parquet",
     )

databao_context_engine/project/info.py

@@ -9,13 +9,29 @@ from databao_context_engine.system.properties import get_dce_path
 
 @dataclass(kw_only=True, frozen=True)
 class DceProjectInfo:
+    """Information about a Databao Context Engine project.
+
+    Attributes:
+        project_path: The root directory of the Databao Context Engine project.
+        is_initialized: Whether the project has been initialized.
+        project_id: The UUID of the project, or None if the project has not been initialized.
+    """
+
     project_path: Path
-    is_initialised: bool
+    is_initialized: bool
     project_id: UUID | None
 
 
 @dataclass(kw_only=True, frozen=True)
 class DceInfo:
+    """Information about the current Databao Context Engine installation and project.
+
+    Attributes:
+        version: The version of the databao_context_engine package installed on the system.
+        dce_path: The path where databao_context_engine stores its global data.
+        project_info: Information about the Databao Context Engine project.
+    """
+
     version: str
     dce_path: Path
 
@@ -23,6 +39,14 @@ class DceInfo:
 
 
 def get_databao_context_engine_info(project_dir: Path) -> DceInfo:
+    """Return information about the current Databao Context Engine installation and project.
+
+    Args:
+        project_dir: The root directory of the Databao Context Project.
+
+    Returns:
+        A DceInfo instance containing information about the Databao Context Engine installation and project.
+    """
     return DceInfo(
         version=get_dce_version(),
         dce_path=get_dce_path(),
@@ -35,10 +59,15 @@ def _get_project_info(project_dir: Path) -> DceProjectInfo:
 
     return DceProjectInfo(
         project_path=project_dir,
-        is_initialised=project_layout is not None,
+        is_initialized=project_layout is not None,
         project_id=project_layout.read_config_file().project_id if project_layout is not None else None,
     )
 
 
 def get_dce_version() -> str:
+    """Return the installed version of the databao_context_engine package.
+
+    Returns:
+        The installed version of the databao_context_engine package.
+    """
     return version("databao_context_engine")

databao_context_engine/project/init_project.py

@@ -13,12 +13,21 @@ from databao_context_engine.project.project_config import ProjectConfig
 
 
 class InitErrorReason(Enum):
+    """Reasons for which project initialization can fail."""
+
     PROJECT_DIR_DOESNT_EXIST = "PROJECT_DIR_DOESNT_EXIST"
     PROJECT_DIR_NOT_DIRECTORY = "PROJECT_DIR_NOT_DIRECTORY"
-    PROJECT_DIR_ALREADY_INITIALISED = "PROJECT_DIR_ALREADY_INITIALISED"
+    PROJECT_DIR_ALREADY_INITIALIZED = "PROJECT_DIR_ALREADY_INITIALIZED"
 
 
 class InitProjectError(Exception):
+    """Raised when a project can't be initialized.
+
+    Attributes:
+        message: The error message.
+        reason: The reason for the initialization failure.
+    """
+
     reason: InitErrorReason
 
     def __init__(self, reason: InitErrorReason, message: str | None):
@@ -65,20 +74,20 @@ class _ProjectCreator:
 
         if self.config_file.is_file() or self.deprecated_config_file.is_file():
             raise InitProjectError(
-                message=f"Can't initialise a Databao Context Engine project in a folder that already contains a config file. [project_dir: {self.project_dir.resolve()}]",
-                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALISED,
+                message=f"Can't initialize a Databao Context Engine project in a folder that already contains a config file. [project_dir: {self.project_dir.resolve()}]",
+                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALIZED,
             )
 
         if self.src_dir.is_dir():
             raise InitProjectError(
-                message=f"Can't initialise a Databao Context Engine project in a folder that already contains a src directory. [project_dir: {self.project_dir.resolve()}]",
-                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALISED,
+                message=f"Can't initialize a Databao Context Engine project in a folder that already contains a src directory. [project_dir: {self.project_dir.resolve()}]",
+                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALIZED,
             )
 
         if self.examples_dir.is_file():
             raise InitProjectError(
-                message=f"Can't initialise a Databao Context Engine project in a folder that already contains an examples dir. [project_dir: {self.project_dir.resolve()}]",
-                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALISED,
+                message=f"Can't initialize a Databao Context Engine project in a folder that already contains an examples dir. [project_dir: {self.project_dir.resolve()}]",
+                reason=InitErrorReason.PROJECT_DIR_ALREADY_INITIALIZED,
             )
 
         return True

databao_context_engine/project/layout.py

@@ -2,8 +2,8 @@ import logging
 from dataclasses import dataclass
 from pathlib import Path
 
+from databao_context_engine.datasources.types import DatasourceId
 from databao_context_engine.project.project_config import ProjectConfig
-from databao_context_engine.project.types import DatasourceId
 
 SOURCE_FOLDER_NAME = "src"
 OUTPUT_FOLDER_NAME = "output"
@@ -96,12 +96,12 @@ class _ProjectValidator:
 
         if self.config_file is None:
             raise ValueError(
-                f"The current project directory has not been initialised. It should contain a config file. [project_dir: {self.project_dir.resolve()}]"
+                f"The current project directory has not been initialized. It should contain a config file. [project_dir: {self.project_dir.resolve()}]"
             )
 
         if not self.is_src_valid():
             raise ValueError(
-                f"The current project directory has not been initialised. It should contain a src directory. [project_dir: {self.project_dir.resolve()}]"
+                f"The current project directory has not been initialized. It should contain a src directory. [project_dir: {self.project_dir.resolve()}]"
             )
 
         return ProjectLayout(project_dir=self.project_dir, config_file=self.config_file)

databao_context_engine/retrieve_embeddings/__init__.py

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

databao_context_engine/retrieve_embeddings/{internal/export_results.py → export_results.py}

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

databao_context_engine/retrieve_embeddings/{internal/retrieve_runner.py → retrieve_runner.py}

@@ -1,9 +1,9 @@
 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.project.layout import get_output_dir
+from databao_context_engine.retrieve_embeddings.export_results import export_retrieve_results
+from databao_context_engine.retrieve_embeddings.retrieve_service import RetrieveService
 from databao_context_engine.storage.repositories.vector_search_repository import VectorSearchResult
 
 logger = logging.getLogger(__name__)
@@ -15,17 +15,13 @@ def retrieve(
     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
-    )
+    retrieve_results = retrieve_service.retrieve(project_id=project_id, text=text, limit=limit)
 
     if export_to_file:
-        export_directory = get_run_dir(project_dir=project_dir, run_name=resolved_run_name)
+        export_directory = get_output_dir(project_dir)
 
         display_texts = [result.display_text for result in retrieve_results]
         export_file = export_retrieve_results(export_directory, display_texts)

databao_context_engine/retrieve_embeddings/{internal/retrieve_service.py → retrieve_service.py}

@@ -2,9 +2,7 @@ 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,
@@ -17,43 +15,34 @@ 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]:
+    def retrieve(self, *, project_id: str, text: 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}")
+        logger.debug(f"Retrieving display texts 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}")
+        logger.debug(f"Retrieved {len(search_results)} display texts in table {table_name}")
 
         if logger.isEnabledFor(logging.DEBUG):
             closest_result = min(search_results, key=lambda result: result.cosine_distance)
@@ -63,6 +52,3 @@ class RetrieveService:
             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/retrieve_wiring.py

@@ -0,0 +1,49 @@
+from duckdb import DuckDBPyConnection
+
+from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
+from databao_context_engine.llm.factory import create_ollama_embedding_provider, create_ollama_service
+from databao_context_engine.project.layout import ProjectLayout, ensure_project_dir
+from databao_context_engine.retrieve_embeddings.retrieve_runner import retrieve
+from databao_context_engine.retrieve_embeddings.retrieve_service import RetrieveService
+from databao_context_engine.services.factories import create_shard_resolver
+from databao_context_engine.storage.connection import open_duckdb_connection
+from databao_context_engine.storage.repositories.factories import create_vector_search_repository
+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,
+    limit: int | None,
+    export_to_file: bool,
+) -> list[VectorSearchResult]:
+    ensure_project_dir(project_layout.project_dir)
+
+    with open_duckdb_connection(get_db_path(project_layout.project_dir)) 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,
+            limit=limit,
+            export_to_file=export_to_file,
+        )
+
+
+def _create_retrieve_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+) -> RetrieveService:
+    vector_search_repo = create_vector_search_repository(conn)
+    shard_resolver = create_shard_resolver(conn)
+
+    return RetrieveService(
+        vector_search_repo=vector_search_repo,
+        shard_resolver=shard_resolver,
+        provider=embedding_provider,
+    )

databao_context_engine/{serialisation → serialization}/yaml.py

@@ -8,7 +8,7 @@ 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 _)
+        # Doesn't serialize "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)

databao_context_engine/services/chunk_embedding_service.py

@@ -5,7 +5,7 @@ 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.serialization.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
@@ -14,9 +14,22 @@ logger = logging.getLogger(__name__)
 
 
 class ChunkEmbeddingMode(Enum):
+    """Mode controlling how chunks are embedded."""
+
     EMBEDDABLE_TEXT_ONLY = "EMBEDDABLE_TEXT_ONLY"
+    """
+    The embedding is generated only from the string defined by the plugin as embeddable for a chunk.
+    """
+
     GENERATED_DESCRIPTION_ONLY = "GENERATED_DESCRIPTION_ONLY"
+    """
+    The embedding is generated only from a description of the chunk generated by a LLM.
+    """
+
     EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION = "EMBEDDABLE_TEXT_AND_GENERATED_DESCRIPTION"
+    """
+    The embedding is generated from both the embeddable string of the chunk and the description of the chunk generated by a LLM.
+    """
 
     def should_generate_description(self) -> bool:
         return self in (
@@ -44,26 +57,24 @@ class ChunkEmbeddingService:
         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
+    def embed_chunks(self, *, chunks: list[EmbeddableChunk], result: str, full_type: str, datasource_id: 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
+        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}"
+            f"Embedding {len(chunks)} chunks for datasource {datasource_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)
+            chunk_display_text = chunk.content if isinstance(chunk.content, str) else to_yaml_string(chunk.content)
 
             generated_description = ""
             match self._chunk_embedding_mode:
@@ -98,7 +109,8 @@ class ChunkEmbeddingService:
         )
 
         self._persistence_service.write_chunks_and_embeddings(
-            datasource_run_id=datasource_run_id,
             chunk_embeddings=enriched_embeddings,
             table_name=table_name,
+            full_type=full_type,
+            datasource_id=datasource_id,
         )

databao_context_engine/services/factories.py

@@ -1,20 +1,15 @@
-from _duckdb import DuckDBPyConnection
+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,
 )
 
 
@@ -46,43 +41,3 @@ def create_chunk_embedding_service(
         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/persistence_service.py

@@ -24,11 +24,13 @@ class PersistenceService:
         self._dim = dim
 
     def write_chunks_and_embeddings(
-        self, *, datasource_run_id: int, chunk_embeddings: list[ChunkEmbedding], table_name: str
+        self, *, chunk_embeddings: list[ChunkEmbedding], table_name: str, full_type: str, datasource_id: str
     ):
-        """
-        Atomically persist chunks and their vectors.
-        Returns the number of embeddings written.
+        """Atomically persist chunks and their vectors.
+
+        Raises:
+            ValueError: If chunk_embeddings is an empty list.
+
         """
         if not chunk_embeddings:
             raise ValueError("chunk_embeddings must be a non-empty list")
@@ -36,21 +38,19 @@ class PersistenceService:
         with transaction(self._conn):
             for chunk_embedding in chunk_embeddings:
                 chunk_dto = self.create_chunk(
-                    datasource_run_id=datasource_run_id,
+                    full_type=full_type,
+                    datasource_id=datasource_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:
+    def create_chunk(self, *, full_type: str, datasource_id: str, embeddable_text: str, display_text: str) -> ChunkDTO:
         return self._chunk_repo.create(
-            datasource_run_id=datasource_run_id,
+            full_type=full_type,
+            datasource_id=datasource_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]):