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

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]):

databao_context_engine/storage/connection.py

@@ -1,24 +1,28 @@
 import logging
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Iterator
 
 import duckdb
-
-from databao_context_engine.system.properties import get_db_path
+from duckdb import DuckDBPyConnection
 
 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.
+def open_duckdb_connection(db_path: str | Path) -> Iterator[DuckDBPyConnection]:
+    """Open a DuckDB connection with vector search enabled and close on exist.
+
+    It also loads the vss extension and enables HNSW experimental persistence on the DuckDB.
 
     Usage:
         with open_duckdb_connection() as conn:
+
+    Yields:
+        The opened DuckDB connection.
+
     """
-    path = str(db_path or get_db_path())
+    path = str(db_path)
     conn = duckdb.connect(path)
     logger.debug(f"Connected to DuckDB database at {path}")
 

databao_context_engine/storage/exceptions/exceptions.py

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

databao_context_engine/storage/migrate.py

@@ -8,12 +8,10 @@ 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:
+def migrate(db_path: str | Path, migration_files: list[Path] | None = None) -> None:
     if migration_files is None:
         migration_files = [
             migration
@@ -21,7 +19,7 @@ def migrate(db_path: str | Path | None = None, migration_files: list[Path] | Non
             if isinstance(migration, Path) and ".sql" == migration.suffix
         ]
 
-    db = Path(db_path or get_db_path()).expanduser().resolve()
+    db = Path(db_path).expanduser().resolve()
     db.parent.mkdir(parents=True, exist_ok=True)
     logger.debug("Running migrations on database: %s", db)
 

databao_context_engine/storage/migrations/V01__init.sql

@@ -2,36 +2,15 @@ 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,
+    chunk_id        BIGINT PRIMARY KEY DEFAULT nextval('chunk_id_seq'),
+    full_type       TEXT NOT NULL,
+    datasource_id   TEXT NOT NULL,
+    embeddable_text TEXT NOT NULL,
+    display_text    TEXT,
+    created_at      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
 );
 
 CREATE TABLE IF NOT EXISTS embedding_model_registry (
@@ -57,7 +36,3 @@ 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

@@ -4,34 +4,13 @@ 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
+    full_type: str
+    datasource_id: str
     embeddable_text: str
     display_text: Optional[str]
-    generated_description: str
     created_at: datetime
 
 

databao_context_engine/storage/repositories/chunk_repository.py

@@ -14,22 +14,22 @@ class ChunkRepository:
     def create(
         self,
         *,
-        datasource_run_id: int,
+        full_type: str,
+        datasource_id: str,
         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)
+                chunk(full_type, datasource_id, embeddable_text, display_text)
             VALUES
                 (?, ?, ?, ?)
             RETURNING
                 *
             """,
-                [datasource_run_id, embeddable_text, display_text, generated_description],
+                [full_type, datasource_id, embeddable_text, display_text],
             ).fetchone()
             if row is None:
                 raise RuntimeError("chunk creation returned no object")
@@ -55,22 +55,26 @@ class ChunkRepository:
         self,
         chunk_id: int,
         *,
+        full_type: Optional[str] = None,
+        datasource_id: Optional[str] = None,
         embeddable_text: Optional[str] = None,
         display_text: Optional[str] = None,
-        generated_description: Optional[str] = None,
     ) -> Optional[ChunkDTO]:
         sets: list[Any] = []
         params: list[Any] = []
 
+        if full_type is not None:
+            sets.append("full_type = ?")
+            params.append(full_type)
+        if datasource_id is not None:
+            sets.append("datasource_id = ?")
+            params.append(datasource_id)
         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)
@@ -119,12 +123,12 @@ class ChunkRepository:
 
     @staticmethod
     def _row_to_dto(row: Tuple) -> ChunkDTO:
-        chunk_id, datasource_run_id, embeddable_text, display_text, created_at, generated_description = row
+        chunk_id, full_type, datasource_id, embeddable_text, display_text, created_at = row
         return ChunkDTO(
             chunk_id=int(chunk_id),
-            datasource_run_id=int(datasource_run_id),
-            embeddable_text=str(embeddable_text),
+            full_type=full_type,
+            datasource_id=datasource_id,
+            embeddable_text=embeddable_text,
             display_text=display_text,
-            generated_description=str(generated_description),
             created_at=created_at,
         )

databao_context_engine/storage/repositories/factories.py

@@ -1,24 +1,13 @@
-from _duckdb import DuckDBPyConnection
+from duckdb import DuckDBPyConnection
 
-from databao_context_engine.services.run_name_policy import RunNamePolicy
 from databao_context_engine.storage.repositories.chunk_repository import ChunkRepository
-from databao_context_engine.storage.repositories.datasource_run_repository import DatasourceRunRepository
 from databao_context_engine.storage.repositories.embedding_model_registry_repository import (
     EmbeddingModelRegistryRepository,
 )
 from databao_context_engine.storage.repositories.embedding_repository import EmbeddingRepository
-from databao_context_engine.storage.repositories.run_repository import RunRepository
 from databao_context_engine.storage.repositories.vector_search_repository import VectorSearchRepository
 
 
-def create_run_repository(conn: DuckDBPyConnection) -> RunRepository:
-    return RunRepository(conn, run_name_policy=RunNamePolicy())
-
-
-def create_datasource_run_repository(conn: DuckDBPyConnection) -> DatasourceRunRepository:
-    return DatasourceRunRepository(conn)
-
-
 def create_chunk_repository(conn: DuckDBPyConnection) -> ChunkRepository:
     return ChunkRepository(conn)
 

databao_context_engine/storage/repositories/vector_search_repository.py

@@ -3,8 +3,8 @@ from dataclasses import dataclass
 
 import duckdb
 
+from databao_context_engine.datasources.types import DatasourceId
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
-from databao_context_engine.project.types import DatasourceId
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -23,32 +23,27 @@ class VectorSearchRepository:
         self._conn = conn
 
     def get_display_texts_by_similarity(
-        self, *, table_name: str, run_id: int, retrieve_vec: Sequence[float], dimension: int, limit: int
+        self, *, table_name: str, retrieve_vec: Sequence[float], dimension: int, limit: int
     ) -> list[VectorSearchResult]:
-        """
-        Read only similarity search on a specific embedding shard table.
-        Returns the display text for the closest matches in a given run
-        """
+        """Read only similarity search on a specific embedding shard table."""
         rows = self._conn.execute(
             f"""
             SELECT
                 COALESCE(c.display_text, c.embeddable_text) AS display_text,
-                c.embeddable_text AS embeddable_text,
+                c.embeddable_text,
                 array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) AS cosine_distance,
