databao-context-engine 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

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/plugins/dbt/dbt_context_extractor.py

@@ -0,0 +1,106 @@
+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

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/plugins/dbt/types.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+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 DbtColumn:
+    name: str
+    type: str | None = None
+    description: str | 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/plugins/dbt/types_artifacts.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Discriminator, Field
+
+
+class DbtManifestNodeConfig(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: DbtManifestNodeConfig | None = None
+    columns: dict[str, DbtManifestColumn]
+    depends_on: dict[str, list[str]] | None = None
+    primary_key: list[str] | None = None
+
+
+class DbtManifestOtherNode(BaseModel):
+    resource_type: Literal["seed", "analysis", "test", "operation", "sql_operation", "snapshot"]
+
+
+DbtManifestNode = Annotated[DbtManifestModel | 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/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.dist-info → databao_context_engine-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databao-context-engine
-Version: 0.1.5
+Version: 0.1.6
 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.dist-info → databao_context_engine-0.1.6.dist-info}/RECORD

@@ -82,10 +82,16 @@ databao_context_engine/plugins/databases/snowflake/snowflake_introspector.py,sha
 databao_context_engine/plugins/databases/sqlite/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/plugins/databases/sqlite/sqlite_db_plugin.py,sha256=MI896G7Yq8NcQ8sTWErAFD7YK8qfDpZZprxoszJ46l8,460
 databao_context_engine/plugins/databases/sqlite/sqlite_introspector.py,sha256=kKFNzclp1NmQ6BF3ylGn86R5PuQh6aqevj6E7_zAKmQ,8361
+databao_context_engine/plugins/dbt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+databao_context_engine/plugins/dbt/dbt_chunker.py,sha256=5OQo0Y9ouelZV8KQBA0Kqksx72wRhTdRztUOSCKwLZc,1516
+databao_context_engine/plugins/dbt/dbt_context_extractor.py,sha256=8ac2Fx5vv0t_AGqNbx9BXj73l9ZqjDK8r-VAAQukcP0,3748
+databao_context_engine/plugins/dbt/dbt_plugin.py,sha256=-oAVh4u9YZ7j1b6w2K2lwcZT_bOdkN-lL1RRBWgB20k,1026
+databao_context_engine/plugins/dbt/types.py,sha256=yJUVSZa5Ohhx6Znl_20gKwjvgxkBBcxlKrAFFV6tfgY,881
+databao_context_engine/plugins/dbt/types_artifacts.py,sha256=OC2_DL7srcniVqrj9uj5XxNV-7VcA974xqgU1Rzfd5A,1349
 databao_context_engine/plugins/duckdb_tools.py,sha256=46rctnTxDPAhHtaiTp1DxMuuDuRKrtKWJFSSM2w7uUU,645
 databao_context_engine/plugins/files/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/plugins/files/unstructured_files_plugin.py,sha256=eqs1anQhYBZh7xu4CwhfkqXQjGE5gJnKEwyJbtUR78E,2384
-databao_context_engine/plugins/plugin_loader.py,sha256=8j7JAKqtG6_VjX3GbfAD6kjrWuuJ7NnC2nnHPEujux8,4074
+databao_context_engine/plugins/plugin_loader.py,sha256=x5cZ8pUwrEYbadqlPddZvBzU2BpdyBWYI4KlEtffFUY,4159
 databao_context_engine/plugins/resources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/plugins/resources/parquet_chunker.py,sha256=R9WCOBqpKRTVN6t5eeOm_mmnKBOxvjIiQ9zTc8vnUb4,848
 databao_context_engine/plugins/resources/parquet_introspector.py,sha256=Cn_yh6E-dOTOZstlavEGAsV6ZRKZXJraVAl_pzJJuGs,5629
@@ -129,7 +135,7 @@ databao_context_engine/system/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5N
 databao_context_engine/system/properties.py,sha256=mQ7-_PZeYSESYn1cMUQ0IK7rJEnbhc7t4WesFjAgo-Q,429
 databao_context_engine/templating/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/templating/renderer.py,sha256=W2-0IGStAp6oxANmsKs_Z-UoIR6Gt_c4ILYFa3Hruo4,662
-databao_context_engine-0.1.5.dist-info/WHEEL,sha256=5w2T7AS2mz1-rW9CNagNYWRCaB0iQqBMYLwKdlgiR4Q,78
-databao_context_engine-0.1.5.dist-info/entry_points.txt,sha256=5EeQJ1W8zEFh4HuF1bs2zBeoP408oiwuM9UrkJiurgI,138
-databao_context_engine-0.1.5.dist-info/METADATA,sha256=HfJBADmvbXEEK710S6zzybfWl-mqnDXfm_18LIgh7kk,7773
-databao_context_engine-0.1.5.dist-info/RECORD,,
+databao_context_engine-0.1.6.dist-info/WHEEL,sha256=5w2T7AS2mz1-rW9CNagNYWRCaB0iQqBMYLwKdlgiR4Q,78
+databao_context_engine-0.1.6.dist-info/entry_points.txt,sha256=5EeQJ1W8zEFh4HuF1bs2zBeoP408oiwuM9UrkJiurgI,138
+databao_context_engine-0.1.6.dist-info/METADATA,sha256=weN9iS4ZtRnnt3tTJY-epHUgNdpMo-dWD4E3F-8dXyk,7773
+databao_context_engine-0.1.6.dist-info/RECORD,,