databao-context-engine 0.1.3__tar.gz → 0.1.4.dev1__tar.gz

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.3
 Name: databao-context-engine
-Version: 0.1.3
-Summary: Add your description here
+Version: 0.1.4.dev1
+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

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "databao-context-engine"
-version = "0.1.3"
-description = "Add your description here"
+version = "0.1.4.dev1"
+description = "Semantic context for your LLMs — generated automatically"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/build_sources/plugin_execution.py

@@ -2,13 +2,13 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, cast
 
+from databao_context_engine.datasources.types import DatasourceId, PreparedConfig, PreparedDatasource
 from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     BuildFilePlugin,
     BuildPlugin,
 )
 from databao_context_engine.pluginlib.plugin_utils import execute_datasource_plugin, execute_file_plugin
-from databao_context_engine.datasources.types import PreparedConfig, PreparedDatasource, DatasourceId
 
 
 @dataclass()

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/cli/info.py

@@ -15,6 +15,7 @@ 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("")
 

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/introspection/property_extract.py

@@ -1,15 +1,15 @@
 import types
 from dataclasses import MISSING, fields, is_dataclass
-from typing import Annotated, Any, ForwardRef, Union, get_origin, get_type_hints, get_args
+from typing import Annotated, Any, ForwardRef, Union, get_args, get_origin, get_type_hints
 
-from pydantic import BaseModel, _internal
+from pydantic import BaseModel
 from pydantic_core import PydanticUndefinedType
 
 from databao_context_engine.pluginlib.config import (
     ConfigPropertyAnnotation,
     ConfigPropertyDefinition,
-    ConfigUnionPropertyDefinition,
     ConfigSinglePropertyDefinition,
+    ConfigUnionPropertyDefinition,
 )
 
 
