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

databao_context_engine/storage/repositories/embedding_repository.py

@@ -0,0 +1,113 @@
+from typing import Optional, Sequence, Tuple
+
+import duckdb
+from _duckdb import ConstraintException
+
+from databao_context_engine.services.table_name_policy import TableNamePolicy
+from databao_context_engine.storage.exceptions.exceptions import IntegrityError
+from databao_context_engine.storage.models import EmbeddingDTO
+
+
+class EmbeddingRepository:
+    def __init__(self, conn: duckdb.DuckDBPyConnection):
+        self._conn = conn
+
+    def create(
+        self,
+        *,
+        table_name: str,
+        chunk_id: int,
+        vec: Sequence[float],
+    ) -> EmbeddingDTO:
+        try:
+            TableNamePolicy.validate_table_name(table_name=table_name)
+            row = self._conn.execute(
+                f"""
+            INSERT INTO
+                {table_name} (chunk_id, vec)
+            VALUES
+                (?, ?)
+            RETURNING
+                *
+            """,
+                [chunk_id, vec],
+            ).fetchone()
+            if row is None:
+                raise RuntimeError("Embedding creation returned no object")
+            return self._row_to_dto(row)
+        except ConstraintException as e:
+            raise IntegrityError from e
+
+    def get(self, *, table_name: str, chunk_id: int) -> Optional[EmbeddingDTO]:
+        TableNamePolicy.validate_table_name(table_name=table_name)
+        row = self._conn.execute(
+            f"""
+            SELECT 
+                *
+            FROM 
+                {table_name}
+            WHERE 
+                chunk_id = ?
+            """,
+            [chunk_id],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def update(
+        self,
+        *,
+        table_name: str,
+        chunk_id: int,
+        vec: Sequence[float],
+    ) -> Optional[EmbeddingDTO]:
+        TableNamePolicy.validate_table_name(table_name=table_name)
+        self._conn.execute(
+            f"""
+            UPDATE 
+                {table_name}
+            SET 
+                vec = ?
+            WHERE 
+                chunk_id = ?
+            """,
+            [list(vec), chunk_id],
+        )
+        return self.get(table_name=table_name, chunk_id=chunk_id)
+
+    def delete(self, *, table_name: str, chunk_id: int) -> int:
+        TableNamePolicy.validate_table_name(table_name=table_name)
+        row = self._conn.execute(
+            f"""
+            DELETE FROM 
+                {table_name}
+            WHERE 
+                chunk_id = ?
+            RETURNING 
+                chunk_id
+            """,
+            [chunk_id],
+        ).fetchone()
+        return 1 if row else 0
+
+    def list(self, table_name: str) -> list[EmbeddingDTO]:
+        TableNamePolicy.validate_table_name(table_name=table_name)
+        rows = self._conn.execute(
+            f"""
+            SELECT
+                *
+            FROM                
+                {table_name}
+            ORDER BY 
+                chunk_id DESC
+            """
+        ).fetchall()
+        return [self._row_to_dto(r) for r in rows]
+
+    @staticmethod
+    def _row_to_dto(row: Tuple) -> EmbeddingDTO:
+        chunk_id, vec, created_at = row
+        return EmbeddingDTO(
+            chunk_id=int(chunk_id),
+            vec=list(vec) if not isinstance(vec, list) else vec,
+            created_at=created_at,
+        )

databao_context_engine/storage/repositories/factories.py

@@ -0,0 +1,35 @@
+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)
+
+
+def create_embedding_repository(conn: DuckDBPyConnection) -> EmbeddingRepository:
+    return EmbeddingRepository(conn)
+
+
+def create_registry_repository(conn: DuckDBPyConnection) -> EmbeddingModelRegistryRepository:
+    return EmbeddingModelRegistryRepository(conn)
+
+
+def create_vector_search_repository(conn: DuckDBPyConnection) -> VectorSearchRepository:
+    return VectorSearchRepository(conn)

databao_context_engine/storage/repositories/run_repository.py

