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

databao_context_engine/plugins/plugin_loader.py

@@ -18,9 +18,8 @@ def load_plugins(exclude_file_plugins: bool = False) -> dict[DatasourceType, Bui
     """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)
 
-    return plugins
+    return _merge_plugins(builtin_plugins, external_plugins)
 
 
 def _load_builtin_plugins(exclude_file_plugins: bool = False) -> list[BuildPlugin]:
@@ -35,7 +34,7 @@ def _load_builtin_plugins(exclude_file_plugins: bool = False) -> list[BuildPlugi
 
 
 def _load_builtin_file_plugins() -> list[BuildFilePlugin]:
-    from databao_context_engine.plugins.unstructured_files_plugin import InternalUnstructuredFilesPlugin
+    from databao_context_engine.plugins.files.unstructured_files_plugin import InternalUnstructuredFilesPlugin
 
     return [
         InternalUnstructuredFilesPlugin(),
@@ -44,57 +43,56 @@ def _load_builtin_file_plugins() -> list[BuildFilePlugin]:
 
 def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
     """Statically register built-in plugins."""
-    from databao_context_engine.plugins.duckdb_db_plugin import DuckDbPlugin
-    from databao_context_engine.plugins.parquet_plugin import ParquetPlugin
+    from databao_context_engine.plugins.databases.duckdb.duckdb_db_plugin import DuckDbPlugin
+    from databao_context_engine.plugins.databases.sqlite.sqlite_db_plugin import SQLiteDbPlugin
+    from databao_context_engine.plugins.dbt.dbt_plugin import DbtPlugin
+    from databao_context_engine.plugins.resources.parquet_plugin import ParquetPlugin
 
     # optional plugins are added to the python environment via extras
     optional_plugins: list[BuildDatasourcePlugin] = []
     try:
-        from databao_context_engine.plugins.mssql_db_plugin import MSSQLDbPlugin
+        from databao_context_engine.plugins.databases.mssql.mssql_db_plugin import MSSQLDbPlugin
 
         optional_plugins = [MSSQLDbPlugin()]
     except ImportError:
         pass
 
     try:
-        from databao_context_engine.plugins.clickhouse_db_plugin import ClickhouseDbPlugin
+        from databao_context_engine.plugins.databases.clickhouse.clickhouse_db_plugin import ClickhouseDbPlugin
 
         optional_plugins.append(ClickhouseDbPlugin())
     except ImportError:
         pass
 
     try:
-        from databao_context_engine.plugins.athena_db_plugin import AthenaDbPlugin
+        from databao_context_engine.plugins.databases.athena.athena_db_plugin import AthenaDbPlugin
 
         optional_plugins.append(AthenaDbPlugin())
     except ImportError:
         pass
 
     try:
-        from databao_context_engine.plugins.snowflake_db_plugin import SnowflakeDbPlugin
+        from databao_context_engine.plugins.databases.snowflake.snowflake_db_plugin import SnowflakeDbPlugin
 
         optional_plugins.append(SnowflakeDbPlugin())
     except ImportError:
         pass
 
     try:
-        from databao_context_engine.plugins.mysql_db_plugin import MySQLDbPlugin
+        from databao_context_engine.plugins.databases.mysql.mysql_db_plugin import MySQLDbPlugin
 
         optional_plugins.append(MySQLDbPlugin())
     except ImportError:
         pass
 
     try:
-        from databao_context_engine.plugins.postgresql_db_plugin import PostgresqlDbPlugin
+        from databao_context_engine.plugins.databases.postgresql.postgresql_db_plugin import PostgresqlDbPlugin
 
         optional_plugins.append(PostgresqlDbPlugin())
     except ImportError:
         pass
 
-    required_plugins: list[BuildDatasourcePlugin] = [
-        DuckDbPlugin(),
-        ParquetPlugin(),
-    ]
+    required_plugins: list[BuildDatasourcePlugin] = [DuckDbPlugin(), ParquetPlugin(), SQLiteDbPlugin(), DbtPlugin()]
     return required_plugins + optional_plugins
 
 

databao_context_engine/plugins/resources/parquet_introspector.py

@@ -12,7 +12,7 @@ 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"
+parquet_type = "parquet"
 
 logger = logging.getLogger(__name__)
 

databao_context_engine/plugins/{parquet_plugin.py → resources/parquet_plugin.py}

@@ -24,9 +24,7 @@ class ParquetPlugin(BuildDatasourcePlugin[ParquetConfigFile]):
         return build_parquet_chunks(context)
 
     def build_context(self, full_type: str, datasource_name: str, file_config: ParquetConfigFile) -> Any:
-        introspection_result = self._introspector.introspect(file_config)
-
-        return introspection_result
+        return self._introspector.introspect(file_config)
 
     def check_connection(self, full_type: str, datasource_name: str, file_config: ParquetConfigFile) -> None:
         self._introspector.check_connection(file_config)

databao_context_engine/project/layout.py

@@ -2,7 +2,6 @@ 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
 
 SOURCE_FOLDER_NAME = "src"
@@ -25,6 +24,14 @@ class ProjectLayout:
     def read_config_file(self) -> ProjectConfig:
         return ProjectConfig.from_file(self.config_file)
 
+    @property
+    def src_dir(self) -> Path:
+        return get_source_dir(self.project_dir)
+
+    @property
+    def output_dir(self) -> Path:
+        return get_output_dir(self.project_dir)
+
 
 def ensure_project_dir(project_dir: Path) -> ProjectLayout:
     return _ProjectValidator(project_dir).ensure_project_dir_valid()
@@ -62,22 +69,14 @@ def get_logs_dir(project_dir: Path) -> Path:
     return project_dir.joinpath(LOGS_FOLDER_NAME)
 
 
-def ensure_datasource_config_file_doesnt_exist(project_dir: Path, datasource_id: DatasourceId) -> Path:
-    config_file = get_source_dir(project_dir).joinpath(datasource_id.relative_path_to_config_file())
-
-    if config_file.is_file():
-        raise ValueError(f"A config file already exists for {str(datasource_id)}")
-
-    return config_file
-
-
 def create_datasource_config_file(
-    project_dir: Path, datasource_id: DatasourceId, config_content: str, overwrite_existing: bool
+    project_layout: ProjectLayout, datasource_relative_name: str, config_content: str, overwrite_existing: bool
 ) -> Path:
+    config_file = project_layout.src_dir / datasource_relative_name
     if not overwrite_existing:
-        ensure_datasource_config_file_doesnt_exist(project_dir, datasource_id)
+        if config_file.is_file():
+            raise ValueError(f"A config file already exists {config_file}")
 
-    config_file = get_source_dir(project_dir).joinpath(datasource_id.relative_path_to_config_file())
     config_file.parent.mkdir(parents=True, exist_ok=True)
 
     config_file.write_text(config_content)

databao_context_engine/retrieve_embeddings/retrieve_runner.py

@@ -1,8 +1,6 @@
 import logging
-from pathlib import Path
 
-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.datasources.types import DatasourceId
 from databao_context_engine.retrieve_embeddings.retrieve_service import RetrieveService
 from databao_context_engine.storage.repositories.vector_search_repository import VectorSearchResult
 
@@ -10,21 +8,10 @@ logger = logging.getLogger(__name__)
 
 
 def retrieve(
-    project_dir: Path,
     *,
     retrieve_service: RetrieveService,
-    project_id: str,
     text: str,
     limit: int | None,
-    export_to_file: bool,
+    datasource_ids: list[DatasourceId] | None = None,
 ) -> list[VectorSearchResult]:
-    retrieve_results = retrieve_service.retrieve(project_id=project_id, text=text, limit=limit)
-
-    if export_to_file:
-        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)
-        logger.info(f"Exported results to {export_file}")
-
-    return retrieve_results
+    return retrieve_service.retrieve(text=text, limit=limit, datasource_ids=datasource_ids)

databao_context_engine/retrieve_embeddings/retrieve_service.py

@@ -1,6 +1,7 @@
 import logging
 from collections.abc import Sequence
 
+from databao_context_engine.datasources.types import DatasourceId
 from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
 from databao_context_engine.services.embedding_shard_resolver import EmbeddingShardResolver
 from databao_context_engine.storage.repositories.vector_search_repository import (
@@ -23,7 +24,9 @@ class RetrieveService:
         self._provider = provider
         self._vector_search_repo = vector_search_repo
 
-    def retrieve(self, *, project_id: str, text: str, limit: int | None = None) -> list[VectorSearchResult]:
+    def retrieve(
+        self, *, text: str, limit: int | None = None, datasource_ids: list[DatasourceId] | None = None
+    ) -> list[VectorSearchResult]:
         if limit is None:
             limit = 10
 
@@ -40,15 +43,19 @@ class RetrieveService:
             retrieve_vec=retrieve_vec,
             dimension=dimension,
             limit=limit,
+            datasource_ids=datasource_ids,
         )
 
         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)
-            logger.debug(f"Best result: ({closest_result.cosine_distance}, {closest_result.embeddable_text})")
-
-            farthest_result = max(search_results, key=lambda result: result.cosine_distance)
-            logger.debug(f"Worst result: ({farthest_result.cosine_distance}, {farthest_result.embeddable_text})")
+            if search_results:
+                closest_result = min(search_results, key=lambda result: result.cosine_distance)
+                logger.debug(f"Best result: ({closest_result.cosine_distance}, {closest_result.embeddable_text})")
+
+                farthest_result = max(search_results, key=lambda result: result.cosine_distance)
+                logger.debug(f"Worst result: ({farthest_result.cosine_distance}, {farthest_result.embeddable_text})")
+            else:
+                logger.debug("No results found")
 
         return search_results

databao_context_engine/retrieve_embeddings/retrieve_wiring.py

@@ -1,8 +1,9 @@
 from duckdb import DuckDBPyConnection
 
+from databao_context_engine.datasources.types import DatasourceId
 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.project.layout import ProjectLayout
 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
@@ -16,21 +17,17 @@ def retrieve_embeddings(
     project_layout: ProjectLayout,
     retrieve_text: str,
     limit: int | None,
-    export_to_file: bool,
+    datasource_ids: list[DatasourceId] | None,
 ) -> 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,
+            datasource_ids=datasource_ids,
         )
 
 

databao_context_engine/serialization/yaml.py

@@ -7,17 +7,17 @@ from yaml import Node, SafeDumper
 def default_representer(dumper: SafeDumper, data: object) -> Node:
     if isinstance(data, Mapping):
         return dumper.represent_dict(data)
-    elif hasattr(data, "__dict__"):
+    if hasattr(data, "__dict__"):
         # 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)
-        else:
-            # If there is no public attributes, we default to the string representation
-            return dumper.represent_str(str(data))
-    else:
+
+        # If there is no public attributes, we default to the string representation
         return dumper.represent_str(str(data))
 
+    return dumper.represent_str(str(data))
+
 
 # Registers our default representer only once, when that file is imported
 yaml.add_multi_representer(object, default_representer, Dumper=SafeDumper)

databao_context_engine/storage/migrate.py

@@ -69,7 +69,7 @@ class _Migration:
 def _create_migration(file: Path) -> _Migration:
     query_bytes = file.read_bytes()
     query = query_bytes.decode("utf-8")
-    checksum = hashlib.md5(query_bytes).hexdigest()
+    checksum = hashlib.md5(query_bytes, usedforsecurity=False).hexdigest()
     version = _extract_version_from_name(file.name)
     return _Migration(name=file.name, version=version, checksum=checksum, query=query)
 

databao_context_engine/storage/repositories/vector_search_repository.py

@@ -1,5 +1,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import Any
 
 import duckdb
 
@@ -23,27 +24,38 @@ class VectorSearchRepository:
         self._conn = conn
 
     def get_display_texts_by_similarity(
-        self, *, table_name: str, retrieve_vec: Sequence[float], dimension: int, limit: int
+        self,
+        *,
+        table_name: str,
+        retrieve_vec: Sequence[float],
+        dimension: int,
+        limit: int,
+        datasource_ids: list[DatasourceId] | None = None,
     ) -> list[VectorSearchResult]:
         """Read only similarity search on a specific embedding shard table."""
+        params: list[Any] = [list(retrieve_vec), self._DEFAULT_DISTANCE_THRESHOLD, list(retrieve_vec), limit]
+        if datasource_ids:
+            params.append([str(datasource_id) for datasource_id in datasource_ids])
+
         rows = self._conn.execute(
             f"""
             SELECT
                 COALESCE(c.display_text, c.embeddable_text) AS display_text,
                 c.embeddable_text,
-                array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) AS cosine_distance,
+                array_cosine_distance(e.vec, CAST($1 AS FLOAT[{dimension}])) AS cosine_distance,
                 c.full_type,
                 c.datasource_id,
             FROM
                 {table_name} e
                 JOIN chunk c ON e.chunk_id = c.chunk_id
             WHERE
