databao-context-engine 0.1.1__tar.gz

databao_context_engine-0.1.1/PKG-INFO

@@ -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
+```
+

databao_context_engine-0.1.1/README.md

@@ -0,0 +1,164 @@
+[![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
+```
+

databao_context_engine-0.1.1/pyproject.toml

@@ -0,0 +1,57 @@
+[project]
+name = "databao-context-engine"
+version = "0.1.1"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "click>=8.3.0",
+    "duckdb>=1.4.3",
+    "pyyaml>=6.0.3",
+    "requests>=2.32.5",
+    "pymysql>=1.1.2",
+    "clickhouse-connect>=0.10.0",
+    "mcp>=1.23.3",
+    "pyathena>=3.22.0",
+    "snowflake-connector-python>=4.1.0",
+    "mssql-python>=1.0.0",
+    "pydantic>=2.12.4",
+    "jinja2>=3.1.6",
+    "asyncpg>=0.31.0",
+    "asyncio>=4.0.0",
+    "asyncpg-stubs>=0.31.1",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.6,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "mypy>=1.18.2",
+    "pre-commit>=4.3.0",
+    "pytest>=8.4.2",
+    "pytest-unordered>=0.7.0",
+    "ruff>=0.14.2",
+    "testcontainers==4.12.0",
+    "types-pyyaml>=6.0.12.20250915",
+    "types-pymysql>=1.1.0.20250916",
+    "pytest-mock>=3.15.1",
+]
+
+[project.scripts]
+dce = "databao_context_engine.main:main"
+generate_configs_schemas = "databao_context_engine.generate_configs_schemas:main"
+
+[tool.uv]
+override-dependencies = ["urllib3>=2.6.3"]
+
+[tool.uv.build-backend]
+source-exclude = ["tests"]
+
+[[tool.mypy.overrides]]
+module = ["requests", "requests.*"]
+ignore_missing_imports = true
+
+[tool.pytest]
+addopts = ["--import-mode=importlib"]

databao_context_engine-0.1.1/src/databao_context_engine/__init__.py

@@ -0,0 +1,35 @@
+from databao_context_engine.build_sources.internal.build_runner import BuildContextResult
+from databao_context_engine.databao_context_project_manager import DatabaoContextProjectManager, DatasourceConfigFile
+from databao_context_engine.databao_engine import ContextSearchResult, DatabaoContextEngine
+from databao_context_engine.datasource_config.check_config import (
+    CheckDatasourceConnectionResult,
+    DatasourceConnectionStatus,
+)
+from databao_context_engine.datasource_config.datasource_context import DatasourceContext
+from databao_context_engine.init_project import init_dce_project, init_or_get_dce_project
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.info import DceInfo, get_databao_context_engine_info
+from databao_context_engine.project.init_project import InitErrorReason, InitProjectError
+from databao_context_engine.project.types import Datasource, DatasourceId
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
+
+__all__ = [
+    "DatabaoContextEngine",
+    "Datasource",
+    "DatasourceId",
+    "DatasourceContext",
+    "ContextSearchResult",
+    "DatabaoContextProjectManager",
+    "ChunkEmbeddingMode",
+    "BuildContextResult",
+    "CheckDatasourceConnectionResult",
+    "DatasourceConnectionStatus",
+    "DatasourceConfigFile",
+    "DatasourceType",
+    "get_databao_context_engine_info",
+    "DceInfo",
+    "init_dce_project",
+    "init_or_get_dce_project",
+    "InitErrorReason",
+    "InitProjectError",
+]

databao_context_engine-0.1.1/src/databao_context_engine/build_sources/internal/build_runner.py

@@ -0,0 +1,111 @@
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+
+from databao_context_engine.build_sources.internal.build_service import BuildService
+from databao_context_engine.build_sources.internal.export_results import (
+    append_result_to_all_results,
+    create_run_dir,
+    export_build_result,
+)
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.plugins.plugin_loader import load_plugins
+from databao_context_engine.project.datasource_discovery import discover_datasources, prepare_source
+from databao_context_engine.project.types import DatasourceId
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BuildContextResult:
+    datasource_id: DatasourceId
+    datasource_type: DatasourceType
+    context_built_at: datetime
+    context_file_path: Path
+
+
+def build(
+    project_dir: Path,
+    *,
+    build_service: BuildService,
+    project_id: str,
+    dce_version: str,
+) -> list[BuildContextResult]:
+    """
+    Build entrypoint.
+
+    1) Load available plugins
+    2) Discover sources
+    3) Create a run
+    4) For each source, call process_source
+    """
+    plugins = load_plugins()
+
+    datasources = discover_datasources(project_dir)
+
+    if not datasources:
+        logger.info("No sources discovered under %s", project_dir)
+        return []
+
+    run = None
+    run_dir = None
+
+    number_of_failed_builds = 0
+    build_result = []
+    for discovered_datasource in datasources:
+        try:
+            prepared_source = prepare_source(discovered_datasource)
+
+            logger.info(
+                f'Found datasource of type "{prepared_source.datasource_type.full_type}" with name {prepared_source.path.stem}'
+            )
+
+            plugin = plugins.get(prepared_source.datasource_type)
+            if plugin is None:
+                logger.warning(
+                    "No plugin for '%s' (datasource=%s) — skipping.",
+                    prepared_source.datasource_type.full_type,
+                    prepared_source.path,
+                )
+                number_of_failed_builds += 1
+                continue
+
+            if run is None or run_dir is None:
+                # Initialiase the run as soon as we found a source
+                run = build_service.start_run(project_id=project_id, dce_version=dce_version)
+                run_dir = create_run_dir(project_dir, run.run_name)
+
+            result = build_service.process_prepared_source(
+                run_id=run.run_id,
+                prepared_source=prepared_source,
+                plugin=plugin,
+            )
+
+            context_file_path = export_build_result(run_dir, result)
+            append_result_to_all_results(run_dir, result)
+
+            build_result.append(
+                BuildContextResult(
+                    datasource_id=DatasourceId.from_string_repr(result.datasource_id),
+                    datasource_type=DatasourceType(full_type=result.datasource_type),
+                    context_built_at=result.context_built_at,
+                    context_file_path=context_file_path,
+                )
+            )
+        except Exception as e:
+            logger.debug(str(e), exc_info=True, stack_info=True)
+            logger.info(f"Failed to build source at ({discovered_datasource.path}): {str(e)}")
+
+            number_of_failed_builds += 1
+
+    if run is not None:
+        build_service.finalize_run(run_id=run.run_id)
+
+    logger.debug(
+        "Successfully built %d datasources. %s",
+        len(build_result),
+        f"Failed to build {number_of_failed_builds}." if number_of_failed_builds > 0 else "",
+    )
+
+    return build_result