@@ -59,18 +59,14 @@ def _get_property_list_from_dataclass(parent_type: type) -> list[ConfigPropertyD
         raise ValueError(f"{parent_type} is not a dataclass")
 
     dataclass_fields = fields(parent_type)
+    type_hints = get_type_hints(parent_type, include_extras=True)
 
     result = []
     for field in dataclass_fields:
         has_field_default = field.default is not None and field.default != MISSING
 
-        if isinstance(field.type, str):
-            try:
-                property_type = _evaluate_type_string(field.type)
-            except Exception:
-                continue
-        else:
-            property_type = field.type
+        # Use the type hints if the field type wasn't resolved (aka. if it is a ForwardRef or a str)
+        property_type = type_hints[field.name] if isinstance(field.type, ForwardRef | str) else field.type
 
         property_for_field = _create_property(
             property_type=property_type,
@@ -89,6 +85,10 @@ def _get_property_list_from_pydantic_base_model(parent_type: type):
     if not issubclass(parent_type, BaseModel):
         raise ValueError(f"{parent_type} is not a Pydantic BaseModel")
 
+    if any(isinstance(field.annotation, ForwardRef) for field in parent_type.model_fields.values()):
+        # If any field's future type wasn't resolved yet, we rebuild the model to resolve them
+        parent_type.model_rebuild(force=True)
+
     pydantic_fields = parent_type.model_fields
     result = []
 
@@ -214,18 +214,3 @@ def compute_default_value(
         return str(property_default)
 
     return None
-
-
-def _evaluate_type_string(property_type: str) -> type:
-    try:
-        # Using a pydantic internal function for this, to avoid having to implement type evaluation manually...
-        return _internal._typing_extra.eval_type(property_type)
-    except Exception as initial_error:
-        try:
-            # Try to convert it ourselves if Pydantic didn't work
-            return ForwardRef(property_type)._evaluate(  # type: ignore[return-value]
-                globalns=globals(), localns=locals(), recursive_guard=frozenset()
-            )
-        except Exception as e:
-            # Ignore if we didn't manage to convert the str to a type
-            raise e from initial_error

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/mcp/mcp_server.py

@@ -1,5 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
+from datetime import date
 from pathlib import Path
 from typing import Literal
 
@@ -7,8 +8,6 @@ from mcp.server import FastMCP
 from mcp.types import ToolAnnotations
 
 from databao_context_engine import DatabaoContextEngine
-from databao_context_engine.mcp.all_results_tool import run_all_results_tool
-from databao_context_engine.mcp.retrieve_tool import run_retrieve_tool
 
 logger = logging.getLogger(__name__)
 
@@ -41,19 +40,23 @@ class McpServer:
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def all_results_tool():
-            return run_all_results_tool(self._databao_context_engine)
+            return self._databao_context_engine.get_all_contexts_formatted()
 
         @mcp.tool(
             description="Retrieve the context built from various resources, including databases, dbt tools, plain and structured files, to retrieve relevant information",
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def retrieve_tool(text: str, limit: int | None):
-            return run_retrieve_tool(
-                databao_context_engine=self._databao_context_engine,
-                text=text,
-                limit=limit or 50,
+            retrieve_results = self._databao_context_engine.search_context(
+                retrieve_text=text, limit=limit, export_to_file=False
             )
 
+            display_results = [context_search_result.context_result for context_search_result in retrieve_results]
+
+            display_results.append(f"\nToday's date is {date.today()}")
+
+            return "\n".join(display_results)
+
         return mcp
 
     def run(self, transport: McpTransport):

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/base_db_plugin.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, TypeVar
+from typing import Annotated, Any, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -8,6 +8,7 @@ from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     EmbeddableChunk,
 )
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector
 from databao_context_engine.plugins.databases.database_chunker import build_database_chunks
 from databao_context_engine.plugins.databases.introspection_scope import IntrospectionScope
@@ -17,7 +18,9 @@ class BaseDatabaseConfigFile(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
     name: str | None = Field(default=None)
     type: str
-    introspection_scope: IntrospectionScope | None = Field(default=None, alias="introspection-scope")
+    introspection_scope: Annotated[
+        IntrospectionScope | None, ConfigPropertyAnnotation(ignored_for_config_wizard=True)
+    ] = Field(default=None, alias="introspection-scope")
 
 
 T = TypeVar("T", bound=BaseDatabaseConfigFile)

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/athena_introspector.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Annotated
+from typing import Annotated, Any
 
 from pyathena import connect
 from pyathena.cursor import DictCursor

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/clickhouse_introspector.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Annotated
+from typing import Annotated, Any
 
 import clickhouse_connect
 from pydantic import BaseModel, Field

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/duckdb_introspector.py

@@ -7,6 +7,7 @@ from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
 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
+from databao_context_engine.plugins.duckdb_tools import fetchall_dicts
 
 
 class DuckDBConfigFile(BaseDatabaseConfigFile):
@@ -319,7 +320,4 @@ class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
 
     def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
         cur = connection.cursor()
-        cur.execute(sql, params or [])
-        columns = [desc[0].lower() for desc in cur.description] if cur.description else []
-        rows = cur.fetchall()
-        return [dict(zip(columns, row)) for row in rows]
+        return fetchall_dicts(cur, sql, params)

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/mssql_introspector.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from typing import Any, Mapping, Annotated
+from typing import Annotated, Any, Mapping
 
 from mssql_python import connect  # type: ignore[import-untyped]
-from pydantic import Field, BaseModel
+from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
 from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/mysql_introspector.py

@@ -1,4 +1,4 @@
-from typing import Any, Annotated
+from typing import Annotated, Any
 
 import pymysql
 from pydantic import BaseModel, Field

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/snowflake_introspector.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Annotated
+from typing import Annotated, Any, Dict, List
 
 import snowflake.connector
 from pydantic import BaseModel, Field

databao_context_engine-0.1.4.dev1/src/databao_context_engine/plugins/duckdb_tools.py

@@ -0,0 +1,18 @@
+from typing import Any
+
+from databao_context_engine.pluginlib.config import DuckDBSecret
+
+
+def generate_create_secret_sql(secret_name, duckdb_secret: DuckDBSecret) -> str:
+    parameters = [("type", duckdb_secret.type)] + list(duckdb_secret.properties.items())
+    return f"""CREATE SECRET {secret_name} (
+    {", ".join([f"{k} '{v}'" for (k, v) in parameters])}
+);
+"""
+
+
+def fetchall_dicts(cur, sql: str, params=None) -> list[dict[str, Any]]:
+    cur.execute(sql, params or [])
+    columns = [desc[0].lower() for desc in cur.description] if cur.description else []
+    rows = cur.fetchall()
+    return [dict(zip(columns, row)) for row in rows]

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/resources/parquet_introspector.py

@@ -10,6 +10,7 @@ from duckdb import DuckDBPyConnection
 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"
 
@@ -49,14 +50,6 @@ class ParquetIntrospectionResult:
     files: list[ParquetFile]
 
 
-def generate_create_secret_sql(secret_name, duckdb_secret: DuckDBSecret) -> str:
-    parameters = [("type", duckdb_secret.type)] + list(duckdb_secret.properties.items())
-    return f"""CREATE SECRET {secret_name} (
-    {", ".join([f"{k} {v}" for (k, v) in parameters])}
-);
-"""
-
-
 @contextlib.contextmanager
 def _create_secret(conn: DuckDBPyConnection, duckdb_secret: DuckDBSecret):
     secret_name = duckdb_secret.name or "gen_secret_" + str(uuid.uuid4()).replace("-", "_")
@@ -97,10 +90,9 @@ class ParquetIntrospector:
         with self._connect(file_config) as conn:
             with conn.cursor() as cur:
                 resolved_url = _resolve_url(file_config)
-                cur.execute(f"SELECT * FROM parquet_file_metadata('{resolved_url}') LIMIT 1")
-                columns = [desc[0].lower() for desc in cur.description] if cur.description else []
-                rows = cur.fetchall()
-                parquet_file_metadata = [dict(zip(columns, row)) for row in rows]
+                parquet_file_metadata = fetchall_dicts(
+                    cur, f"SELECT * FROM parquet_file_metadata('{resolved_url}') LIMIT 1"
+                )
                 if not parquet_file_metadata:
                     raise ValueError(f"No parquet files found by url {resolved_url}")
                 if not parquet_file_metadata or not parquet_file_metadata[0]["file_name"]:
@@ -110,10 +102,7 @@ class ParquetIntrospector:
         with self._connect(file_config) as conn:
             with conn.cursor() as cur:
                 resolved_url = _resolve_url(file_config)
-                cur.execute(f"SELECT * from parquet_metadata('{resolved_url}')")
-                cols = [desc[0].lower() for desc in cur.description] if cur.description else []
-                rows = cur.fetchall()
-                file_metas = [dict(zip(cols, row)) for row in rows]
+                file_metas = fetchall_dicts(cur, f"SELECT * from parquet_metadata('{resolved_url}')")
 
                 columns_per_file: dict[str, dict[int, ParquetColumn]] = defaultdict(defaultdict)
                 for file_meta in file_metas:

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/project/info.py

@@ -3,6 +3,7 @@ from importlib.metadata import version
 from pathlib import Path
 from uuid import UUID
 
+from databao_context_engine.plugins.plugin_loader import load_plugins
 from databao_context_engine.project.layout import validate_project_dir
 from databao_context_engine.system.properties import get_dce_path
 
@@ -34,6 +35,7 @@ class DceInfo:
 
     version: str
     dce_path: Path
+    plugin_ids: list[str]
 
     project_info: DceProjectInfo
 
@@ -50,6 +52,7 @@ def get_databao_context_engine_info(project_dir: Path) -> DceInfo:
     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),
     )
 

databao_context_engine-0.1.3/src/databao_context_engine/mcp/all_results_tool.py

@@ -1,5 +0,0 @@
-from databao_context_engine import DatabaoContextEngine
-
-
-def run_all_results_tool(databao_context_engine: DatabaoContextEngine) -> str:
-    return databao_context_engine.get_all_contexts_formatted()

databao_context_engine-0.1.3/src/databao_context_engine/mcp/retrieve_tool.py

@@ -1,14 +0,0 @@
-import datetime
-
-from databao_context_engine import DatabaoContextEngine
-
-
-def run_retrieve_tool(*, databao_context_engine: DatabaoContextEngine, text: str, limit: int | None = None) -> str:
-    """Execute the retrieve flow for MCP and return the matching display texts."""
-    retrieve_results = databao_context_engine.search_context(retrieve_text=text, limit=limit, export_to_file=False)
-
-    display_results = [context_search_result.context_result for context_search_result in retrieve_results]
-
-    display_results.append(f"\nToday's date is {datetime.date.today()}")
-
-    return "\n".join(display_results)

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/cli/add_datasource_config.py

@@ -6,11 +6,11 @@ from typing import Any
 import click
 
 from databao_context_engine import (
+    ConfigPropertyDefinition,
     DatabaoContextPluginLoader,
     DatabaoContextProjectManager,
     DatasourceId,
     DatasourceType,
-    ConfigPropertyDefinition,
 )
 from databao_context_engine.pluginlib.config import ConfigUnionPropertyDefinition
 

{databao_context_engine-0.1.3 → databao_context_engine-0.1.4.dev1}/src/databao_context_engine/plugins/databases/introspection_model_builder.py

@@ -5,13 +5,13 @@ from databao_context_engine.plugins.databases.databases_types import (
     CheckConstraint,
     DatabaseColumn,
     DatabasePartitionInfo,
+    DatabaseSchema,
     DatabaseTable,
     DatasetKind,
     ForeignKey,
     ForeignKeyColumnMap,
     Index,
     KeyConstraint,
-    DatabaseSchema,
 )