@@ -0,0 +1,157 @@
+from datetime import datetime
+from typing import Any, Optional
+
+import duckdb
+
+from databao_context_engine.services.run_name_policy import RunNamePolicy
+from databao_context_engine.storage.models import RunDTO
+
+
+class RunRepository:
+    def __init__(self, conn: duckdb.DuckDBPyConnection, run_name_policy: RunNamePolicy):
+        self._conn = conn
+        self._run_name_policy = run_name_policy
+
+    def create(
+        self, *, project_id: str, dce_version: Optional[str] = None, started_at: datetime | None = None
+    ) -> RunDTO:
+        if started_at is None:
+            started_at = datetime.now()
+        run_name = self._run_name_policy.build(run_started_at=started_at)
+
+        row = self._conn.execute(
+            """
+            INSERT INTO 
+                run (project_id, nemory_version, started_at, run_name)
+            VALUES 
+                (?, ?, ?, ?)
+            RETURNING 
+                *
+            """,
+            [project_id, dce_version, started_at, run_name],
+        ).fetchone()
+        if row is None:
+            raise RuntimeError("Run creation returned no object")
+        return self._row_to_dto(row)
+
+    def get(self, run_id: int) -> Optional[RunDTO]:
+        row = self._conn.execute(
+            """
+            SELECT 
+                *
+            FROM
+                run
+            WHERE
+                run_id = ?
+            """,
+            [run_id],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def get_by_run_name(self, *, project_id: str, run_name: str) -> RunDTO | None:
+        row = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                run
+            WHERE
+                run.project_id = ? AND run_name = ?
+            """,
+            [project_id, run_name],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def get_latest_run_for_project(self, project_id: str) -> RunDTO | None:
+        row = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                run
+            WHERE
+                run.project_id = ?
+            ORDER BY run.started_at DESC
+            LIMIT 1
+            """,
+            [project_id],
+        ).fetchone()
+        return self._row_to_dto(row) if row else None
+
+    def update(
+        self,
+        run_id: int,
+        *,
+        project_id: Optional[str] = None,
+        ended_at: Optional[datetime] = None,
+        dce_version: Optional[str] = None,
+    ) -> Optional[RunDTO]:
+        sets: list[Any] = []
+        params: list[Any] = []
+
+        if project_id is not None:
+            sets.append("project_id = ?")
+            params.append(project_id)
+        if ended_at is not None:
+            sets.append("ended_at = ?")
+            params.append(ended_at)
+        if dce_version is not None:
+            sets.append("nemory_version = ?")
+            params.append(dce_version)
+
+        if not sets:
+            return self.get(run_id)
+
+        params.append(run_id)
+        self._conn.execute(
+            f"""
+                    UPDATE
+                        run
+                    SET 
+                        {", ".join(sets)}
+                    WHERE
+                        run_id = ?
+                    """,
+            params,
+        )
+
+        return self.get(run_id)
+
+    def delete(self, run_id: int) -> int:
+        row = self._conn.execute(
+            """
+            DELETE FROM
+                run
+            WHERE
+                run_id = ?
+            RETURNING
+                run_id
+            """,
+            [run_id],
+        ).fetchone()
+        return 1 if row else 0
+
+    def list(self) -> list[RunDTO]:
+        rows = self._conn.execute(
+            """
+            SELECT
+                *
+            FROM
+                run
+            ORDER BY
+                run_id DESC
+            """
+        ).fetchall()
+        return [self._row_to_dto(r) for r in rows]
+
+    @staticmethod
+    def _row_to_dto(row: tuple) -> RunDTO:
+        run_id, project_id, started_at, ended_at, dce_version, run_name = row
+        return RunDTO(
+            run_id=int(run_id),
+            run_name=run_name,
+            project_id=str(project_id),
+            started_at=started_at,
+            ended_at=ended_at,
+            nemory_version=dce_version,
+        )

databao_context_engine/storage/repositories/vector_search_repository.py

@@ -0,0 +1,63 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import duckdb
+
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.types import DatasourceId
+
+
+@dataclass(kw_only=True, frozen=True)
+class VectorSearchResult:
+    display_text: str
+    embeddable_text: str
+    cosine_distance: float
+    datasource_type: DatasourceType
+    datasource_id: DatasourceId
+
+
+class VectorSearchRepository:
+    _DEFAULT_DISTANCE_THRESHOLD = 0.75
+
+    def __init__(self, conn: duckdb.DuckDBPyConnection):
+        self._conn = conn
+
+    def get_display_texts_by_similarity(
+        self, *, table_name: str, run_id: int, 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
+        """
+        rows = self._conn.execute(
+            f"""
+            SELECT
+                COALESCE(c.display_text, c.embeddable_text) AS display_text,
+                c.embeddable_text AS embeddable_text,
+                array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) AS cosine_distance,
+                dr.full_type,
+                dr.source_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 < ?
+            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],
+        ).fetchall()
+
+        return [
+            VectorSearchResult(
+                display_text=row[0],
+                embeddable_text=row[1],
+                cosine_distance=row[2],
+                datasource_type=DatasourceType(full_type=row[3]),
+                datasource_id=DatasourceId.from_string_repr(row[4]),
+            )
+            for row in rows
+        ]

databao_context_engine/storage/transaction.py

@@ -0,0 +1,14 @@
+from contextlib import contextmanager
+
+import duckdb
+
+
+@contextmanager
+def transaction(conn: duckdb.DuckDBPyConnection):
+    conn.execute("BEGIN")
+    try:
+        yield
+        conn.execute("COMMIT")
+    except Exception:
+        conn.execute("ROLLBACK")
+        raise

