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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databao-context-engine
-Version: 0.1.6
+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.6 → databao_context_engine-0.1.7}/pyproject.toml

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

{databao_context_engine-0.1.6 → 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_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.6 → databao_context_engine-0.1.7}/src/databao_context_engine/plugins/dbt/types.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -19,11 +20,39 @@ class DbtMaterialization(str, Enum):
         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)

{databao_context_engine-0.1.6 → databao_context_engine-0.1.7}/src/databao_context_engine/plugins/dbt/types_artifacts.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Discriminator, Field
 
 
-class DbtManifestNodeConfig(BaseModel):
+class DbtManifestModelConfig(BaseModel):
     materialized: str
 
 
@@ -21,17 +21,36 @@ class DbtManifestModel(BaseModel):
     database: str
     schema_: str = Field(alias="schema")
     description: str | None = None
-    config: DbtManifestNodeConfig | 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", "test", "operation", "sql_operation", "snapshot"]
+    resource_type: Literal["seed", "analysis", "operation", "sql_operation", "snapshot"]
 
 
-DbtManifestNode = Annotated[DbtManifestModel | DbtManifestOtherNode, Discriminator("resource_type")]
+DbtManifestNode = Annotated[DbtManifestModel | DbtManifestTest | DbtManifestOtherNode, Discriminator("resource_type")]
 
 
 class DbtManifest(BaseModel):

{databao_context_engine-0.1.6 → 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.6/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)

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

@@ -1,106 +0,0 @@
-from pathlib import Path
-
-from databao_context_engine.plugins.dbt.types import (
-    DbtColumn,
-    DbtConfigFile,
-    DbtContext,
-    DbtMaterialization,
-    DbtModel,
-)
-from databao_context_engine.plugins.dbt.types_artifacts import (
-    DbtArtifacts,
-    DbtCatalog,
-    DbtCatalogColumn,
-    DbtCatalogNode,
-    DbtManifest,
-    DbtManifestColumn,
-    DbtManifestModel,
-)
-
-
-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_model
-        for manifest_model in artifacts.manifest.nodes.values()
-        if isinstance(manifest_model, DbtManifestModel)
-    ]
-
-    catalog_nodes = artifacts.catalog.nodes if artifacts.catalog else {}
-
-    # TODO: Extract the stages? Or at least the "highest-level" models (= marts?)
-    # TODO: Extract the constraints
-    # 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))
-            for manifest_model in manifest_models
-        ],
-    )
-
-
-def _manifest_model_to_dbt_model(manifest_model: DbtManifestModel, catalog_node: DbtCatalogNode | None) -> 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))
-            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
-) -> DbtColumn:
-    return DbtColumn(
-        name=manifest_column.name,
-        description=manifest_column.description,
-        type=catalog_column.type if catalog_column else manifest_column.data_type,
-    )
-
-
-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