databao_context_engine-0.1.1/src/databao_context_engine/build_sources/internal/build_service.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from databao_context_engine.build_sources.internal.plugin_execution import BuiltDatasourceContext, execute
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildPlugin,
+)
+from databao_context_engine.project.types import PreparedDatasource
+from databao_context_engine.serialisation.yaml import to_yaml_string
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingService
+from databao_context_engine.storage.models import RunDTO
+from databao_context_engine.storage.repositories.datasource_run_repository import DatasourceRunRepository
+from databao_context_engine.storage.repositories.run_repository import RunRepository
+
+logger = logging.getLogger(__name__)
+
+
+class BuildService:
+    def __init__(
+        self,
+        *,
+        run_repo: RunRepository,
+        datasource_run_repo: DatasourceRunRepository,
+        chunk_embedding_service: ChunkEmbeddingService,
+    ) -> None:
+        self._run_repo = run_repo
+        self._datasource_run_repo = datasource_run_repo
+        self._chunk_embedding_service = chunk_embedding_service
+
+    def start_run(self, *, project_id: str, dce_version: str) -> RunDTO:
+        """
+        Create a new run row and return (run_id, started_at).
+        """
+        return self._run_repo.create(project_id=project_id, dce_version=dce_version)
+
+    def finalize_run(self, *, run_id: int):
+        """
+        Mark the run as complete (sets ended_at).
+        """
+        self._run_repo.update(run_id=run_id, ended_at=datetime.now())
+
+    def process_prepared_source(
+        self,
+        *,
+        run_id: int,
+        prepared_source: PreparedDatasource,
+        plugin: BuildPlugin,
+    ) -> BuiltDatasourceContext:
+        """
+        Process a single source.
+
+        1) Execute the plugin
+        2) Divide the results into chunks
+        3) Embed and persist the chunks
+        """
+        result = execute(prepared_source, plugin)
+
+        chunks = plugin.divide_context_into_chunks(result.context)
+        if not chunks:
+            logger.info("No chunks for %s — skipping.", prepared_source.path.name)
+            return result
+
+        datasource_run = self._datasource_run_repo.create(
+            run_id=run_id,
+            plugin=plugin.name,
+            full_type=prepared_source.datasource_type.full_type,
+            source_id=result.datasource_id,
+            storage_directory=str(prepared_source.path.parent),
+        )
+
+        self._chunk_embedding_service.embed_chunks(
+            datasource_run_id=datasource_run.datasource_run_id, chunks=chunks, result=to_yaml_string(result.context)
+        )
+
+        return result

databao_context_engine-0.1.1/src/databao_context_engine/build_sources/internal/build_wiring.py

@@ -0,0 +1,52 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.build_sources.internal.build_runner import BuildContextResult, build
+from databao_context_engine.llm.factory import (
+    create_ollama_description_provider,
+    create_ollama_embedding_provider,
+    create_ollama_service,
+)
+from databao_context_engine.project.info import get_dce_version
+from databao_context_engine.project.layout import ensure_project_dir
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
+from databao_context_engine.services.factories import (
+    create_build_service,
+)
+from databao_context_engine.storage.connection import open_duckdb_connection
+from databao_context_engine.system.properties import get_db_path
+
+logger = logging.getLogger(__name__)
+
+
+def build_all_datasources(project_dir: Path, chunk_embedding_mode: ChunkEmbeddingMode) -> list[BuildContextResult]:
+    """
+    Public build entrypoint
+    - Instantiates the build service
+    - Delegates the actual build logic to the build runner
+    """
+    project_layout = ensure_project_dir(project_dir)
+
+    logger.debug(f"Starting to build datasources in project {project_dir.resolve()}")
+
+    with open_duckdb_connection(get_db_path()) as conn:
+        ollama_service = create_ollama_service()
+        embedding_provider = create_ollama_embedding_provider(ollama_service)
+        description_provider = (
+            create_ollama_description_provider(ollama_service)
+            if chunk_embedding_mode.should_generate_description()
+            else None
+        )
+        build_service = create_build_service(
+            conn,
+            embedding_provider=embedding_provider,
+            description_provider=description_provider,
+            chunk_embedding_mode=chunk_embedding_mode,
+        )
+        dce_config = project_layout.read_config_file()
+        return build(
+            project_dir=project_dir,
+            build_service=build_service,
+            project_id=str(dce_config.project_id),
+            dce_version=get_dce_version(),
+        )