databao-context-engine 0.1.2__py3-none-any.whl → 0.1.4.dev1__py3-none-any.whl

databao_context_engine/plugins/databases/snowflake_introspector.py

@@ -1,22 +1,57 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Mapping
+from typing import Annotated, Any, Dict, List
 
 import snowflake.connector
-from pydantic import Field
+from pydantic import BaseModel, Field
 from snowflake.connector import DictCursor
 
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
 from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
 
 
+class SnowflakePasswordAuth(BaseModel):
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+
+
+class SnowflakeKeyPairAuth(BaseModel):
+    private_key_file: str | None = None
+    private_key_file_pwd: str | None = None
+    private_key: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+
+
+class SnowflakeSSOAuth(BaseModel):
+    authenticator: str = Field(description='e.g. "externalbrowser"')
+
+
+class SnowflakeConnectionProperties(BaseModel):
+    account: Annotated[str, ConfigPropertyAnnotation(required=True)]
+    warehouse: str | None = None
+    database: str | None = None
+    user: str | None = None
+    role: str | None = None
+    auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth
+    additional_properties: dict[str, Any] = {}
+
+    def to_snowflake_kwargs(self) -> dict[str, Any]:
+        kwargs = self.model_dump(
+            exclude={
+                "additional_properties": True,
+            },
+            exclude_none=True,
+        )
+        auth_fields = kwargs.pop("auth", {})
+        kwargs.update(auth_fields)
+        kwargs.update(self.additional_properties)
+        return kwargs
+
+
 class SnowflakeConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/snowflake")
-    connection: dict[str, Any] = Field(
-        description="Connection parameters for Snowflake. It can contain any of the keys supported by the Snowflake connection library"
-    )
+    type: str = Field(default="snowflake")
+    connection: SnowflakeConnectionProperties
 
 
 class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
@@ -28,11 +63,9 @@ class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
 
     def _connect(self, file_config: SnowflakeConfigFile):
         connection = file_config.connection
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
         snowflake.connector.paramstyle = "qmark"
         return snowflake.connector.connect(
-            **connection,
+            **connection.to_snowflake_kwargs(),
         )
 
     def _connect_to_catalog(self, file_config: SnowflakeConfigFile, catalog: str):
@@ -41,7 +74,7 @@ class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
         return snowflake.connector.connect(**cfg)
 
     def _get_catalogs(self, connection, file_config: SnowflakeConfigFile) -> list[str]:
-        database = file_config.connection.get("database")
+        database = file_config.connection.database
         if database:
             return [database]
 

databao_context_engine/plugins/duckdb_tools.py

@@ -0,0 +1,18 @@
+from typing import Any
+
+from databao_context_engine.pluginlib.config import DuckDBSecret
+
+
+def generate_create_secret_sql(secret_name, duckdb_secret: DuckDBSecret) -> str:
+    parameters = [("type", duckdb_secret.type)] + list(duckdb_secret.properties.items())
+    return f"""CREATE SECRET {secret_name} (
+    {", ".join([f"{k} '{v}'" for (k, v) in parameters])}
+);
+"""
+
+
+def fetchall_dicts(cur, sql: str, params=None) -> list[dict[str, Any]]:
+    cur.execute(sql, params or [])
+    columns = [desc[0].lower() for desc in cur.description] if cur.description else []
+    rows = cur.fetchall()
+    return [dict(zip(columns, row)) for row in rows]

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,16 +43,9 @@ 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.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
 
     # optional plugins are added to the python environment via extras
     optional_plugins: list[BuildDatasourcePlugin] = []
@@ -81,31 +56,57 @@ def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
     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] = [
-        AthenaDbPlugin(),
-        ClickhouseDbPlugin(),
         DuckDbPlugin(),
-        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,10 +6,11 @@ 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
+from databao_context_engine.plugins.duckdb_tools import fetchall_dicts, generate_create_secret_sql
 
 parquet_type = "resources/parquet"
 
@@ -18,9 +19,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",
     )
@@ -50,14 +50,6 @@ class ParquetIntrospectionResult:
     files: list[ParquetFile]
 
 
-def generate_create_secret_sql(secret_name, duckdb_secret: DuckDBSecret) -> str:
-    parameters = [("type", duckdb_secret.type)] + list(duckdb_secret.properties.items())
-    return f"""CREATE SECRET {secret_name} (
-    {", ".join([f"{k} {v}" for (k, v) in parameters])}
-);
-"""
-
-
 @contextlib.contextmanager
 def _create_secret(conn: DuckDBPyConnection, duckdb_secret: DuckDBSecret):
     secret_name = duckdb_secret.name or "gen_secret_" + str(uuid.uuid4()).replace("-", "_")
@@ -98,10 +90,9 @@ class ParquetIntrospector:
         with self._connect(file_config) as conn:
             with conn.cursor() as cur:
                 resolved_url = _resolve_url(file_config)
-                cur.execute(f"SELECT * FROM parquet_file_metadata('{resolved_url}') LIMIT 1")
-                columns = [desc[0].lower() for desc in cur.description] if cur.description else []
-                rows = cur.fetchall()
-                parquet_file_metadata = [dict(zip(columns, row)) for row in rows]
+                parquet_file_metadata = fetchall_dicts(
+                    cur, f"SELECT * FROM parquet_file_metadata('{resolved_url}') LIMIT 1"
+                )
                 if not parquet_file_metadata:
                     raise ValueError(f"No parquet files found by url {resolved_url}")
                 if not parquet_file_metadata or not parquet_file_metadata[0]["file_name"]:
@@ -111,10 +102,7 @@ class ParquetIntrospector:
         with self._connect(file_config) as conn:
             with conn.cursor() as cur:
                 resolved_url = _resolve_url(file_config)
-                cur.execute(f"SELECT * from parquet_metadata('{resolved_url}')")
-                cols = [desc[0].lower() for desc in cur.description] if cur.description else []
-                rows = cur.fetchall()
-                file_metas = [dict(zip(cols, row)) for row in rows]
+                file_metas = fetchall_dicts(cur, f"SELECT * from parquet_metadata('{resolved_url}')")
 
                 columns_per_file: dict[str, dict[int, ParquetColumn]] = defaultdict(defaultdict)
                 for file_meta in file_metas:

databao_context_engine/project/info.py

@@ -3,29 +3,56 @@ from importlib.metadata import version
 from pathlib import Path
 from uuid import UUID
 
+from databao_context_engine.plugins.plugin_loader import load_plugins
 from databao_context_engine.project.layout import validate_project_dir
 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
+    plugin_ids: list[str]
 
     project_info: DceProjectInfo
 
 
 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(),
+        plugin_ids=[plugin.id for plugin in load_plugins().values()],
         project_info=_get_project_info(project_dir),
     )
 
@@ -35,10 +62,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)