databao-context-engine 0.1.5__tar.gz → 0.1.7__tar.gz

{databao_context_engine-0.1.5 → databao_context_engine-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databao-context-engine
-Version: 0.1.5
+Version: 0.1.7
 Summary: Semantic context for your LLMs — generated automatically
 Requires-Dist: click>=8.3.0
 Requires-Dist: duckdb>=1.4.3

{databao_context_engine-0.1.5 → databao_context_engine-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databao-context-engine"
-version = "0.1.5"
+version = "0.1.7"
 description = "Semantic context for your LLMs — generated automatically"
 readme = "README.md"
 requires-python = ">=3.12"

{databao_context_engine-0.1.5 → databao_context_engine-0.1.7}/src/databao_context_engine/__init__.py

@@ -23,7 +23,12 @@ from databao_context_engine.pluginlib.build_plugin import (
     DatasourceType,
 )
 from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
-from databao_context_engine.project.info import DceInfo, DceProjectInfo, get_databao_context_engine_info
+from databao_context_engine.project.info import (
+    DceInfo,
+    DceProjectInfo,
+    get_databao_context_engine_info,
+    get_databao_context_engine_project_info,
+)
 from databao_context_engine.project.init_project import InitErrorReason, InitProjectError
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
@@ -41,6 +46,7 @@ __all__ = [
     "DatasourceConfigFile",
     "DatasourceType",
     "get_databao_context_engine_info",
+    "get_databao_context_engine_project_info",
     "DceInfo",
     "DceProjectInfo",
     "init_dce_project",

databao_context_engine-0.1.7/src/databao_context_engine/cli/info.py

@@ -0,0 +1,36 @@
+import os
+import sys
+from pathlib import Path
+
+import click
+
+from databao_context_engine import DceInfo, DceProjectInfo, get_databao_context_engine_info
+from databao_context_engine.project.info import get_databao_context_engine_project_info
+
+
+def echo_info(project_dir: Path) -> None:
+    dce_info = get_databao_context_engine_info()
+    dce_project_info = get_databao_context_engine_project_info(project_dir=project_dir)
+    click.echo(_generate_info_string(dce_info, dce_project_info))
+
+
+def _generate_info_string(dce_info: DceInfo, project_info: DceProjectInfo) -> str:
+    info_lines = []
+    info_lines.append(f"Databao context engine version: {dce_info.version}")
+    info_lines.append(f"Databao context engine storage dir: {dce_info.dce_path}")
+    info_lines.append(f"Databao context engine plugins: {dce_info.plugin_ids}")
+
+    info_lines.append("")
+
+    info_lines.append(f"OS name: {sys.platform}")
+    info_lines.append(f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}")
+
+    info_lines.append("")
+
+    if project_info.is_initialized:
+        info_lines.append(f"Project dir: {project_info.project_path.resolve()}")
+        info_lines.append(f"Project ID: {str(project_info.project_id)}")
+    else:
+        info_lines.append(f"Project not initialized at {project_info.project_path.resolve()}")
+
+    return os.linesep.join(info_lines)

databao_context_engine-0.1.7/src/databao_context_engine/plugins/dbt/dbt_chunker.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.dbt.types import DbtColumn, DbtContext, DbtModel
+
+
+@dataclass
+class DbtColumnChunkContent:
+    database_name: str
+    schema_name: str
+    model_name: str
+    column: DbtColumn
+
+
+def build_dbt_chunks(context: DbtContext) -> list[EmbeddableChunk]:
+    chunks = []
+
+    for model in context.models:
+        chunks.append(_create_model_chunk(model))
+
+        for column in model.columns:
+            chunks.append(_create_column_chunk(model, column))
+
+    return chunks
+
+
+def _create_model_chunk(model: DbtModel) -> EmbeddableChunk:
+    return EmbeddableChunk(embeddable_text=_build_model_chunk_text(model), content=model)
+
+
+def _build_model_chunk_text(model: DbtModel) -> str:
+    # TODO: Use description and potentially other infos?
+    return f"Model {model.name} in database {model.database} and schema {model.schema}, with unique id {model.id}"
+
+
+def _create_column_chunk(model: DbtModel, column: DbtColumn) -> EmbeddableChunk:
+    return EmbeddableChunk(
+        embeddable_text=_build_column_chunk_text(model, column),
+        content=DbtColumnChunkContent(
+            database_name=model.database, schema_name=model.schema, model_name=model.name, column=column
+        ),
+    )
+
+
+def _build_column_chunk_text(model: DbtModel, column: DbtColumn) -> str:
+    # TODO: Use description and potentially other infos?
+    return f"Column {column.name} in model {model.id}"

databao_context_engine-0.1.7/src/databao_context_engine/plugins/dbt/dbt_context_extractor.py

@@ -0,0 +1,212 @@
+import re
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.plugins.dbt.types import (
+    DbtAcceptedValuesConstraint,
+    DbtColumn,
+    DbtConfigFile,
+    DbtConstraint,
+    DbtContext,
+    DbtMaterialization,
+    DbtModel,
+    DbtRelationshipConstraint,
+    DbtSimpleConstraint,
+)
+from databao_context_engine.plugins.dbt.types_artifacts import (
+    DbtArtifacts,
+    DbtCatalog,
+    DbtCatalogColumn,
+    DbtCatalogNode,
+    DbtManifest,
+    DbtManifestColumn,
+    DbtManifestModel,
+    DbtManifestTest,
+)
+
+
+def check_connection(config_file: DbtConfigFile) -> None:
+    _read_dbt_artifacts(config_file.dbt_target_folder_path.expanduser())
+
+
+def extract_context(config_file: DbtConfigFile) -> DbtContext:
+    artifacts = _read_dbt_artifacts(config_file.dbt_target_folder_path.expanduser())
+
+    return _extract_context_from_artifacts(artifacts)
+
+
+def _read_dbt_artifacts(dbt_target_folder_path: Path) -> DbtArtifacts:
+    if not dbt_target_folder_path.is_dir():
+        raise ValueError(f'Invalid "dbt_target_folder_path": not a directory ({dbt_target_folder_path})')
+
+    # TODO: Check the manifest schema version?
+    manifest_file = dbt_target_folder_path.joinpath("manifest.json")
+    if not manifest_file.is_file():
+        raise ValueError(f'Invalid "dbt_target_folder_path": missing manifest.json file ({manifest_file})')
+
+    manifest = DbtManifest.model_validate_json(manifest_file.read_text())
+
+    catalog_file = dbt_target_folder_path.joinpath("catalog.json")
+    catalog = DbtCatalog.model_validate_json(catalog_file.read_text()) if catalog_file.is_file() else None
+
+    return DbtArtifacts(manifest=manifest, catalog=catalog)
+
+
+def _extract_context_from_artifacts(artifacts: DbtArtifacts) -> DbtContext:
+    manifest_models = [
+        manifest_node
+        for manifest_node in artifacts.manifest.nodes.values()
+        if isinstance(manifest_node, DbtManifestModel)
+    ]
+
+    manifest_tests_by_model_and_column = _get_manifest_tests(artifacts.manifest)
+
+    catalog_nodes = artifacts.catalog.nodes if artifacts.catalog else {}
+
+    # TODO: Extract the stages? Or at least the "highest-level" models (= marts?)
+    # TODO: Organize the models by schemas? Or by stages?
+    return DbtContext(
+        models=[
+            _manifest_model_to_dbt_model(
+                manifest_model,
+                catalog_nodes.get(manifest_model.unique_id, None),
+                manifest_tests_by_model_and_column.get(manifest_model.unique_id, {}),
+            )
+            for manifest_model in manifest_models
+        ],
+    )
+
+
+def _get_manifest_tests(manifest: DbtManifest) -> dict[Any, dict[Any, list]]:
+    """Extract all tests nodes in the manifest and groups them by model and column."""
+    manifest_tests_by_model_and_column: dict[str, dict[str, list[DbtManifestTest]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+    for manifest_node in manifest.nodes.values():
+        if isinstance(manifest_node, DbtManifestTest) and manifest_node.attached_node and manifest_node.column_name:
+            manifest_tests_by_model_and_column[manifest_node.attached_node][manifest_node.column_name].append(
+                manifest_node
+            )
+    return manifest_tests_by_model_and_column
+
+
+def _manifest_model_to_dbt_model(
+    manifest_model: DbtManifestModel,
+    catalog_node: DbtCatalogNode | None,
+    test_nodes_by_column_name: dict[str, list[DbtManifestTest]],
+) -> DbtModel:
+    catalog_columns = catalog_node.columns if catalog_node else {}
+
+    return DbtModel(
+        id=manifest_model.unique_id,
+        name=manifest_model.name,
+        database=manifest_model.database,
+        schema=manifest_model.schema_,
+        description=manifest_model.description,
+        columns=[
+            _manifest_column_to_dbt_column(
+                manifest_column,
+                catalog_columns.get(manifest_column.name),
+                test_nodes_by_column_name.get(manifest_column.name, []),
+            )
+            for manifest_column in manifest_model.columns.values()
+        ],
+        materialization=_manifest_materialization_to_dbt_materializaton(
+            manifest_model.config.materialized if manifest_model.config else None
+        ),
+        primary_key=manifest_model.primary_key,
+        depends_on_nodes=manifest_model.depends_on.get("nodes", []) if manifest_model.depends_on else [],
+    )
+
+
+def _manifest_column_to_dbt_column(
+    manifest_column: DbtManifestColumn, catalog_column: DbtCatalogColumn | None, test_nodes: list[DbtManifestTest]
+) -> DbtColumn:
+    constraints = _manifest_test_to_dbt_constraint(test_nodes)
+
+    return DbtColumn(
+        name=manifest_column.name,
+        description=manifest_column.description,
+        type=catalog_column.type if catalog_column else manifest_column.data_type,
+        constraints=constraints,
+    )
+
+
+def _manifest_test_to_dbt_constraint(test_nodes: list[DbtManifestTest]) -> list[DbtConstraint]:
+    result: list[DbtConstraint] = []
+
+    for test_node in test_nodes:
+        is_enforced = test_node.config.severity == "ERROR" if test_node.config else False
+
+        if test_node.test_metadata is None:
+            continue
+
+        match test_node.test_metadata.name:
+            case "not_null":
+                result.append(
+                    DbtSimpleConstraint(type="not_null", is_enforced=is_enforced, description=test_node.description)
+                )
+            case "unique":
+                result.append(
+                    DbtSimpleConstraint(type="unique", is_enforced=is_enforced, description=test_node.description)
+                )
+            case "accepted_values":
+                if test_node.test_metadata.kwargs is None:
+                    continue
+
+                accepted_values = test_node.test_metadata.kwargs.get("values", None)
+                if accepted_values is None:
+                    continue
+                result.append(
+                    DbtAcceptedValuesConstraint(
+                        type="accepted_values",
+                        is_enforced=is_enforced,
+                        description=test_node.description,
+                        accepted_values=accepted_values,
+                    )
+                )
+            case "relationships":
+                if test_node.test_metadata.kwargs is None:
+                    continue
+
+                target_model = _extract_ref_model(test_node.test_metadata.kwargs.get("to", None))
+                if target_model is None:
+                    continue
+                target_column = test_node.test_metadata.kwargs.get("field", None)
+                if target_column is None:
+                    continue
+                result.append(
+                    DbtRelationshipConstraint(
+                        type="relationships",
+                        is_enforced=is_enforced,
+                        description=test_node.description,
+                        target_model=target_model,
+                        target_column=target_column,
+                    )
+                )
+            case _:
+                continue
+
+    return result
+
+
+def _extract_ref_model(target_model_with_ref: str | None) -> str | None:
+    if target_model_with_ref is None:
+        return None
+
+    match = re.fullmatch(r"ref\(['\"]([\w.]+)['\"]\)", target_model_with_ref)
+    if match:
+        return match.group(1)
+
+    return None
+
+
+def _manifest_materialization_to_dbt_materializaton(materialized: str | None) -> DbtMaterialization | None:
+    if materialized is None:
+        return None
+
+    try:
+        return DbtMaterialization(materialized)
+    except ValueError:
+        return None

databao_context_engine-0.1.7/src/databao_context_engine/plugins/dbt/dbt_plugin.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from databao_context_engine import BuildDatasourcePlugin
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.dbt.dbt_chunker import build_dbt_chunks
+from databao_context_engine.plugins.dbt.dbt_context_extractor import check_connection, extract_context
+from databao_context_engine.plugins.dbt.types import DbtConfigFile
+
+
+class DbtPlugin(BuildDatasourcePlugin[DbtConfigFile]):
+    id = "jetbrains/dbt"
+    name = "Dbt Plugin"
+    config_file_type = DbtConfigFile
+
+    def supported_types(self) -> set[str]:
+        return {"dbt"}
+
+    def build_context(self, full_type: str, datasource_name: str, file_config: DbtConfigFile) -> Any:
+        return extract_context(file_config)
+
+    def check_connection(self, full_type: str, datasource_name: str, file_config: DbtConfigFile) -> None:
+        check_connection(file_config)
+
+    def divide_context_into_chunks(self, context: Any) -> list[EmbeddableChunk]:
+        return build_dbt_chunks(context)

databao_context_engine-0.1.7/src/databao_context_engine/plugins/dbt/types.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class DbtConfigFile(BaseModel):
+    name: str | None = Field(default=None)
+    type: str = Field(default="dbt")
+    dbt_target_folder_path: Path
+
+
+class DbtMaterialization(str, Enum):
+    TABLE = "table"
+    VIEW = "view"
+
+    def __str__(self):
+        return self.value
+
+
+@dataclass(kw_only=True)
+class DbtSimpleConstraint:
+    type: Literal["unique", "not_null"]
+    is_enforced: bool
+    description: str | None = None
+
+
+@dataclass(kw_only=True)
+class DbtAcceptedValuesConstraint:
+    type: Literal["accepted_values"]
+    is_enforced: bool
+    description: str | None = None
+    accepted_values: list[str]
+
+
+@dataclass(kw_only=True)
+class DbtRelationshipConstraint:
+    type: Literal["relationships"]
+    is_enforced: bool
+    description: str | None = None
+    target_model: str
+    target_column: str
+
+
+DbtConstraint = DbtSimpleConstraint | DbtAcceptedValuesConstraint | DbtRelationshipConstraint
+
+
+@dataclass(kw_only=True)
+class DbtColumn:
+    name: str
+    type: str | None = None
+    description: str | None = None
+    constraints: list[DbtConstraint] | None = None
+
+
+@dataclass(kw_only=True)
+class DbtModel:
+    id: str
+    name: str
+    database: str
+    schema: str
+    columns: list[DbtColumn]
+    description: str | None = None
+    materialization: DbtMaterialization | None = None
+    primary_key: list[str] | None = None
+    depends_on_nodes: list[str]
+
+
+@dataclass(kw_only=True)
+class DbtContext:
+    models: list[DbtModel]

databao_context_engine-0.1.7/src/databao_context_engine/plugins/dbt/types_artifacts.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Discriminator, Field
+
+
+class DbtManifestModelConfig(BaseModel):
+    materialized: str
+
+
+class DbtManifestColumn(BaseModel):
+    name: str
+    description: str | None = None
+    data_type: str | None = None
+
+
+class DbtManifestModel(BaseModel):
+    resource_type: Literal["model"]
+    unique_id: str
+    name: str
+    database: str
+    schema_: str = Field(alias="schema")
+    description: str | None = None
+    config: DbtManifestModelConfig | None = None
+    columns: dict[str, DbtManifestColumn]
+    depends_on: dict[str, list[str]] | None = None
+    primary_key: list[str] | None = None
+
+
+class DbtManifestTestConfig(BaseModel):
+    severity: str
+
+
+class DbtManifestTestMetadata(BaseModel):
+    name: str | None = None
+    kwargs: dict[str, Any] | None = None
+
+
+class DbtManifestTest(BaseModel):
+    resource_type: Literal["test"]
+    unique_id: str
+    attached_node: str | None = None
+    column_name: str | None = None
+    description: str | None = None
+    test_metadata: DbtManifestTestMetadata | None = None
+    config: DbtManifestTestConfig | None = None
+
+
+class DbtManifestOtherNode(BaseModel):
+    resource_type: Literal["seed", "analysis", "operation", "sql_operation", "snapshot"]
+
+
+DbtManifestNode = Annotated[DbtManifestModel | DbtManifestTest | DbtManifestOtherNode, Discriminator("resource_type")]
+
+
+class DbtManifest(BaseModel):
+    nodes: dict[str, DbtManifestNode]
+
+
+class DbtCatalogColumn(BaseModel):
+    name: str
+    type: str
+
+
+class DbtCatalogNode(BaseModel):
+    unique_id: str | None = None
+    columns: dict[str, DbtCatalogColumn]
+
+
+class DbtCatalog(BaseModel):
+    nodes: dict[str, DbtCatalogNode]
+
+
+@dataclass(kw_only=True)
+class DbtArtifacts:
+    manifest: DbtManifest
+    catalog: DbtCatalog | None

{databao_context_engine-0.1.5 → databao_context_engine-0.1.7}/src/databao_context_engine/plugins/plugin_loader.py

@@ -45,6 +45,7 @@ def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
     """Statically register built-in plugins."""
     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
@@ -91,7 +92,7 @@ def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
     except ImportError:
         pass
 
-    required_plugins: list[BuildDatasourcePlugin] = [DuckDbPlugin(), ParquetPlugin(), SQLiteDbPlugin()]
+    required_plugins: list[BuildDatasourcePlugin] = [DuckDbPlugin(), ParquetPlugin(), SQLiteDbPlugin(), DbtPlugin()]
     return required_plugins + optional_plugins
 
 

{databao_context_engine-0.1.5 → databao_context_engine-0.1.7}/src/databao_context_engine/project/info.py

@@ -30,33 +30,38 @@ class DceInfo:
     Attributes:
         version: The version of the databao_context_engine package installed on the system.
         dce_path: The path where databao_context_engine stores its global data.
-        project_info: Information about the Databao Context Engine project.
     """
 
     version: str
     dce_path: Path
     plugin_ids: list[str]
 
-    project_info: DceProjectInfo
 
-
-def get_databao_context_engine_info(project_dir: Path) -> DceInfo:
+def get_databao_context_engine_info() -> DceInfo:
     """Return information about the current Databao Context Engine installation and project.
 
-    Args:
-        project_dir: The root directory of the Databao Context Project.
-
     Returns:
-        A DceInfo instance containing information about the Databao Context Engine installation and project.
+        A DceInfo instance containing information about the Databao Context Engine installation.
     """
     return DceInfo(
         version=get_dce_version(),
         dce_path=get_dce_path(),
         plugin_ids=[plugin.id for plugin in load_plugins().values()],
-        project_info=_get_project_info(project_dir),
     )
 
 
+def get_databao_context_engine_project_info(project_dir: Path) -> DceProjectInfo:
+    """Return information about the current Databao Context Engine project.
+
+    Args:
+        project_dir: The root directory of the Databao Context Project.
+
+    Returns:
+        A DceProjectInfo instance containing information about the Databao Context Engine project.
+    """
+    return _get_project_info(project_dir)
+
+
 def _get_project_info(project_dir: Path) -> DceProjectInfo:
     project_layout = validate_project_dir(project_dir)
 

databao_context_engine-0.1.5/src/databao_context_engine/cli/info.py

@@ -1,33 +0,0 @@
-import os
-import sys
-from pathlib import Path
-
-import click
-
-from databao_context_engine import DceInfo, get_databao_context_engine_info
-
-
-def echo_info(project_dir: Path) -> None:
-    click.echo(_generate_info_string(get_databao_context_engine_info(project_dir=project_dir)))
-
-
-def _generate_info_string(command_info: DceInfo) -> str:
-    info_lines = []
-    info_lines.append(f"Databao context engine version: {command_info.version}")
-    info_lines.append(f"Databao context engine storage dir: {command_info.dce_path}")
-    info_lines.append(f"Databao context engine plugins: {command_info.plugin_ids}")
-
-    info_lines.append("")
-
-    info_lines.append(f"OS name: {sys.platform}")
-    info_lines.append(f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}")
-
-    info_lines.append("")
-
-    if command_info.project_info.is_initialized:
-        info_lines.append(f"Project dir: {command_info.project_info.project_path.resolve()}")
-        info_lines.append(f"Project ID: {str(command_info.project_info.project_id)}")
-    else:
-        info_lines.append(f"Project not initialized at {command_info.project_info.project_path.resolve()}")
-
-    return os.linesep.join(info_lines)