databao-context-engine 0.1.2__tar.gz → 0.1.4.dev1__tar.gz

databao_context_engine-0.1.4.dev1/PKG-INFO

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

databao_context_engine-0.1.4.dev1/README.md

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

{databao_context_engine-0.1.2 → databao_context_engine-0.1.4.dev1}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "databao-context-engine"
-version = "0.1.2"
-description = "Add your description here"
+version = "0.1.4.dev1"
+description = "Semantic context for your LLMs — generated automatically"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
@@ -9,22 +9,34 @@ dependencies = [
     "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",
     "pydantic>=2.12.4",
     "jinja2>=3.1.6",
-    "asyncpg>=0.31.0",
-    "asyncio>=4.0.0",
-    "asyncpg-stubs>=0.31.1",
+]
+default-optional-dependency-keys = [
+    "mysql",
+    "postgresql"
 ]
 
 [project.optional-dependencies]
 mssql = [
     "mssql-python>=1.0.0"
 ]
+clickhouse = [
+    "clickhouse-connect>=0.10.0"
+]
+athena = [
+    "pyathena>=3.25.0"
+]
+snowflake = [
+    "snowflake-connector-python>=4.2.0",
+]
+mysql = [
+    "pymysql>=1.1.2"
+]
+postgresql = [
+    "asyncpg>=0.31.0",
+]
 
 [build-system]
 requires = ["uv_build>=0.9.6,<0.10.0"]
@@ -33,14 +45,15 @@ build-backend = "uv_build"
 [dependency-groups]
 dev = [
     "mypy>=1.18.2",
-    "pre-commit>=4.3.0",
-    "pytest>=8.4.2",
+    "pre-commit>=4.5.1",
+    "pytest>=9.0.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",
+    "asyncpg-stubs>=0.31.1",
 ]
 
 [project.scripts]
@@ -48,7 +61,7 @@ dce = "databao_context_engine.main:main"
 generate_configs_schemas = "databao_context_engine.generate_configs_schemas:main"
 
 [tool.uv]
-override-dependencies = ["urllib3>=2.6.3", "filelock>=3.20.3"]
+override-dependencies = ["urllib3>=2.6.3", "filelock>=3.20.3", "virtualenv>=20.36.1"]
 
 [tool.uv.build-backend]
 source-exclude = ["tests"]

{databao_context_engine-0.1.2 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/__init__.py

@@ -1,16 +1,23 @@
-from databao_context_engine.build_sources.internal.build_runner import BuildContextResult
+from databao_context_engine.build_sources import BuildContextResult
+from databao_context_engine.databao_context_engine import ContextSearchResult, DatabaoContextEngine
 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 (
+from databao_context_engine.datasources.check_config import (
     CheckDatasourceConnectionResult,
     DatasourceConnectionStatus,
 )
-from databao_context_engine.datasource_config.datasource_context import DatasourceContext
+from databao_context_engine.datasources.datasource_context import DatasourceContext
+from databao_context_engine.datasources.types import Datasource, DatasourceId
 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.plugin_loader import DatabaoContextPluginLoader
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    BuildFilePlugin,
+    BuildPlugin,
+    DatasourceType,
+)
+from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
 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__ = [
@@ -32,4 +39,9 @@ __all__ = [
     "init_or_get_dce_project",
     "InitErrorReason",
     "InitProjectError",
+    "DatabaoContextPluginLoader",
+    "ConfigPropertyDefinition",
+    "BuildPlugin",
+    "BuildDatasourcePlugin",
+    "BuildFilePlugin",
 ]

databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources/__init__.py

@@ -0,0 +1,4 @@
+from databao_context_engine.build_sources.build_runner import BuildContextResult
+from databao_context_engine.build_sources.build_wiring import build_all_datasources
+
+__all__ = ["build_all_datasources", "BuildContextResult"]

{databao_context_engine-0.1.2/src/databao_context_engine/build_sources/internal → databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources}/build_runner.py

@@ -3,22 +3,32 @@ 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 (
+from databao_context_engine.build_sources.build_service import BuildService
+from databao_context_engine.build_sources.export_results import (
     append_result_to_all_results,
-    create_run_dir,
     export_build_result,
+    reset_all_results,
 )
+from databao_context_engine.datasources.datasource_discovery import discover_datasources, prepare_source
+from databao_context_engine.datasources.types import DatasourceId
 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
+from databao_context_engine.project.layout import get_output_dir
 
 logger = logging.getLogger(__name__)
 
 
 @dataclass
 class BuildContextResult:
+    """Result of a single datasource context build.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        datasource_type: The type of the datasource.
+        context_built_at: The timestamp when the context was built.
+        context_file_path: The path to the generated context file.
+    """
+
     datasource_id: DatasourceId
     datasource_type: DatasourceType
     context_built_at: datetime
@@ -29,16 +39,19 @@ def build(
     project_dir: Path,
     *,
     build_service: BuildService,
-    project_id: str,
-    dce_version: str,
 ) -> list[BuildContextResult]:
-    """
-    Build entrypoint.
+    """Build the context for all datasources in the project.
+
+    Unless you already have access to BuildService, this should not be called directly.
+    Instead, internal callers should go through the build_wiring module or directly use DatabaoContextProjectManager.build_context().
 
     1) Load available plugins
     2) Discover sources
     3) Create a run
     4) For each source, call process_source