databao_context_engine/system/properties.py

@@ -0,0 +1,13 @@
+import os
+from pathlib import Path
+
+# 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()
+
+
+def get_dce_path() -> Path:
+    return _dce_path
+
+
+def get_db_path() -> Path:
+    return get_dce_path() / "dce.duckdb"

databao_context_engine/templating/renderer.py

@@ -0,0 +1,29 @@
+import os
+
+from jinja2.sandbox import SandboxedEnvironment
+
+
+class DceTemplateError(Exception):
+    pass
+
+
+class UnknownEnvVarTemplateError(DceTemplateError):
+    pass
+
+
+def render_template(source: str) -> str:
+    env = SandboxedEnvironment()
+
+    return env.from_string(source=str(source)).render(env_var=resolve_env_var)
+
+
+def resolve_env_var(env_var: str, default: str | None = None) -> str:
+    if env_var in os.environ:
+        return os.environ[env_var]
+
+    if default is not None:
+        return default
+
+    raise UnknownEnvVarTemplateError(
+        f"Error in template. The environment variable {env_var} is missing and no default was provided"
+    )

databao_context_engine-0.1.1.dist-info/METADATA

@@ -0,0 +1,186 @@
+Metadata-Version: 2.3
+Name: databao-context-engine
+Version: 0.1.1
+Summary: Add your description here
+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: pymysql>=1.1.2
+Requires-Dist: clickhouse-connect>=0.10.0
+Requires-Dist: mcp>=1.23.3
+Requires-Dist: pyathena>=3.22.0
+Requires-Dist: snowflake-connector-python>=4.1.0
+Requires-Dist: mssql-python>=1.0.0
+Requires-Dist: pydantic>=2.12.4
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: asyncpg>=0.31.0
+Requires-Dist: asyncio>=4.0.0
+Requires-Dist: asyncpg-stubs>=0.31.1
+Requires-Python: >=3.12
+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             |
+
+# Prerequisites
+
+This README assumes you will use `uv` as your package manager.
+
+You can install it following the instructions [here](https://docs.astral.sh/uv/getting-started/installation/)
+
+If you are going to push to the repository, please make sure to install git pre-commit hooks by running
+
+```bash
+  uv run pre-commit install
+```
+
+# How to run?
+
+You can run it with:
+
+```bash
+  uv run dce info
+```
+
+Not providing the `info` subcommand or using the `--help` flag will show the help screen for the command.
+
+## Using the dce command directly
+
+To be able to use the `dce` command directly (without using `uv run` or `python`) there are two options.
+
+### Installing dce locally
+
+For that one needs to:
+
+1. Build the project by running
+
+```bash
+  uv build
+```
+
+2. Installing the project on our machine by running:
+
+```bash
+  uv tool install -e .
+```
+
+This second step will install the `dce` script on your machine and add it into your path.
+
+### Create dce alias using nix
+
+This method will simply create a new shell environment with `dce` alias. For that one needs to install `nix` package
+manager (https://nixos.org/download/). After that one could simply run in the project root
+
+```bash
+$ nix-shell
+```
+
+which is a short version of `$ nix-shell shell.nix`.
+
+Alternatively, one could specify the path to the project repository
+
+```bash
+$ nix-shell {path_to_dce_repository}
+```
+
+After that, you can then directly use:
+
+```bash
+  dce --help
+```
+
+Note: when we actually release our built Python package, users that don't use `uv` will still be able to install the CLI
+by using `pipx install` instead.
+
+# Running Mypy
+
+[mypy](https://mypy.readthedocs.io/en/stable/getting_started.html) has been added to the project for type checking.
+
+You can run it with the following:
+
+```bash
+  uv run mypy src --exclude "test_*" --exclude dist
+```
+
+NB: the above runs type checking on all files within the `src` directory, excluding all test files.
+
+# Running tests
+
+You can run the tests with:
+
+```bash
+  uv run pytest
+```
+
+(there is currently one test succeeding and one test failing in the project)
+
+# Generating JSON Schemas for our plugin's config files
+
+To be able to build a datasource, each plugin requires a yaml config file that describes how to connect to the
+datasource,
+as well as other information needed to customise the plugin.
+
+To document what each config file should look like, we can generate a JSON schema describing the fields allowed in that
+file.
+
+You can generate all JSON schemas for all plugins by running:
+
+```bash
+  uv run generate_configs_schemas
+```
+
+Some options can be provided to the command to choose which plugins to include or exclude from the generation.
+To see the options available, you can refer to the help:
+
+```bash
+  uv run generate_configs_schemas --help
+```
+