-                cosine_distance < ?
+                cosine_distance < $2
+                {"AND c.datasource_id IN $5" if datasource_ids else ""}
             ORDER BY
-                array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) ASC
-            LIMIT ?
+                array_cosine_distance(e.vec, CAST($3 AS FLOAT[{dimension}])) ASC
+            LIMIT $4
             """,
-            [list(retrieve_vec), self._DEFAULT_DISTANCE_THRESHOLD, list(retrieve_vec), limit],
+            params,
         ).fetchall()
 
         return [

databao_context_engine-0.1.6.dist-info/METADATA

@@ -0,0 +1,228 @@
+Metadata-Version: 2.3
+Name: databao-context-engine
+Version: 0.1.6
+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://github.com/JetBrains#jetbrains-on-github)
+[![PyPI version](https://img.shields.io/pypi/v/databao-context-engine.svg)](https://pypi.org/project/databao-context-engine)
+[![License](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](https://github.com/JetBrains/databao-context-engine?tab=License-1-ov-file)
+
+[//]: # ([![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 is a CLI tool that **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             |
+
+## Installation
+
+Databao Context Engine is [available on PyPI](https://pypi.org/project/databao-context-engine/) and can be installed with uv, pip, or another package manage.
+
+### Using uv
+
+1. Install Databao Context Engine:
+
+   ```bash
+   uv tool install databao-context-engine
+   ```
+
+1. Add it to your PATH:
+
+   ```bash
+   uv tool update-shell
+   ```
+
+1. Verify the installation:
+
+   ```bash
+   dce --help
+   ```
+
+### Using pip
+
+1. Install Databao Context Engine:
+
+   ```bash
+   pip install databao-context-engine
+   ```
+
+1. Verify the installation:
+
+   ```bash
+   dce --help
+   ```
+
+##  Supported data sources
+
+* <img src="https://cdn.simpleicons.org/postgresql/316192" width="16" height="16" alt=""> PostgreSQL
+* <img src="https://cdn.simpleicons.org/mysql/4479A1" width="16" height="16" alt=""> MySQL
+* <img src="https://cdn.simpleicons.org/sqlite/003B57" width="16" height="16" alt=""> SQLite
+* <img src="https://cdn.simpleicons.org/duckdb/FFF000" width="16" height="16" alt=""> DuckDB
+* <img src="https://cdn.simpleicons.org/dbt/FF694B" width="16" height="16" alt=""> dbt projects
+* 📄 Documents & spreadsheets *(coming soon)*
+
+##  Supported LLMs
+
+| Provider      | Configuration                                |
+|---------------|----------------------------------------------|
+| **Ollama**    | `languageModel: OLLAMA`: runs locally, free  |
+| **OpenAI**    | `languageModel: OPENAI`: requires an API key |
+| **Anthropic** | `languageModel: CLAUDE`: requires an API key |
+| **Google**    | `languageModel: GEMINI`: requires an API key |
+
+## Quickstart
+
+### 1. Create a project
+
+1. Create a new directory for your project and navigate to it:
+
+   ```bash
+   mkdir dce-project && cd dce-project
+   ```
+
+1. Initialize a new project:
+
+   ```bash
+   dce init
+   ```
+
+### 2. Configure data sources
+
+1. When prompted, agree to create a new datasource.
+   You can also use the `dce datasource add` command.
+
+1. Provide the data source type and its name.
+
+1. Open the config file that was created for you in your editor and fill in the connection details.
+
+1. Repeat these steps for all data sources you want to include in your project.
+
+1. If you have data in Markdown or text files,
+   you can add them to the `dce/src/files` directory.
+
+### 3. Build context
+
+1. To build the context, run the following command:
+
+   ```bash
+   dce build
+   ```
+
+### 4. Use Context with Your LLM
+
+**Option A: Dynamic via MCP Server**
+
+Databao Context Engine exposes the context through a local MCP Server, so your agent can access the latest context at runtime.
+
+1. In **Claude Desktop**, **Cursor**, or another MCP-compatible agent, add the following configuration.
+   Replace `dce-project/` with the path to your project directory:
+  
+   ```json 
+   # claude_desktop_config.json, mcp.json, or similar
+   
+   {
+     "mcpServers": {
+       "dce": {
+         "command": "dce mcp",
+         "args": ["--project-dir", "dce-project/"]
+       }
+     }
+   }
+   ```
+
+1. Save the file and restart your agent.
+
+1. Open a new chat, in the chat window, select the `dce` server, and ask questions related to your project context.
+
+**Option B: Static artifact**
+
+Even if you don’t have Claude or Cursor installed on your local machine,
+you can still use the context built by Databao Context Engine by pasting it directly into your chat with an AI assistant.
+
+1. Navigate to `dce-project/output/` and open the directory with the latest run.
+
+1. Attach the `all_results.yaml` file to your chat with the AI assistant or copy and paste its contents into your chat.
+
+##  Contributing
+
+We’d love your help! Here’s how to get involved:
+
+- ⭐ **Star this repo** — it helps others find us!
+- 🐛 **Found a bug?** [Open an issue](https://github.com/JetBrains/databao-context-engine/issues)
+- 💡 **Have an idea?** We’re all ears — create a feature request
+- 👍 **Upvote issues** you care about — helps us prioritize
+- 🔧 **Submit a PR**
+- 📝 **Improve docs** — typos, examples, tutorials — everything helps!
+
+New to open source? No worries! We're friendly and happy to help you get started. 🌱
+
+For more details, see [CONTRIBUTING](CONTRIBUTING.md).
+
+## 📄 License
+
+Apache 2.0 — use it however you want. See the [LICENSE](LICENSE.md) file for details.
+
+---
+
+<p align="center">
+ <b>Like Databao Context Engine?</b> Give us a ⭐ — it means a lot!
+</p>
+
+<p align="center">
+ <a href="https://databao.app">Website</a> •
+ <a href="https://discord.gg/databao">Discord</a>
+</p>