databao-context-engine 0.6.0__tar.gz → 0.6.1.dev1__tar.gz

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databao-context-engine
-Version: 0.6.0
+Version: 0.6.1.dev1
 Summary: Semantic context for your LLMs — generated automatically
 License-Expression: Apache-2.0 AND LicenseRef-Additional-Terms
 License-File: LICENSE.md

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/pyproject.toml

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

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/__init__.py

@@ -55,6 +55,11 @@ from databao_context_engine.plugins.databases.snowflake.config_file import (
     SnowflakeSSOAuth,
 )
 from databao_context_engine.plugins.databases.sqlite.config_file import SQLiteConfigFile, SQLiteConnectionConfig
+from databao_context_engine.plugins.dbt.context_filtering import (
+    DbtContextFilter,
+    DbtContextFilterRule,
+    DbtContextFilterStructuredRule,
+)
 from databao_context_engine.plugins.dbt.types import DbtConfigFile
 from databao_context_engine.plugins.plugin_loader import DatabaoContextPluginLoader
 from databao_context_engine.plugins.resources.types import ParquetConfigFile
@@ -129,5 +134,8 @@ __all__ = [
     "SQLiteConfigFile",
     "SQLiteConnectionConfig",
     "DbtConfigFile",
+    "DbtContextFilter",
+    "DbtContextFilterStructuredRule",
+    "DbtContextFilterRule",
     "ParquetConfigFile",
 ]

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/build_sources/build_runner.py