+
+    Returns:
+        A list of all the contexts built.
     """
     plugins = load_plugins()
 
@@ -48,11 +61,9 @@ def build(
         logger.info("No sources discovered under %s", project_dir)
         return []
 
-    run = None
-    run_dir = None
-
     number_of_failed_builds = 0
     build_result = []
+    reset_all_results(get_output_dir(project_dir))
     for discovered_datasource in datasources:
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -71,19 +82,15 @@ def build(
                 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)
+            output_dir = get_output_dir(project_dir)
+
+            context_file_path = export_build_result(output_dir, result)
+            append_result_to_all_results(output_dir, result)
 
             build_result.append(
                 BuildContextResult(
@@ -99,9 +106,6 @@ def build(
 
             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),

databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources/build_service.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import logging
+
+from databao_context_engine.build_sources.plugin_execution import BuiltDatasourceContext, execute
+from databao_context_engine.datasources.types import PreparedDatasource
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildPlugin,
+)
+from databao_context_engine.serialization.yaml import to_yaml_string
+from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingService
+
+logger = logging.getLogger(__name__)
+
+
+class BuildService:
+    def __init__(
+        self,
+        *,
+        chunk_embedding_service: ChunkEmbeddingService,
+    ) -> None:
+        self._chunk_embedding_service = chunk_embedding_service
+
+    def process_prepared_source(
+        self,
+        *,
+        prepared_source: PreparedDatasource,
+        plugin: BuildPlugin,
+    ) -> BuiltDatasourceContext:
+        """Process a single source to build its context.
+
+        1) Execute the plugin
+        2) Divide the results into chunks
+        3) Embed and persist the chunks
+
+        Returns:
+            The built context.
+        """
+        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
+
+        self._chunk_embedding_service.embed_chunks(
+            chunks=chunks,
+            result=to_yaml_string(result.context),
+            full_type=prepared_source.datasource_type.full_type,
+            datasource_id=result.datasource_id,
+        )
+
+        return result

databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources/build_wiring.py

@@ -0,0 +1,84 @@
+import logging
+from pathlib import Path
+
+from duckdb import DuckDBPyConnection
+
+from databao_context_engine.build_sources.build_runner import BuildContextResult, build
+from databao_context_engine.build_sources.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.llm.factory import (
+    create_ollama_description_provider,
+    create_ollama_embedding_provider,
+    create_ollama_service,
+)
+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_chunk_embedding_service
+from databao_context_engine.storage.connection import open_duckdb_connection
+from databao_context_engine.storage.migrate import migrate
+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]:
+    """Build the context for all datasources in the project.
+
+    - Instantiates the build service
+    - Delegates the actual build logic to the build runner
+
+    Returns:
+        A list of all the contexts built.
+    """
+    ensure_project_dir(project_dir)
+
+    logger.debug(f"Starting to build datasources in project {project_dir.resolve()}")
+
+    # Think about alternative solutions. This solution will mirror the current behaviour
+    # The current behaviour only builds what is currently in the /src folder
+    # This will need to change in the future when we can pick which datasources to build
+    db_path = get_db_path(project_dir)
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    if db_path.exists():
+        db_path.unlink()
+
+    with open_duckdb_connection(db_path) as conn:
+        migrate(db_path)
+
+        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,
+        )
+        return build(
+            project_dir=project_dir,
+            build_service=build_service,
+        )
+
+
+def _create_build_service(
+    conn: DuckDBPyConnection,
+    *,
+    embedding_provider: EmbeddingProvider,
+    description_provider: DescriptionProvider | None,
+    chunk_embedding_mode: ChunkEmbeddingMode,
+) -> BuildService:
+    chunk_embedding_service = create_chunk_embedding_service(
+        conn,
+        embedding_provider=embedding_provider,
+        description_provider=description_provider,
+        chunk_embedding_mode=chunk_embedding_mode,
+    )
+
+    return BuildService(
+        chunk_embedding_service=chunk_embedding_service,
+    )

databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources/export_results.py

@@ -0,0 +1,41 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.build_sources.plugin_execution import BuiltDatasourceContext
+from databao_context_engine.datasources.datasource_context import get_context_header_for_datasource
+from databao_context_engine.datasources.types import DatasourceId
+from databao_context_engine.project.layout import ALL_RESULTS_FILE_NAME
+from databao_context_engine.serialization.yaml import write_yaml_to_stream
+
+logger = logging.getLogger(__name__)
+
+
+def export_build_result(output_dir: Path, result: BuiltDatasourceContext) -> Path:
+    datasource_id = DatasourceId.from_string_repr(result.datasource_id)
+    export_file_path = output_dir.joinpath(datasource_id.relative_path_to_context_file())
+
+    # Make sure the parent folder exists
+    export_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with export_file_path.open("w") as export_file:
+        write_yaml_to_stream(data=result, file_stream=export_file)
+
+    logger.info(f"Exported result to {export_file_path.resolve()}")
+
+    return export_file_path
+
+
+def append_result_to_all_results(output_dir: Path, result: BuiltDatasourceContext):
+    path = output_dir.joinpath(ALL_RESULTS_FILE_NAME)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a", encoding="utf-8") as export_file:
+        export_file.write(get_context_header_for_datasource(DatasourceId.from_string_repr(result.datasource_id)))
+        write_yaml_to_stream(data=result, file_stream=export_file)
+        export_file.write("\n")
+
+
+# Here temporarily because it needs to be reset between runs.
+# A subsequent PR will remove the existence of the all_results file
+def reset_all_results(output_dir: Path):
+    path = output_dir.joinpath(ALL_RESULTS_FILE_NAME)
+    path.unlink(missing_ok=True)

{databao_context_engine-0.1.2/src/databao_context_engine/build_sources/internal → databao_context_engine-0.1.4.dev1/src/databao_context_engine/build_sources}/plugin_execution.py

@@ -2,20 +2,18 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, cast
 
+from databao_context_engine.datasources.types import DatasourceId, PreparedConfig, PreparedDatasource
 from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     BuildFilePlugin,
     BuildPlugin,
 )
 from databao_context_engine.pluginlib.plugin_utils import execute_datasource_plugin, execute_file_plugin
-from databao_context_engine.project.types import PreparedConfig, PreparedDatasource, DatasourceId
 
 
 @dataclass()
 class BuiltDatasourceContext:
-    """
-    Dataclass defining the result of building a datasource's context.
-    """
+    """Dataclass defining the result of building a datasource's context."""
 
     datasource_id: str
     """
@@ -53,9 +51,7 @@ def execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> Bui
 
 
 def _execute(prepared_datasource: PreparedDatasource, plugin: BuildPlugin) -> Any:
-    """
-    Run a prepared source through the plugin
-    """
+    """Run a prepared source through the plugin."""
     if isinstance(prepared_datasource, PreparedConfig):
         ds_plugin = cast(BuildDatasourcePlugin, plugin)