-                dr.full_type,
-                dr.source_id,
+                c.full_type,
+                c.datasource_id,
             FROM
                 {table_name} e
                 JOIN chunk c ON e.chunk_id = c.chunk_id
-                JOIN datasource_run dr ON c.datasource_run_id = dr.datasource_run_id
             WHERE
-                dr.run_id = ?
-                AND cosine_distance < ?
+                cosine_distance < ?
             ORDER BY
                 array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) ASC
             LIMIT ?
             """,
-            [list(retrieve_vec), run_id, self._DEFAULT_DISTANCE_THRESHOLD, list(retrieve_vec), limit],
+            [list(retrieve_vec), self._DEFAULT_DISTANCE_THRESHOLD, list(retrieve_vec), limit],
         ).fetchall()
 
         return [

databao_context_engine/system/properties.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
 
+from databao_context_engine.project.layout import get_output_dir
+
 # it's private, so it doesn't get imported directy. This value is mocked in tests
 _dce_path = Path(os.getenv("DATABAO_CONTEXT_ENGINE_PATH") or "~/.dce").expanduser().resolve()
 
@@ -9,5 +11,5 @@ def get_dce_path() -> Path:
     return _dce_path
 
 
-def get_db_path() -> Path:
-    return get_dce_path() / "dce.duckdb"
+def get_db_path(project_dir: Path) -> Path:
+    return get_output_dir(project_dir) / "dce.duckdb"

databao_context_engine-0.1.4.dev1.dist-info/METADATA

@@ -0,0 +1,75 @@
+Metadata-Version: 2.3
+Name: databao-context-engine
+Version: 0.1.4.dev1
+Summary: Semantic context for your LLMs — generated automatically
+Requires-Dist: click>=8.3.0
+Requires-Dist: duckdb>=1.4.3
+Requires-Dist: pyyaml>=6.0.3
+Requires-Dist: requests>=2.32.5
+Requires-Dist: mcp>=1.23.3
+Requires-Dist: pydantic>=2.12.4
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: pyathena>=3.25.0 ; extra == 'athena'
+Requires-Dist: clickhouse-connect>=0.10.0 ; extra == 'clickhouse'
+Requires-Dist: mssql-python>=1.0.0 ; extra == 'mssql'
+Requires-Dist: pymysql>=1.1.2 ; extra == 'mysql'
+Requires-Dist: asyncpg>=0.31.0 ; extra == 'postgresql'
+Requires-Dist: snowflake-connector-python>=4.2.0 ; extra == 'snowflake'
+Requires-Python: >=3.12
+Provides-Extra: athena
+Provides-Extra: clickhouse
+Provides-Extra: mssql
+Provides-Extra: mysql
+Provides-Extra: postgresql
+Provides-Extra: snowflake
+Description-Content-Type: text/markdown
+
+[![official project](https://jb.gg/badges/official.svg)](https://confluence.jetbrains.com/display/ALL/JetBrains+on+GitHub)
+[![License](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](https://github.com/JetBrains/databao-context-engine/blob/main/LICENSE)
+
+[//]: # ([![PyPI version](https://img.shields.io/pypi/v/databao-context-engine.svg)](https://pypi.org/project/databao-context-engine))
+
+[//]: # ([![Python versions](https://img.shields.io/pypi/pyversions/databao-context-engine.svg)](https://pypi.org/project/databao-context-engine/))
+
+
+<h1 align="center">Databao Context Engine</h1>
+<p align="center">
+ <b>Semantic context for your LLMs — generated automatically.</b><br/>
+ No more copying schemas. No manual documentation. Just accurate answers.
+</p>
+<p align="center">
+ <a href="https://databao.app">Website</a> 
+
+[//]: # (•)
+
+[//]: # ( <a href="#quickstart">Quickstart</a> •)
+
+[//]: # ( <a href="#supported-data-sources">Data Sources</a> •)
+
+[//]: # ( <a href="#contributing">Contributing</a>)
+</p>
+
+---
+
+## What is Databao Context Engine?
+
+Databao Context Engine **automatically generates governed semantic context** from your databases, BI tools, documents, and spreadsheets.
+
+Integrate it with any LLM to deliver **accurate, context-aware answers** — without copying schemas or writing documentation by hand.
+
+```
+Your data sources → Context Engine → Unified semantic graph → Any LLM
+```
+
+## Why choose Databao Context Engine?
+
+| Feature                    | What it means for you                                          |
+|----------------------------|----------------------------------------------------------------|
+| **Auto-generated context** | Extracts schemas, relationships, and semantics automatically   |
+| **Runs locally**           | Your data never leaves your environment                        |
+| **MCP integration**        | Works with Claude Desktop, Cursor, and any MCP-compatible tool |
+| **Multiple sources**       | Databases, dbt projects, spreadsheets, documents               |
+| **Built-in benchmarks**    | Measure and improve context quality over time                  |
+| **LLM agnostic**           | OpenAI, Anthropic, Ollama, Gemini — use any model              |
+| **Governed & versioned**   | Track, version, and share context across your team             |
+| **Dynamic or static**      | Serve context via MCP server or export as artifact             |