@@ -3,9 +3,8 @@ import logging
 import databao_context_engine.perf.core as perf
 from databao_context_engine.build_sources.build_service import BuildService
 from databao_context_engine.build_sources.export_results import (
-    append_result_to_all_results,
+    delete_all_results_file,
     export_build_result,
-    reset_all_results,
 )
 from databao_context_engine.build_sources.types import (
     BuildDatasourceResult,
@@ -17,6 +16,7 @@ from databao_context_engine.datasources.datasource_context import (
     read_datasource_type_from_context,
 )
 from databao_context_engine.datasources.datasource_discovery import discover_datasources, prepare_source
+from databao_context_engine.datasources.types import PreparedConfig
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
 from databao_context_engine.plugins.plugin_loader import DatabaoContextPluginLoader
 from databao_context_engine.project.layout import ProjectLayout
@@ -34,7 +34,7 @@ def build(
     project_layout: ProjectLayout,
     plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
-    generate_embeddings: bool = True,
+    should_index: bool,
 ) -> list[BuildDatasourceResult]:
     """Build the context for all datasources in the project.
 
@@ -57,7 +57,7 @@ def build(
     results: list[BuildDatasourceResult] = []
     failed = 0
     skipped = 0
-    reset_all_results(project_layout.output_dir)
+    delete_all_results_file(project_layout)
     for datasource_id in datasource_ids:
         try:
             result = _build_one_datasource(
@@ -65,7 +65,7 @@ def build(
                 plugin_loader=plugin_loader,
                 build_service=build_service,
                 datasource_id=datasource_id,
-                generate_embeddings=generate_embeddings,
+                should_index=should_index,
             )
             results.append(result)
             if result.status == DatasourceStatus.SKIPPED:
@@ -100,9 +100,12 @@ def _build_one_datasource(
     plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     datasource_id,
-    generate_embeddings: bool,
+    should_index: bool,
 ) -> BuildDatasourceResult:
     prepared_source = prepare_source(project_layout, datasource_id)
+    if isinstance(prepared_source, PreparedConfig) and prepared_source.config.get("enabled") is False:
+        logger.info(f"Skipping disabled datasource {prepared_source.datasource_id.datasource_path}")
+        return BuildDatasourceResult(datasource_id=datasource_id, status=DatasourceStatus.SKIPPED)
 
     perf.set_attribute("datasource_type", prepared_source.datasource_type.full_type)
 
@@ -119,10 +122,10 @@ def _build_one_datasource(
         )
         return BuildDatasourceResult(datasource_id=datasource_id, status=DatasourceStatus.SKIPPED)
 
-    result = build_service.process_prepared_source(
+    result = build_service.build_context(
         prepared_source=prepared_source,
         plugin=plugin,
-        generate_embeddings=generate_embeddings,
+        should_index=should_index,
     )
 
     output_dir = project_layout.output_dir
@@ -130,8 +133,6 @@ def _build_one_datasource(
 
     perf.set_attribute("context_size_bytes", context_file_path.stat().st_size)
 
-    append_result_to_all_results(output_dir, result)
-
     return BuildDatasourceResult(
         datasource_id=datasource_id,
         status=DatasourceStatus.OK,

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/build_sources/build_service.py

@@ -30,8 +30,8 @@ class BuildService:
         self._project_layout = project_layout
         self._chunk_embedding_service = chunk_embedding_service
 
-    def process_prepared_source(
-        self, *, prepared_source: PreparedDatasource, plugin: BuildPlugin, generate_embeddings: bool = True
+    def build_context(
+        self, *, prepared_source: PreparedDatasource, plugin: BuildPlugin, should_index: bool
     ) -> BuiltDatasourceContext:
         """Process a single source to build its context.
 
@@ -44,23 +44,10 @@ class BuildService:
         """
         result = self._execute_plugin(prepared_source=prepared_source, plugin=plugin)
 
-        if not generate_embeddings:
+        if not should_index:
             return result
 
-        chunks = plugin.divide_context_into_chunks(result.context)
-
-        perf.set_attribute("chunk_count", len(chunks))
-
-        if not chunks:
-            logger.info("No chunks for %s — skipping.", prepared_source.datasource_id.relative_path_to_config_file())
-            return result
-
-        self._chunk_embedding_service.embed_chunks(
-            chunks=chunks,
-            result=result,
-            full_type=prepared_source.datasource_type.full_type,
-            datasource_id=result.datasource_id,
-        )
+        self._index_context(built_context=result, plugin=plugin)
 
         return result
 
@@ -78,21 +65,24 @@ class BuildService:
         """
         built = self._deserialize_built_context(context=context, context_type=plugin.context_type)
 
-        chunks = plugin.divide_context_into_chunks(built.context)
+        self._index_context(built_context=built, plugin=plugin, override=True)
+
+    def _index_context(
+        self, *, built_context: BuiltDatasourceContext, plugin: BuildPlugin, override: bool = False
+    ) -> None:
+        chunks = plugin.divide_context_into_chunks(built_context.context)
         perf.set_attribute("chunk_count", len(chunks))
 
         if not chunks:
-            logger.info(
-                "No chunks for %s — skipping indexing.", context.datasource_id.relative_path_to_context_file().name
-            )
+            logger.info("No chunks for %s — skipping indexing.", built_context.datasource_id)
             return
 
         self._chunk_embedding_service.embed_chunks(
             chunks=chunks,
-            result=built,
-            full_type=built.datasource_type,
-            datasource_id=built.datasource_id,
-            override=True,
+            result=built_context,
+            full_type=built_context.datasource_type,
+            datasource_id=built_context.datasource_id,
+            override=override,
         )
 
     def _deserialize_built_context(

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/build_sources/build_wiring.py

@@ -9,8 +9,6 @@ from databao_context_engine.build_sources.build_runner import (
 from databao_context_engine.build_sources.build_service import BuildService
 from databao_context_engine.build_sources.types import BuildDatasourceResult, IndexDatasourceResult
 from databao_context_engine.datasources.datasource_context import DatasourceContext
-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,
@@ -30,7 +28,7 @@ def build_all_datasources(
     project_layout: ProjectLayout,
     plugin_loader: DatabaoContextPluginLoader,
     chunk_embedding_mode: ChunkEmbeddingMode,
-    generate_embeddings: bool = True,
+    should_index: bool,
 ) -> list[BuildDatasourceResult]:
     """Build the context for all datasources in the project.
 
@@ -48,31 +46,21 @@ def build_all_datasources(
     db_path = project_layout.db_path
     db_path.parent.mkdir(parents=True, exist_ok=True)
     if db_path.exists():
+        logger.debug("Deleting existing database %s", db_path)
         db_path.unlink()
 
     migrate(db_path)
     with open_duckdb_connection(db_path) as conn:
-        ollama_service = create_ollama_service()
-        embedding_provider = create_ollama_embedding_provider(
-            ollama_service, model_details=project_layout.project_config.ollama_embedding_model_details
-        )
-        description_provider = (
-            create_ollama_description_provider(ollama_service)
-            if chunk_embedding_mode.should_generate_description()
-            else None
-        )
         build_service = _create_build_service(
             conn,
             project_layout=project_layout,
-            embedding_provider=embedding_provider,
-            description_provider=description_provider,
             chunk_embedding_mode=chunk_embedding_mode,
         )
         return build(
             project_layout=project_layout,
             plugin_loader=plugin_loader,
             build_service=build_service,
-            generate_embeddings=generate_embeddings,
+            should_index=should_index,
         )
 
 
@@ -98,21 +86,9 @@ def index_built_contexts(
         migrate(db_path)
 
     with open_duckdb_connection(db_path) as conn:
-        ollama_service = create_ollama_service()
-        embedding_provider = create_ollama_embedding_provider(
-            ollama_service, model_details=project_layout.project_config.ollama_embedding_model_details
-        )
-        description_provider = (
-            create_ollama_description_provider(ollama_service)
-            if chunk_embedding_mode.should_generate_description()
-            else None
-        )
-
         build_service = _create_build_service(
             conn,
             project_layout=project_layout,
-            embedding_provider=embedding_provider,
-            description_provider=description_provider,
             chunk_embedding_mode=chunk_embedding_mode,
         )
         return run_indexing(
@@ -124,10 +100,18 @@ def _create_build_service(
     conn: DuckDBPyConnection,
     *,
     project_layout: ProjectLayout,
-    embedding_provider: EmbeddingProvider,
-    description_provider: DescriptionProvider | None,
     chunk_embedding_mode: ChunkEmbeddingMode,
 ) -> BuildService:
+    ollama_service = create_ollama_service()
+    embedding_provider = create_ollama_embedding_provider(
+        ollama_service, model_details=project_layout.project_config.ollama_embedding_model_details
+    )
+    description_provider = (
+        create_ollama_description_provider(ollama_service)
+        if chunk_embedding_mode.should_generate_description()
+        else None
+    )
+
     chunk_embedding_service = create_chunk_embedding_service(
         conn,
         embedding_provider=embedding_provider,

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/build_sources/export_results.py

@@ -2,9 +2,8 @@ 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.project.layout import DEPRECATED_ALL_RESULTS_FILE_NAME, ProjectLayout
 from databao_context_engine.serialization.yaml import write_yaml_to_stream
 
 logger = logging.getLogger(__name__)
@@ -25,17 +24,13 @@ def export_build_result(output_dir: Path, result: BuiltDatasourceContext) -> Pat
     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")
+def delete_all_results_file(project_layout: ProjectLayout) -> None:
+    """Deletes the all_results.yaml file we were previously writing.
 
+    We're keeping this method for now, to make sure that older projects
+    don't keep an outdated "all_results.yaml" file in their output folder.
 
-# 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)
+    We should be able to remove this code once we don't expect projects with an "all_results.yaml" file to exist anymore.
+    """
+    path = project_layout.output_dir.joinpath(DEPRECATED_ALL_RESULTS_FILE_NAME)
     path.unlink(missing_ok=True)

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/databao_context_domain_manager.py

@@ -97,7 +97,7 @@ class DatabaoContextDomainManager:
             project_layout=self._project_layout,
             plugin_loader=self._plugin_loader,
             chunk_embedding_mode=chunk_embedding_mode,
-            generate_embeddings=should_index,
+            should_index=should_index,
         )
 
     def index_built_contexts(
@@ -118,11 +118,9 @@ class DatabaoContextDomainManager:
             The summary of the index operation.
         """
         engine: DatabaoContextEngine = self.get_engine_for_domain()
-        contexts: list[DatasourceContext] = engine.get_all_contexts()
-
-        if datasource_ids is not None:
-            wanted_paths = {d.datasource_path for d in datasource_ids}
-            contexts = [c for c in contexts if c.datasource_id.datasource_path in wanted_paths]
+        contexts: list[DatasourceContext] = (
+            engine.get_all_contexts() if datasource_ids is None else engine.get_datasource_contexts(datasource_ids)
+        )
 
         return index_built_contexts(
             project_layout=self._project_layout,

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/databao_context_engine.py

@@ -1,13 +1,14 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, Collection
 
 from databao_context_engine.datasources.datasource_context import (
     DatasourceContext,
     get_all_contexts,
     get_context_header_for_datasource,
     get_datasource_context,
+    get_datasource_contexts,
     get_introspected_datasource_list,
 )
 from databao_context_engine.datasources.execute_sql_query import run_sql
@@ -81,9 +82,20 @@ class DatabaoContextEngine:
 
         Returns:
             The context for this datasource.
-        """
+        """  # noqa: DOC501
         return get_datasource_context(project_layout=self._project_layout, datasource_id=datasource_id)
 
+    def get_datasource_contexts(self, datasource_ids: Collection[DatasourceId]) -> list[DatasourceContext]:
+        """Return the context available for a given list of datasources.
+
+        Args:
+            datasource_ids: The list of datasources IDs to get the context for.
+
+        Returns:
+            The context for those datasources.
+        """
+        return get_datasource_contexts(project_layout=self._project_layout, datasource_ids=datasource_ids)
+
     def get_all_contexts(self) -> list[DatasourceContext]:
         """Return all contexts generated in the domain.
 

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/datasources/config_wizard.py

@@ -37,7 +37,7 @@ def _build_config_content_from_properties(
 ) -> dict[str, Any]:
     config_content: dict[str, Any] = {}
     for config_file_property in properties:
-        if config_file_property.property_key in ["type", "name"] and len(properties_prefix) == 0:
+        if config_file_property.property_key in ["type", "name", "enabled"] and len(properties_prefix) == 0:
             # We ignore type and name properties as they've already been filled
             continue
         if in_union and config_file_property.property_key == "type":

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/datasources/datasource_context.py

@@ -2,12 +2,12 @@ import logging
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Iterable
+from typing import Collection, Iterable
 
 import yaml
 
 from databao_context_engine.datasources.types import Datasource, DatasourceId, DatasourceType
-from databao_context_engine.project.layout import ALL_RESULTS_FILE_NAME, ProjectLayout
+from databao_context_engine.project.layout import DEPRECATED_ALL_RESULTS_FILE_NAME, ProjectLayout
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +55,7 @@ def _get_datasources_with_context(project_layout: ProjectLayout) -> list[Datasou
             relative_context_file = context_file.relative_to(project_layout.output_dir)
 
             if context_file.suffix not in DatasourceId.ALLOWED_YAML_SUFFIXES or (
-                len(relative_context_file.parts) == 1 and context_file_name == ALL_RESULTS_FILE_NAME
+                len(relative_context_file.parts) == 1 and context_file_name == DEPRECATED_ALL_RESULTS_FILE_NAME
             ):
                 continue
 
@@ -85,6 +85,12 @@ def get_introspected_datasource_list(project_layout: ProjectLayout) -> list[Data
     return result
 
 
+def get_datasource_contexts(
+    project_layout: ProjectLayout, datasource_ids: Collection[DatasourceId]
+) -> list[DatasourceContext]:
+    return [get_datasource_context(project_layout, datasource_id) for datasource_id in datasource_ids]
+
+
 def get_datasource_context(project_layout: ProjectLayout, datasource_id: DatasourceId) -> DatasourceContext:
     context_path = datasource_id.absolute_path_to_context_file(project_layout)
     if not context_path.is_file():

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/introspection/property_extract.py

@@ -143,7 +143,7 @@ def _create_property(
 
     actual_property_types = _read_actual_property_type(property_type)
 
-    required = annotation.required if annotation else is_property_required
+    required = annotation.required if annotation and annotation.required is not None else is_property_required
     secret = annotation.secret if annotation else False
 
     if len(actual_property_types) > 1:
@@ -175,7 +175,6 @@ def _create_property(
 
     resolved_type = actual_property_type if not nested_properties else None
     default_value = compute_default_value(
-        annotation=annotation,
         property_default=property_default,
         has_nested_properties=nested_properties is not None and len(nested_properties) > 0,
     )
@@ -223,15 +222,10 @@ def _read_actual_property_type(property_type: type) -> tuple[type, ...]:
     return (property_type,)
 
 
-def compute_default_value(
-    *, annotation: ConfigPropertyAnnotation | None, property_default: Any | None = None, has_nested_properties: bool
-) -> str | None:
+def compute_default_value(*, property_default: Any | None = None, has_nested_properties: bool) -> str | None:
     if has_nested_properties:
         return None
 
-    if annotation is not None and annotation.default_value is not None:
-        return str(annotation.default_value)
-
     if property_default is not None:
         return str(property_default)
 

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/mcp/mcp_server.py

@@ -37,7 +37,7 @@ class McpServer:
         mcp = FastMCP(host=host or "127.0.0.1", port=port or 8000, lifespan=mcp_server_lifespan)
 
         @mcp.tool(
-            description="Retrieve the contents of the all_results file",
+            description="Read all available contexts",
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def all_results_tool():

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/pluginlib/build_plugin.py

@@ -181,6 +181,7 @@ class DatasourceType:
 class AbstractConfigFile(ABC):
     type: str
     name: str
+    enabled: bool = True
 
 
 # Config files can either be defined:

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/pluginlib/config.py

@@ -41,8 +41,7 @@ ConfigPropertyDefinition = ConfigSinglePropertyDefinition | ConfigUnionPropertyD
 
 @dataclass(kw_only=True)
 class ConfigPropertyAnnotation:
-    required: bool = False
-    default_value: str | None = None
+    required: bool | None = None
     ignored_for_config_wizard: bool = False
     secret: bool = False
 

{databao_context_engine-0.6.0 → databao_context_engine-0.6.1.dev1}/src/databao_context_engine/plugins/databases/athena/athena_introspector.py

@@ -4,11 +4,10 @@ from typing import Any
 
 from pyathena import connect
 from pyathena.cursor import DictCursor
+from typing_extensions import override
 
 from databao_context_engine.plugins.databases.athena.config_file import AthenaConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
-from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
-from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
 
 
 class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
@@ -36,36 +35,12 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
         sql = f"SELECT schema_name, catalog_name FROM {catalog}.information_schema.schemata"
         return SQLQuery(sql, None)
 
-    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
-        if not schemas:
-            return []
-
-        comps = self._component_queries(catalog, schemas)
-        results: dict[str, list[dict]] = {}
-
-        for name, q in comps.items():
-            results[name] = self._fetchall_dicts(connection, q, None)
-
-        return IntrospectionModelBuilder.build_schemas_from_components(
-            schemas=schemas,
-            rels=results.get("relations", []),
-            cols=results.get("columns", []),
-            pk_cols=[],
-            uq_cols=[],
-            checks=[],
-            fk_cols=[],
-            idx_cols=[],
-        )
-
-    def _component_queries(self, catalog: str, schemas: list[str]) -> dict[str, str]:
+    @override
+    def get_relations_sql_query(self, catalog: str, schemas: list[str]) -> SQLQuery:
         schemas_in = ", ".join(self._quote_literal(s) for s in schemas)
-        return {
-            "relations": self._sql_relations(catalog, schemas_in),
-            "columns": self._sql_columns(catalog, schemas_in),
-        }
 
-    def _sql_relations(self, catalog: str, schemas_in: str) -> str:
-        return f"""
+        return SQLQuery(
+            sql=f"""
             SELECT
                 table_schema AS schema_name,
                 table_name,
@@ -80,9 +55,14 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
             WHERE 
                 table_schema IN ({schemas_in})
         """
+        )
 
-    def _sql_columns(self, catalog: str, schemas_in: str) -> str:
-        return f"""
+    @override
+    def get_columns_sql_query(self, catalog: str, schemas: list[str]) -> SQLQuery:
+        schemas_in = ", ".join(self._quote_literal(s) for s in schemas)
+
+        return SQLQuery(
+            sql=f"""
         SELECT 
             table_schema AS schema_name,
             table_name, 
@@ -99,6 +79,7 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
             table_name,
             ordinal_position
         """
+        )
 
     def _resolve_pseudo_catalog_name(self, file_config: AthenaConfigFile) -> str:
         return "awsdatacatalog"