diff --git a/demos/agent_mcp_app/.gitignore b/demos/agent_mcp_app/.gitignore new file mode 100644 index 00000000..2085cdf5 --- /dev/null +++ b/demos/agent_mcp_app/.gitignore @@ -0,0 +1,9 @@ + +.databricks + +# Synced/built by build_wheel.sh from the repo's own src/ -- not checked in, +# since they'd otherwise drift from the repo's actual source. Run +# build_wheel.sh before deploying (see README.md). +impulse_reporting/ +impulse_query_engine/ +wheels/*.whl diff --git a/demos/agent_mcp_app/README.md b/demos/agent_mcp_app/README.md new file mode 100644 index 00000000..bcb169e9 --- /dev/null +++ b/demos/agent_mcp_app/README.md @@ -0,0 +1,85 @@ +# mcp-impulse-agent + +A custom MCP server, hosted as a Databricks App, exposing ad-hoc Impulse +queries as agent tools: `list_channels`, `list_containers`, +`preview_histogram`, `preview_histogram_2d`, `preview_stats`, and +`preview_point_values`. See `demos/agent_mcp_query` for the notebook this +was built from, including a full write-up of the design decisions (why not +Genie/managed MCP, the safety model for virtual-signal expression trees, and +the latency work that got steady-state calls down to 2-7s). + +## Prerequisites + +- A Databricks workspace with Unity Catalog and serverless compute enabled. +- An already-loaded Impulse silver layer to point this at. +- `databricks` CLI authenticated to your workspace. + +## Setup + +1. **Build the wheel** (ships this repo's `impulse_reporting`/ + `impulse_query_engine` to the remote serverless workers TSAL compiles + Python UDFs onto, and to the app's own process): + ```bash + ./build_wheel.sh + ``` + Re-run this whenever `src/` changes -- it always resyncs from the + canonical source, so the bundled copies here never drift. + +2. **Configure `app.yaml`**: replace the `CATALOG`/`SCHEMA`/`TABLE_PREFIX` + placeholders with wherever your silver layer lives. + +3. **Create and deploy the app**: + ```bash + databricks apps create mcp-impulse-agent # name must start with mcp- + databricks sync . /Workspace/Users//mcp-impulse-agent + databricks apps deploy mcp-impulse-agent \ + --source-code-path /Workspace/Users//mcp-impulse-agent + ``` + +4. **Grant the app's service principal Unity Catalog access** (client ID + from `databricks apps get mcp-impulse-agent`): + ```sql + GRANT USE CATALOG, BROWSE ON CATALOG TO ``; + GRANT USE SCHEMA, SELECT, MODIFY, CREATE TABLE ON SCHEMA . TO ``; + ``` + +5. **Test it**: the app is automatically discoverable as a custom MCP + server in AI Playground (Workspace sidebar → Playground → Tools + dropdown), since its name starts with `mcp-`. + +## Implementation notes + +- **Dependency pins matter.** `databricks-connect==18.2.*` is required -- + `17.0.1` has a pyspark incompatibility and `18.3.1` is unsupported on + serverless compute. `databricks-sdk==0.106.0` must match Impulse's own + pin, since its telemetry code reads a private `Config._product_info` + attribute that only exists on that version. Don't pin `pyspark`/ + `delta-spark` separately -- let `databricks-connect` provide them. +- **Why the wheel exists at all:** TSAL compiles to Python UDFs that run on + remote serverless workers, not in the app's own process, so + `impulse_reporting`/`impulse_query_engine` have to be shipped there + explicitly via `DatabricksEnv().withDependencies("local:")` -- + see `server/main.py`'s `get_spark()`. +- **Latency:** AI Playground enforces a ~55s per-call timeout. Persisting + results to Delta and reading them back took 33-43s per call -- too slow. + Reading `report.aggregation_dfs[...]` directly (already in final fact + schema, no join needed) plus going sinkless (`unity_sink` is `Optional` + on Impulse's `Report` config) cut steady-state latency to 2-7s. + +## Known limitations + +- **No per-user (on-behalf-of-user) authorization.** All callers share this + app's single service principal's Unity Catalog permissions. On-behalf-of- + user auth was attempted and hit a platform inconsistency between the + Databricks Apps `user_api_scopes` API and the Databricks Connect runtime's + actual scope requirements -- parked pending further investigation. +- **Distance/custom-weighted histograms are not exposed.** Impulse supports + weighting histograms by distance or an arbitrary signal + (`HistogramDistance`/`HistogramCustomWeights`), not just duration. + Verification found these produce numerically incorrect results (off by + several orders of magnitude) when the weight signal is derived via + `resample()`+`cumtrapz()` -- traced to the `synchronized()`+`diff()` + interaction inside `HistogramCustomWeights.build()` in the query engine, + not something fixable from this MCP layer. `preview_histogram`/ + `preview_histogram_2d` raise a clear error if you ask for anything other + than duration weighting, rather than silently returning wrong numbers. diff --git a/demos/agent_mcp_app/app.yaml b/demos/agent_mcp_app/app.yaml new file mode 100644 index 00000000..11bb4fe6 --- /dev/null +++ b/demos/agent_mcp_app/app.yaml @@ -0,0 +1,12 @@ +command: ['uv', 'run', 'mcp-server'] +# CATALOG/SCHEMA/TABLE_PREFIX must point at an already-loaded Impulse silver +# layer. Replace these placeholders before deploying, or override them via +# the app's "Environment variables" configuration in the Databricks Apps UI +# without editing this file. +env: + - name: CATALOG + value: "" + - name: SCHEMA + value: "" + - name: TABLE_PREFIX + value: "" diff --git a/demos/agent_mcp_app/build_wheel.sh b/demos/agent_mcp_app/build_wheel.sh new file mode 100755 index 00000000..0ec27f7c --- /dev/null +++ b/demos/agent_mcp_app/build_wheel.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Syncs impulse_reporting/impulse_query_engine from this repo's src/ (so the +# app's own process can import them -- Databricks Apps only deploy the +# source-code-path directory itself, not sibling paths, so these can't just +# be sys.path-referenced at runtime) and builds a wheel from the same source +# into ./wheels/ (so the app can also ship this code to remote serverless +# workers -- TSAL compiles to Python UDFs that run there, not in the app's +# own process; see server/main.py's get_spark()). +# +# Run this once before `databricks sync` + `databricks apps deploy`, and +# again whenever the repo's src/ changes -- it always overwrites the local +# copies, so they can't silently drift from the canonical source. Assumes +# this directory lives two levels under the repo root (e.g. +# demos/agent_mcp_app/) -- adjust REPO_ROOT if you've moved it. +set -euo pipefail +cd "$(dirname "$0")" +REPO_ROOT="../.." + +echo "Syncing impulse_reporting/impulse_query_engine from $REPO_ROOT/src..." +rm -rf impulse_reporting impulse_query_engine +cp -r "$REPO_ROOT/src/impulse_reporting" . +cp -r "$REPO_ROOT/src/impulse_query_engine" . + +echo "Building wheel..." +rm -f wheels/databricks_impulse-*.whl +mkdir -p wheels +python3 -m pip install -q build +python3 -m build --wheel --outdir wheels "$REPO_ROOT" +echo "Built: $(ls wheels/databricks_impulse-*.whl)" diff --git a/demos/agent_mcp_app/pyproject.toml b/demos/agent_mcp_app/pyproject.toml new file mode 100644 index 00000000..eb33f10d --- /dev/null +++ b/demos/agent_mcp_app/pyproject.toml @@ -0,0 +1,29 @@ +[project] +name = "mcp-impulse-agent" +version = "0.1.0" +requires-python = "==3.12.*" +dependencies = [ + "mcp>=1.0.0", + "databricks-connect==18.2.*", + "databricks-sdk==0.106.0", + "lz4==4.4.5", + "nptyping==2.5.0", + "pandas==2.2.3", + "pyarrow==19.0.1", + "pydantic==2.11.7", + "scipy==1.15.1", + "zstandard>=0.25.0", + "uvicorn>=0.30.0", + "starlette>=0.37.0", +] + +[project.scripts] +mcp-server = "server.main:main" + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["."] +include = ["server*", "impulse_reporting*", "impulse_query_engine*"] diff --git a/demos/agent_mcp_app/requirements.txt b/demos/agent_mcp_app/requirements.txt new file mode 100644 index 00000000..60cc5e6a --- /dev/null +++ b/demos/agent_mcp_app/requirements.txt @@ -0,0 +1 @@ +uv diff --git a/demos/agent_mcp_app/server/__init__.py b/demos/agent_mcp_app/server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/demos/agent_mcp_app/server/main.py b/demos/agent_mcp_app/server/main.py new file mode 100644 index 00000000..cdafb270 --- /dev/null +++ b/demos/agent_mcp_app/server/main.py @@ -0,0 +1,700 @@ +import operator +import os +import uuid +from typing import Literal + +import pyspark.sql.functions as F +from mcp.server.fastmcp import FastMCP + +_spark = None + + +def get_spark(): + """Lazily create a serverless Spark session via Databricks Connect. + + Impulse's TSAL layer compiles expressions into Python UDFs that run on + the remote serverless workers, not in this app's own process -- so the + bundled impulse_query_engine/impulse_reporting source (not a PyPI + package) has to be shipped to those workers too via a wheel built from + the same source (see wheels/), referenced through withDependencies. + """ + global _spark + if _spark is None: + import glob + from databricks.connect import DatabricksEnv, DatabricksSession + wheels_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "wheels") + matches = glob.glob(os.path.join(wheels_dir, "databricks_impulse-*.whl")) + if not matches: + raise RuntimeError( + f"No databricks_impulse-*.whl found in {wheels_dir} -- run build_wheel.sh " + "(see README) before deploying this app." + ) + env = DatabricksEnv().withDependencies(f"local:{matches[0]}") + _spark = DatabricksSession.builder.serverless().withEnvironment(env).getOrCreate() + return _spark + + +CATALOG = os.environ["CATALOG"] +SCHEMA = os.environ["SCHEMA"] +TABLE_PREFIX = os.environ["TABLE_PREFIX"] +PFX = f"{CATALOG}.{SCHEMA}.{TABLE_PREFIX}" + +PORT = int(os.environ.get("DATABRICKS_APP_PORT", "8000")) +mcp = FastMCP("impulse-agent", host="0.0.0.0", port=PORT) + + +def _adhoc_config() -> dict: + return { + "source": { + "container_metrics_table": f"{PFX}_container_metrics", + "channel_metrics_table": f"{PFX}_channel_metrics", + "channels_uri": f"{PFX}_channels", + "container_tags_table": f"{PFX}_container_tags", + "channel_tags_table": f"{PFX}_channel_tags", + }, + # No unity_sink: results are read straight from report.aggregation_dfs + # in memory (see each tool below) and never persisted -- going + # sinkless also skips several per-call Unity Catalog round trips + # (gold-layer-exists checks, temp-table cleanup scans) that only + # matter for persistence. + "query_engine": {"solver": "DefaultSolver", "data_type": "RAW"}, + "measurement_dimensions": ["container_id", "vehicle_key", "start_ts", "stop_ts"], + } + + +def _new_report(name_prefix: str): + from databricks.sdk import WorkspaceClient + from impulse_reporting.core.report import Report + + report = Report( + name=f"{name_prefix}_{uuid.uuid4().hex[:8]}", + spark=get_spark(), config=_adhoc_config(), workspace_client=WorkspaceClient(), + ) + return report, report.get_db() + + +def _aggregation_df(report, agg_type: str): + """Read an aggregation's result straight from memory (see _adhoc_config).""" + dfs = report.aggregation_dfs[agg_type] + return dfs.get("changed") if dfs.get("changed") is not None else dfs["unchanged"] + + +# --------------------------------------------------------------------------- +# Virtual signal expression trees +# --------------------------------------------------------------------------- + +_OPS = { + ">": operator.gt, + "<": operator.lt, + ">=": operator.ge, + "<=": operator.le, + "==": operator.eq, +} + +# Whitelist-only binary operators for expression trees. Every node maps to a +# fixed, known Python operator applied to TSAL objects -- there is no +# eval()/exec() of user-provided strings anywhere. +_EXPR_OPS = { + "add": operator.add, + "sub": operator.sub, + "mul": operator.mul, + "div": operator.truediv, + "gt": operator.gt, + "lt": operator.lt, + "ge": operator.ge, + "le": operator.le, + "eq": operator.eq, + "ne": operator.ne, + "and": operator.and_, + "or": operator.or_, +} + +# Whitelist-only unary TSAL methods (SampleSeries API) reachable from an +# expression tree's "method" node. Same safety property as _EXPR_OPS: only +# names literally in this set are ever passed to getattr(), so there's no way +# to reach an unintended method (e.g. a dunder) even in principle. +_EXPR_METHODS = { + # SampleSeries (raw/derived numeric signals) + "resample", "cumtrapz", "diff", "where", + "rising_edges", "falling_edges", "rising_edge", "falling_edge", + "intervals_between_falling_edges", "rolling_average", + # Intervals (comparison-derived conditions, e.g. channel > threshold) -- + # start_points/end_points are how you turn "RPM is above 2000" into + # discrete instants ("each time RPM crosses above/below 2000"), as + # opposed to rising_edges/falling_edges above, which detect transitions + # in a raw signal that's already boolean-valued (e.g. a brake switch). + "start_points", "end_points", +} + +_MAX_EXPR_DEPTH = 10 + + +def _resolve_arg(arg, db, depth): + """A method node's args/kwargs may be a literal (number/bool/string) or + a nested expression node -- resolve the latter recursively.""" + if isinstance(arg, dict) and ({"channel", "const", "op", "method"} & arg.keys()): + return build_expr(arg, db, depth) + return arg + + +def build_expr(node: dict, db, _depth: int = 0): + """Recursively build a TSAL expression from a constrained JSON tree. + + Node shapes: + {"channel": "", "tags": {...}} -- leaf: channel reference + {"const": } -- leaf: constant + {"op": "", "left": , "right": } -- one of _EXPR_OPS (binary) + {"method": "", "operand": , + "args": [...], "kwargs": {...}} -- one of _EXPR_METHODS (unary, + called on the built operand) + + Every op/method maps to a fixed, known Python operator or TSAL method on + the same objects list_channels/db.query.channel already produce -- never + to arbitrary generated code. args/kwargs may themselves be nested + expression nodes (e.g. the condition passed to a "where" method) or plain + literals (e.g. the sample_rate passed to "resample"). + """ + if _depth > _MAX_EXPR_DEPTH: + raise ValueError(f"Expression tree exceeds max depth of {_MAX_EXPR_DEPTH}") + if "channel" in node: + return db.query.channel(channel_name=node["channel"], **node.get("tags", {})) + if "const" in node: + return node["const"] + if "method" in node: + method = node["method"] + if method not in _EXPR_METHODS: + raise ValueError(f"Unsupported method {method!r}; must be one of {sorted(_EXPR_METHODS)}") + if "operand" not in node: + raise ValueError(f"Method node must have 'operand': {node}") + operand = build_expr(node["operand"], db, _depth + 1) + args = [_resolve_arg(a, db, _depth + 1) for a in node.get("args", [])] + kwargs = {k: _resolve_arg(v, db, _depth + 1) for k, v in node.get("kwargs", {}).items()} + return getattr(operand, method)(*args, **kwargs) + if "op" not in node: + raise ValueError(f"Expression node must have 'channel', 'const', 'method', or 'op': {node}") + op = node["op"] + if op not in _EXPR_OPS: + raise ValueError(f"Unsupported op {op!r}; must be one of {sorted(_EXPR_OPS)}") + left = build_expr(node["left"], db, _depth + 1) + right = build_expr(node["right"], db, _depth + 1) + return _EXPR_OPS[op](left, right) + + +_EXPR_DOC = """Expression trees (signal_expr / condition_expr / event.signal_expr): +{"channel": "", "tags": {...}} references a channel; {"const": } is a +constant; {"op": "add"|"sub"|"mul"|"div"|"gt"|"lt"|"ge"|"le"|"eq"|"ne"|"and"|"or", +"left": , "right": } combines two nodes (a "gt"/"lt"/etc. +comparison produces an Intervals, e.g. "channel > 2000" means "the time +ranges where this holds", not a per-sample boolean); {"method": +"resample"|"cumtrapz"|"diff"|"where"|"rolling_average", "operand": , +"args": [...]} transforms a numeric signal; {"method": +"start_points"|"end_points", "operand": } turns a comparison's +Intervals into discrete instants (e.g. each time a condition starts/stops +holding); {"method": "rising_edges"|"falling_edges", "operand": } finds +transitions in an already-boolean raw channel (e.g. a brake switch), as +opposed to a comparison. Args may themselves be nested nodes, e.g. the +condition passed to "where". Example -- average of two channels: {"op": "div", +"left": {"op": "add", "left": {"channel": "A"}, "right": {"channel": "B"}}, +"right": {"const": 2}}. Example -- cumulative distance from speed: {"method": +"cumtrapz", "operand": {"method": "resample", "operand": {"channel": "Vehicle Speed"}, +"args": [1000000]}}.""" + + +# --------------------------------------------------------------------------- +# Events +# --------------------------------------------------------------------------- + +def build_event(event_spec: dict | None, db): + """Build an Event from a constrained spec dict. + + Shapes: + None or {"type": "container"} -- ContainerEvent: the whole recording, + no condition. Default when event is omitted. + {"type": "basic", "condition_expr": } -- BasicEvent from a + compound boolean expression tree (see _EXPR_DOC). + {"type": "basic", "channel_name"/"signal_expr", "tags", "condition_op", + "condition_value"} -- BasicEvent from a single threshold. + {"type": "points_in_time", "signal_expr": threshold>}>} -- PointsInTimeEvent, + for discrete instants (pairs with preview_point_values; + duration-weighted tools like preview_histogram/preview_stats need + Intervals, not instants). + """ + from impulse_reporting.events.basic_event import BasicEvent + from impulse_reporting.events.container_event import ContainerEvent + from impulse_reporting.events.points_in_time_event import PointsInTimeEvent + + if event_spec is None: + return ContainerEvent(name="adhoc_event", desc="Full measurement") + + etype = event_spec.get("type", "basic") + + if etype == "container": + return ContainerEvent(name="adhoc_event", desc="Full measurement") + + if etype == "basic": + if "condition_expr" in event_spec: + cond = build_expr(event_spec["condition_expr"], db) + desc = "custom condition" + else: + if "signal_expr" in event_spec: + channel = build_expr(event_spec["signal_expr"], db) + elif "channel_name" in event_spec: + channel = db.query.channel( + channel_name=event_spec["channel_name"], **event_spec.get("tags", {}) + ) + else: + raise ValueError( + "basic event needs condition_expr, or (channel_name/signal_expr " + "+ condition_op + condition_value)" + ) + op = event_spec.get("condition_op") + value = event_spec.get("condition_value") + if op not in _OPS or value is None: + raise ValueError(f"condition_op must be one of {list(_OPS)} and condition_value must be set") + cond = _OPS[op](channel, value) + desc = f"{event_spec.get('channel_name', 'signal')} {op} {value}" + return BasicEvent(name="adhoc_event", expr=cond, desc=desc) + + if etype == "points_in_time": + if "signal_expr" not in event_spec: + raise ValueError( + "points_in_time event needs signal_expr (must evaluate to PointsInTime, " + "e.g. a rising_edges/falling_edges method node)" + ) + expr = build_expr(event_spec["signal_expr"], db) + return PointsInTimeEvent(name="adhoc_event", expr=expr, desc="points in time") + + raise ValueError(f"Unsupported event type {etype!r}; must be one of: container, basic, points_in_time") + + +_EVENT_DOC = """event scopes which time range an aggregation applies to. Omit +for the whole recording. {"type": "basic", "channel_name": "...", "tags": {...}, +"condition_op": ">", "condition_value": 2000} scopes to a threshold; {"type": +"basic", "condition_expr": } scopes to a compound boolean condition +(see the expression tree docs). {"type": "points_in_time", "signal_expr": } +scopes to discrete instants (e.g. {"method": "start_points", "operand": +{"op": "gt", "left": {"channel": "Engine RPM"}, "right": {"const": 2000}}} +for each moment RPM crosses above 2000) -- required for preview_point_values, +not usable with duration-weighted tools.""" + + +# --------------------------------------------------------------------------- +# Grounding tools +# --------------------------------------------------------------------------- + +@mcp.tool(name="list_channels") +def list_channels_tool() -> list[dict]: + """List available measurement channels and their tags (e.g. brand, model, + unit -- whatever tag vocabulary this dataset uses). Read-only and + instant, with no effect on the report being built -- distinct from + search_aliases/add_physical_signal, which look up or commit a channel to + the persisted report definition. Call this to ground yourself before + preview_histogram/preview_histogram_2d/preview_stats/preview_point_values + (so channel_name/tags are real values, not guesses), or whenever the user + asks an exploratory question ("what signals do we have?") without yet + committing to add anything.""" + spark = get_spark() + tag_keys = [r["key"] for r in spark.table(f"{PFX}_channel_tags").select("key").distinct().collect()] + order_col = "channel_name" if "channel_name" in tag_keys else tag_keys[0] + df = ( + spark.table(f"{PFX}_channel_tags") + .groupBy("container_id", "channel_id") + .pivot("key", tag_keys) + .agg(F.first("value")) + .select(*tag_keys) + .distinct() + .orderBy(order_col) + ) + return [row.asDict() for row in df.collect()] + + +@mcp.tool(name="list_containers") +def list_containers_tool() -> list[dict]: + """List available measurement containers (recording sessions) and their + tags (e.g. vehicle, duration, condition). Read-only and instant -- + distinct from set_vehicle, which commits a container/time-range to the + persisted report. Use for questions like "how many test drives do we + have?" or "what's the total recorded duration?" without adding anything + to the report.""" + spark = get_spark() + tag_keys = [r["key"] for r in spark.table(f"{PFX}_container_tags").select("key").distinct().collect()] + tags_df = ( + spark.table(f"{PFX}_container_tags") + .groupBy("container_id") + .pivot("key", tag_keys) + .agg(F.first("value")) + ) + df = ( + spark.table(f"{PFX}_container_metrics") + .join(tags_df, on="container_id", how="left") + .orderBy("container_id") + ) + return [row.asDict() for row in df.collect()] + + +# --------------------------------------------------------------------------- +# Histogram weighting +# --------------------------------------------------------------------------- + +def _resolve_weight(weight: dict | None, db): + """Returns (kind, weights_expr_or_None, extra_kwargs) for a weight spec. + Only "duration" (the default) is currently supported -- see the + NotImplementedError message below for why distance/custom are gated off. + """ + if weight is None: + return "duration", None, {} + wtype = weight.get("type", "duration") + if wtype == "duration": + return "duration", None, {} + if wtype in ("distance", "custom"): + raise NotImplementedError( + f"weight type {wtype!r} is not supported yet: verification found " + "HistogramDistance/HistogramCustomWeights produce numerically incorrect " + "results (off by several orders of magnitude) when the weight signal is " + "derived via resample()+cumtrapz() -- traced to the synchronized()+diff() " + "interaction inside the query engine, not something fixable from this MCP " + "server. Duration weighting (the default -- omit weight, or {\"type\": " + "\"duration\"}) is fully verified and safe to use." + ) + raise ValueError(f"weight type must be 'duration' (distance/custom not yet supported), got {wtype!r}") + + +_WEIGHT_DOC = """weight controls what the histogram accumulates per bin. +Currently only plain duration weighting is supported (time spent in each +bin, in seconds) -- omit weight, or pass {"type": "duration"}. Distance/ +custom weighting exists in Impulse but is not exposed here yet pending a +numerical discrepancy fix in the underlying query engine.""" + + +# --------------------------------------------------------------------------- +# Preview tools +# --------------------------------------------------------------------------- + +@mcp.tool() +def preview_histogram( + bins: list[float], + channel_name: str | None = None, + tags: dict[str, str] | None = None, + signal_expr: dict | None = None, + event: dict | None = None, + weight: dict | None = None, + bins_unit: str | None = None, + values_unit: str = "s", +) -> list[dict]: + """Compute a 1D histogram RIGHT NOW and return the actual numeric result + -- an instant, read-only preview, not a persisted report aggregation. + Use this when the user asks a question like "what's the distribution of + X?" or wants to sanity-check bin edges before committing to anything. Do + NOT use this in place of add_histogram: add_histogram only registers a + definition in the wizard's Aggregations step for the report that later + gets deployed as a job -- this tool runs the computation immediately and + returns the answer inline, nothing is persisted. A natural flow: call + this a few times to try different bin choices interactively, then call + add_histogram with the finalized parameters once the user is happy with + what they saw here. + + The histogrammed value is either a plain channel (channel_name + optional + tags to disambiguate, from list_channels) or a derived "virtual signal" + via signal_expr. """ + _EXPR_DOC + """ + + """ + _EVENT_DOC + """ + + """ + _WEIGHT_DOC + """ + + Provide either (channel_name [+ tags]) or signal_expr. bins_unit/ + values_unit are display labels only (e.g. "rpm") -- the returned value is + always in seconds regardless of values_unit. Returns one row per bin with + bin_name, lower_bound, and duration_s.""" + from impulse_reporting.core.page import Page + from impulse_reporting.aggregations.histogram import HistogramDuration + + report, db = _new_report("adhoc") + + if signal_expr is not None: + channel = build_expr(signal_expr, db) + elif channel_name is not None: + channel = db.query.channel(channel_name=channel_name, **(tags or {})) + else: + raise ValueError("Either channel_name or signal_expr must be provided") + + ev = build_event(event, db) + report.add_event(ev) + + bins_f = [float(b) for b in bins] + display_name = channel_name or "virtual_signal" + _resolve_weight(weight, db) # duration is the only supported kind; raises otherwise + + agg = HistogramDuration( + name="adhoc_histogram", base_expr=channel, bins=bins_f, event=ev, + channel_name=display_name, bins_unit=bins_unit or "", values_unit=values_unit, + ) + + page = Page(page_number=1) + page.add_aggregation(agg) + report.add_page(page) + report.determine_report() + + hist_df = _aggregation_df(report, "HISTOGRAM") + result_df = ( + hist_df + .groupBy("bin_name", "lower_bound") + .agg(F.sum("hist_value").alias("duration_us")) + .orderBy("lower_bound") + .toPandas() + ) + result_df["duration_s"] = result_df["duration_us"] / 1e6 + + return result_df[["bin_name", "lower_bound", "duration_s"]].to_dict(orient="records") + + +@mcp.tool() +def preview_histogram_2d( + x_bins: list[float], + y_bins: list[float], + x_channel_name: str | None = None, + y_channel_name: str | None = None, + tags: dict[str, str] | None = None, + x_signal_expr: dict | None = None, + y_signal_expr: dict | None = None, + event: dict | None = None, + weight: dict | None = None, + x_bins_unit: str | None = None, + y_bins_unit: str | None = None, + values_unit: str = "s", +) -> list[dict]: + """Compute a 2D heatmap of two signals (x vs y) RIGHT NOW and return the + actual numeric result -- an instant, read-only preview, not a persisted + report aggregation. Use when the user wants to see how two signals + correlate (e.g. "how does RPM relate to speed while RPM is above 2000?"). + Do NOT use this in place of add_histogram_2d: that only registers a + definition in the wizard's Aggregations step for the report that later + gets deployed as a job -- this tool runs the computation immediately and + returns the answer inline, nothing is persisted. + + x/y axes: provide x_channel_name/y_channel_name (+ optional tags), or + x_signal_expr/y_signal_expr for derived signals. """ + _EXPR_DOC + """ + + """ + _EVENT_DOC + """ + + """ + _WEIGHT_DOC + """ + + x_bins_unit/y_bins_unit/values_unit are display labels only -- the + returned value is always in seconds regardless of values_unit. Returns + one row per (x_bin, y_bin) with x_bin_name, y_bin_name, x_lower_bound, + y_lower_bound, and duration_s.""" + from impulse_reporting.core.page import Page + from impulse_reporting.aggregations.histogram2d import Histogram2DDuration + + report, db = _new_report("adhoc2d") + + if x_signal_expr is not None: + x_channel = build_expr(x_signal_expr, db) + elif x_channel_name is not None: + x_channel = db.query.channel(channel_name=x_channel_name, **(tags or {})) + else: + raise ValueError("Either x_channel_name or x_signal_expr must be provided") + + if y_signal_expr is not None: + y_channel = build_expr(y_signal_expr, db) + elif y_channel_name is not None: + y_channel = db.query.channel(channel_name=y_channel_name, **(tags or {})) + else: + raise ValueError("Either y_channel_name or y_signal_expr must be provided") + + ev = build_event(event, db) + report.add_event(ev) + + x_bins_f = [float(b) for b in x_bins] + y_bins_f = [float(b) for b in y_bins] + x_name = x_channel_name or "x_signal" + y_name = y_channel_name or "y_signal" + _resolve_weight(weight, db) # duration is the only supported kind; raises otherwise + + agg = Histogram2DDuration( + name="adhoc_histogram2d", x_expr=x_channel, y_expr=y_channel, + x_bins=x_bins_f, y_bins=y_bins_f, event=ev, + x_channel_name=x_name, y_channel_name=y_name, + x_bins_unit=x_bins_unit, y_bins_unit=y_bins_unit, values_unit=values_unit, + ) + + page = Page(page_number=1) + page.add_aggregation(agg) + report.add_page(page) + report.determine_report() + + hist_df = _aggregation_df(report, "HISTOGRAM2D") + result_df = ( + hist_df + .groupBy("x_bin_name", "y_bin_name", "x_lower_bound", "y_lower_bound") + .agg(F.sum("hist_value").alias("duration_us")) + .orderBy("x_lower_bound", "y_lower_bound") + .toPandas() + ) + result_df["duration_s"] = result_df["duration_us"] / 1e6 + + return result_df[ + ["x_bin_name", "y_bin_name", "x_lower_bound", "y_lower_bound", "duration_s"] + ].to_dict(orient="records") + + +_SIGNALS_DOC = """signals: list of {"label": "...", "channel_name": "...", +"tags": {...}} (a plain channel) or {"label": "...", "signal_expr": } +(a derived virtual signal) -- label names the signal in the returned rows.""" + +_STATS = {"min", "max", "mean", "median"} + + +@mcp.tool() +def preview_stats( + signals: list[dict], + statistics: list[Literal["min", "max", "mean", "median"]] | None = None, + event: dict | None = None, +) -> list[dict]: + """Compute summary statistics for one or more signals RIGHT NOW and + return the actual numeric result -- an instant, read-only preview, not a + persisted report aggregation. Use for questions like "what's the average + X?" or "what were the min/max of X and Y during this event?" -- this is + the tool for statistics across MULTIPLE signals at once. Do NOT use this + in place of add_statistics: that only registers a definition in the + wizard's Aggregations step for the report that later gets deployed as a + job. + + """ + _SIGNALS_DOC + """ + statistics: which to compute (default: all four -- min, max, mean, median). + + """ + _EVENT_DOC + """ + + Statistics are computed per container (a test drive/recording session), + not collapsed across containers -- averaging an already-averaged value + across sessions without knowing sample counts wouldn't be statistically + valid, so each container gets its own row. Returns one row per + (container_id, label, statistic) with columns container_id, label, + statistic, value.""" + from impulse_reporting.core.page import Page + from impulse_reporting.aggregations.stats_aggregator import StatsAggregator + + if not signals: + raise ValueError("signals must have at least one entry") + stats = statistics or sorted(_STATS) + bad = set(stats) - _STATS + if bad: + raise ValueError(f"Unsupported statistics {bad}; must be a subset of {sorted(_STATS)}") + + report, db = _new_report("adhoc_stats") + + ev = build_event(event, db) + report.add_event(ev) + + exprs, labels = [], [] + for sig in signals: + if "signal_expr" in sig: + exprs.append(build_expr(sig["signal_expr"], db)) + elif "channel_name" in sig: + exprs.append(db.query.channel(channel_name=sig["channel_name"], **sig.get("tags", {}))) + else: + raise ValueError(f"signal {sig!r} needs channel_name or signal_expr") + labels.append(sig["label"]) + + page = Page(page_number=1) + page.add_aggregation(StatsAggregator( + name="adhoc_stats", input_expressions=exprs, channel_names=labels, + statistics=list(stats), event=ev, + )) + report.add_page(page) + report.determine_report() + + stats_df = _aggregation_df(report, "STATS_AGGREGATOR") + result_df = ( + stats_df + .select("container_id", "channel_name", "aggregation_label", "statistic_value") + .orderBy("container_id", "channel_name", "aggregation_label") + .toPandas() + .rename(columns={"channel_name": "label", "aggregation_label": "statistic", "statistic_value": "value"}) + ) + return result_df.to_dict(orient="records") + + +@mcp.tool() +def preview_point_values( + signals: list[dict], + event: dict, +) -> list[dict]: + """Sample one or more signals at each instant of a points-in-time event + RIGHT NOW and return the actual values -- an instant, read-only preview, + not a persisted report aggregation. Use for questions like "what was the + speed and RPM at each 10 km milestone?" or "what was the temperature at + every rising edge of the brake signal?" -- this is the one case + preview_histogram/preview_stats can't cover, since those need time + intervals (durations), not discrete instants. + + """ + _SIGNALS_DOC + """ + + event: REQUIRED, and must be a points_in_time event -- {"type": + "points_in_time", "signal_expr": }. + + Returns one row per (container_id, event_instance_id, label) with the + sampled value.""" + from impulse_reporting.core.page import Page + from impulse_reporting.aggregations.point_value_aggregator import PointValueAggregator + + if not signals: + raise ValueError("signals must have at least one entry") + if not event or event.get("type") != "points_in_time": + raise ValueError('event must be a points_in_time event: {"type": "points_in_time", "signal_expr": ...}') + + report, db = _new_report("adhoc_pv") + + ev = build_event(event, db) + report.add_event(ev) + + exprs, labels = [], [] + for sig in signals: + if "signal_expr" in sig: + exprs.append(build_expr(sig["signal_expr"], db)) + elif "channel_name" in sig: + exprs.append(db.query.channel(channel_name=sig["channel_name"], **sig.get("tags", {}))) + else: + raise ValueError(f"signal {sig!r} needs channel_name or signal_expr") + labels.append(sig["label"]) + + page = Page(page_number=1) + page.add_aggregation(PointValueAggregator( + name="adhoc_point_values", input_expressions=exprs, channel_names=labels, event=ev, + )) + report.add_page(page) + report.determine_report() + + pv_df = _aggregation_df(report, "POINT_VALUE_AGGREGATOR") + result_df = ( + pv_df + .select("container_id", "event_instance_id", "channel_name", "statistic_value") + .orderBy("container_id", "event_instance_id", "channel_name") + .toPandas() + .rename(columns={"channel_name": "label", "statistic_value": "value"}) + ) + return result_df.to_dict(orient="records") + + +def _keep_warm(): + """Ping the serverless session periodically so the underlying compute + doesn't scale down from inactivity between questions.""" + import time + while True: + try: + get_spark().sql("SELECT 1").collect() + except Exception: + pass + time.sleep(120) + + +def main(): + import threading + threading.Thread(target=_keep_warm, daemon=True).start() + mcp.run(transport="streamable-http") + + +if __name__ == "__main__": + main() diff --git a/demos/agent_mcp_query.ipynb b/demos/agent_mcp_query.ipynb new file mode 100644 index 00000000..56ed855e --- /dev/null +++ b/demos/agent_mcp_query.ipynb @@ -0,0 +1,1159 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Impulse \u2014 Ad-hoc Agent Queries via MCP\n", + "\n", + "This notebook demonstrates an **ad-hoc** query pattern: answering questions\n", + "(\"what fraction of time was the engine above 3000 RPM?\") where the channels,\n", + "thresholds, and bins aren't known ahead of time, no report has been defined\n", + "in code for them, and no Gold table exists yet for the answer. Instead, the\n", + "report is built and computed on the fly, in response to a caller's\n", + "parameters, and the result is returned directly without persisting a\n", + "Gold-layer star schema.\n", + "\n", + "**Why not just point Genie/AI-BI at the Gold layer?** Genie does text-to-SQL\n", + "against tables that already exist with a semantic model configured up front.\n", + "Impulse's value is defining a *new* TSAL event/aggregation per question \u2014\n", + "there's no way to pre-populate Gold tables for arbitrary thresholds, channel\n", + "combinations, and bin sizes a user might ask about. The report has to be\n", + "**computed on demand**, which means something has to hold a live Spark\n", + "session and run Impulse's `Report` API synchronously in response to a request.\n", + "\n", + "**What this notebook builds:**\n", + "1. **Grounding tools** \u2014 read channel/container metadata so a caller knows what\n", + " exists before building a query (same role schema-introspection plays in\n", + " text-to-SQL).\n", + "2. **A constrained ad-hoc execution function** \u2014 takes a small, validated set\n", + " of parameters (channel, condition, bins) and drives the real `Report` /\n", + " `BasicEvent` / `HistogramDuration` API against a **throwaway scratch\n", + " prefix**, so ad-hoc runs never collide with or pollute each other or the\n", + " real Gold layer.\n", + "3. **The same two functions exposed as MCP tools**, with an in-notebook MCP\n", + " client call proving the round trip works end to end against real data.\n", + "\n", + "Deliberately *not* done here: letting an LLM generate and `exec()` raw TSAL/\n", + "Python. TSAL expressions are live Python objects wired into a Spark execution\n", + "plan \u2014 free-form generated code run against a cluster with Unity Catalog\n", + "access is a real code-injection risk. Instead the tool surface below only\n", + "accepts a small typed parameter set that this notebook's own code translates\n", + "into Impulse API calls." + ], + "id": "300f1354752c4037" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1. Setup\n", + "\n", + "Pins `databricks-sdk==0.106.0` to match Impulse's `pyproject.toml` \u2014 its\n", + "telemetry code reads a private `Config._product_info` attribute that only\n", + "exists on that version. Job/serverless compute can default to an older SDK\n", + "(observed: 0.20.0) that lacks it, so pin explicitly rather than relying on\n", + "whatever happens to be preinstalled." + ], + "id": "c0479f1bbae94789" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install mcp \"databricks-sdk==0.106.0\" scipy -q\n", + "dbutils.library.restartPython()" + ], + "id": "7652cc5a45534163" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Configure target location\n", + "\n", + "Fill in **Catalog**, **Schema**, and **Table Prefix** in the widgets above,\n", + "then run the next cells. This notebook loads its own copy of the demo\n", + "silver layer, so it's self-contained and doesn't depend on any other\n", + "notebook having been run first." + ], + "id": "c4c1389023334ba8" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dbutils.widgets.text(\"catalog\", \"\", \"Catalog\")\n", + "dbutils.widgets.text(\"schema\", \"\", \"Schema\")\n", + "dbutils.widgets.text(\"table_prefix\", \"agent\", \"Table Prefix\")" + ], + "id": "5cd1b27bedc04e75" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys, os\n", + "import pandas as pd\n", + "\n", + "CATALOG = dbutils.widgets.get(\"catalog\")\n", + "SCHEMA = dbutils.widgets.get(\"schema\")\n", + "TABLE_PREFIX = dbutils.widgets.get(\"table_prefix\") or \"agent\"\n", + "\n", + "if not CATALOG or not SCHEMA or not TABLE_PREFIX:\n", + " raise ValueError(\"Please set Catalog, Schema, and Table Prefix widgets above before running.\")\n", + "\n", + "nb_path = dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()\n", + "DEMOS_DIR = \"/Workspace\" + \"/\".join(nb_path.split(\"/\")[:-1])\n", + "REPO_ROOT = \"/Workspace\" + \"/\".join(nb_path.split(\"/\")[:-2])\n", + "sys.path.insert(0, os.path.join(REPO_ROOT, \"src\"))\n", + "\n", + "pfx = f\"{CATALOG}.{SCHEMA}.{TABLE_PREFIX}\"\n", + "print(f\"Silver layer target: {pfx}_*\")" + ], + "id": "b3d14e93d9854981" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Load demo data into Silver layer\n", + "\n", + "Same bootstrap step as the other two demos: 5 silver tables from CSV." + ], + "id": "e5682a927dc04243" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "spark.sql(f\"CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}\")\n", + "\n", + "csv_dir = os.path.join(DEMOS_DIR, \"data\", \"reporting\")\n", + "SILVER = [\"container_metrics\", \"container_tags\", \"channel_metrics\", \"channel_tags\", \"channels\"]\n", + "for t in SILVER:\n", + " (spark.createDataFrame(pd.read_csv(f\"{csv_dir}/{t}.csv\"))\n", + " .write.mode(\"overwrite\")\n", + " .saveAsTable(f\"{pfx}_{t}\"))\n", + "print(f\"Loaded {len(SILVER)} silver-layer tables under {pfx}_*\")" + ], + "id": "995765e278df42e4" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Grounding tools\n", + "\n", + "Before an agent can build a query, it needs to know what channels and\n", + "containers actually exist \u2014 real channel names and whatever tag vocabulary\n", + "this particular dataset happens to use (brand/model here, but Impulse\n", + "doesn't fix the tag schema \u2014 a different silver layer could tag channels\n", + "by `sensor_type`, `line_id`, `plant`, anything). These read directly from\n", + "the silver metadata tables; no Impulse `Report` needed, so they're cheap.\n", + "\n", + "Channel/container tags are stored as EAV (`key`, `value`) pairs per\n", + "`(container_id, channel_id)` \u2014 pivoted here into one row per channel/\n", + "container. The set of tag keys is **discovered at query time**, not\n", + "hardcoded, so this works against any Impulse dataset's tag vocabulary,\n", + "not just this demo's." + ], + "id": "914a90ebf21f4d9a" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pyspark.sql.functions as F\n", + "\n", + "def list_channels() -> list[dict]:\n", + " \"\"\"List distinct measurement channels available in the silver layer, with all their tags.\"\"\"\n", + " tag_keys = [r[\"key\"] for r in spark.table(f\"{pfx}_channel_tags\").select(\"key\").distinct().collect()]\n", + " order_col = \"channel_name\" if \"channel_name\" in tag_keys else tag_keys[0]\n", + " df = (\n", + " spark.table(f\"{pfx}_channel_tags\")\n", + " .groupBy(\"container_id\", \"channel_id\")\n", + " .pivot(\"key\", tag_keys)\n", + " .agg(F.first(\"value\"))\n", + " .select(*tag_keys)\n", + " .distinct()\n", + " .orderBy(order_col)\n", + " )\n", + " return [row.asDict() for row in df.collect()]\n", + "\n", + "\n", + "def list_containers() -> list[dict]:\n", + " \"\"\"List available measurement containers (recording sessions), with all their tags.\"\"\"\n", + " tag_keys = [r[\"key\"] for r in spark.table(f\"{pfx}_container_tags\").select(\"key\").distinct().collect()]\n", + " tags_df = (\n", + " spark.table(f\"{pfx}_container_tags\")\n", + " .groupBy(\"container_id\")\n", + " .pivot(\"key\", tag_keys)\n", + " .agg(F.first(\"value\"))\n", + " )\n", + " df = (\n", + " spark.table(f\"{pfx}_container_metrics\")\n", + " .join(tags_df, on=\"container_id\", how=\"left\")\n", + " .orderBy(\"container_id\")\n", + " )\n", + " return [row.asDict() for row in df.collect()]" + ], + "id": "f1ffbd91eecd4034" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Quick sanity check before wiring these into anything else.\n", + "display(pd.DataFrame(list_channels()))\n", + "display(pd.DataFrame(list_containers()))" + ], + "id": "43d303bbbb764473" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "%md\n", + "# 3. Ad-hoc execution\n", + "\n", + "This is the part that actually drives Impulse. Each call builds a fresh\n", + "`Report`, defines an event and an aggregation from the caller's\n", + "parameters, runs `determine_report()`, then reads the answer straight out\n", + "of `report.aggregation_dfs` in memory. No `persist_results()`, no\n", + "`unity_sink` in the config, no Delta round-trip for a throwaway one-off\n", + "answer \u2014 the fact schemas already have the final columns directly, so\n", + "nothing needs to be written and read back just to get a result that's\n", + "already sitting in memory. This is what took per-call latency from 33-43s\n", + "down to 2-7s in the deployed version.\n", + "\n", + "**Comprehensive coverage.** Impulse's real capability surface is bigger\n", + "than \"one channel, one threshold, one histogram\" \u2014 this section covers\n", + "everything the framework supports today:\n", + "- **Virtual signals**: not just arithmetic on channels, but the full\n", + " `SampleSeries` transform surface`resample`, `cumtrapz`, `diff`,\n", + " `where`, edge/point detection.\n", + "- **Events**: all four kinds \u2014 `BasicEvent` (threshold \u2192 intervals),\n", + " `ContainerEvent` (whole recording, no condition), `PointsInTimeEvent`\n", + " (discrete instants, e.g. each time a condition starts holding).\n", + "- **Aggregations**: histograms and 2D heatmaps, plus\n", + " `StatsAggregator` (min/max/mean/median across *multiple* signals at\n", + " once) and `PointValueAggregator` (sample values at instants \u2014 the one\n", + " case duration-weighted tools can't cover).\n", + "\n", + "**Virtual signal / condition expression trees** (`signal_expr`/\n", + "`condition_expr`/an event's `signal_expr`): `{\"channel\": \"\", \"tags\":\n", + "{...}}` references a channel; `{\"const\": }` is a constant; `{\"op\":\n", + "\"add\"|\"sub\"|\"mul\"|\"div\"|\"gt\"|\"lt\"|\"ge\"|\"le\"|\"eq\"|\"ne\"|\"and\"|\"or\", \"left\":\n", + ", \"right\": }` combines two nodes (a comparison like \"gt\"\n", + "produces an *Intervals* \u2014 \"channel > 2000\" means \"the time ranges where\n", + "this holds\", not a per-sample boolean); `{\"method\": \"resample\"|\"cumtrapz\"|\n", + "\"diff\"|\"where\"|\"rolling_average\", \"operand\": , \"args\": [...]}`\n", + "transforms a numeric signal; `{\"method\": \"start_points\"|\"end_points\",\n", + "\"operand\": }` turns a comparison's Intervals into discrete instants\n", + "(e.g. each time a condition starts/stops holding); `{\"method\":\n", + "\"rising_edges\"|\"falling_edges\", \"operand\": }` finds transitions in\n", + "an already-boolean raw channel (e.g. a brake switch), as opposed to a\n", + "comparison. Every op/method maps to a fixed, known Python operator or TSAL\n", + "method \u2014 there is no `eval()`/`exec()` of generated code anywhere, even\n", + "though this now covers arithmetic, resampling, integration, and edge\n", + "detection. Example \u2014 cumulative distance from speed: `{\"method\":\n", + "\"cumtrapz\", \"operand\": {\"method\": \"resample\", \"operand\": {\"channel\":\n", + "\"Vehicle Speed\"}, \"args\": [1000000]}}`.\n", + "\n", + "**Events**: omit for the whole recording (`ContainerEvent`). `{\"type\":\n", + "\"basic\", \"channel_name\": \"...\", \"tags\": {...}, \"condition_op\": \">\",\n", + "\"condition_value\": 2000}` scopes to a threshold; `{\"type\": \"basic\",\n", + "\"condition_expr\": }` scopes to a compound boolean condition.\n", + "`{\"type\": \"points_in_time\", \"signal_expr\": }` scopes to discrete\n", + "instants, e.g. `{\"method\": \"start_points\", \"operand\": {\"op\": \"gt\", \"left\":\n", + "{\"channel\": \"Engine RPM\"}, \"right\": {\"const\": 2000}}}` for each moment RPM\n", + "crosses above 2000 \u2014 required for point-value sampling, not usable with\n", + "duration-weighted aggregations.\n", + "\n", + "**Known limitation \u2014 distance/custom-weighted histograms.** Impulse also\n", + "supports weighting histograms by distance or an arbitrary custom signal\n", + "(`HistogramDistance`/`HistogramCustomWeights`), not just duration. Testing\n", + "found these produce numerically incorrect results (off by several orders\n", + "of magnitude) when the weight signal is derived via `resample()` +\n", + "`cumtrapz()` \u2014 verified by checking the same signal's magnitude directly\n", + "via `StatsAggregator` (correct, tens of km) versus as a histogram weight\n", + "(near-zero), reproducible with no event scoping at all. The discrepancy\n", + "traces to the `synchronized()` + `diff()` interaction inside\n", + "`HistogramCustomWeights.build()` in the query engine \u2014 not something\n", + "fixable from this notebook/MCP layer. Duration weighting (the only kind\n", + "exposed below) is fully verified and safe to use." + ], + "id": "44ca9ac9f94c47c6" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import operator\n", + "import uuid\n", + "\n", + "from databricks.sdk import WorkspaceClient\n", + "from impulse_reporting.core.report import Report\n", + "from impulse_reporting.core.page import Page\n", + "from impulse_reporting.aggregations.histogram import HistogramDuration\n", + "from impulse_reporting.aggregations.histogram2d import Histogram2DDuration\n", + "from impulse_reporting.aggregations.stats_aggregator import StatsAggregator\n", + "from impulse_reporting.aggregations.point_value_aggregator import PointValueAggregator\n", + "from impulse_reporting.events.basic_event import BasicEvent\n", + "from impulse_reporting.events.container_event import ContainerEvent\n", + "from impulse_reporting.events.points_in_time_event import PointsInTimeEvent\n", + "\n", + "def _adhoc_config():\n", + " return {\n", + " \"source\": {\n", + " \"container_metrics_table\": f\"{pfx}_container_metrics\",\n", + " \"channel_metrics_table\": f\"{pfx}_channel_metrics\",\n", + " \"channels_uri\": f\"{pfx}_channels\",\n", + " \"container_tags_table\": f\"{pfx}_container_tags\",\n", + " \"channel_tags_table\": f\"{pfx}_channel_tags\",\n", + " },\n", + " # No unity_sink: read the result from report.aggregation_dfs in\n", + " # memory instead of persisting -- see markdown above.\n", + " \"query_engine\": {\"solver\": \"DefaultSolver\", \"data_type\": \"RAW\"},\n", + " \"measurement_dimensions\": [\"container_id\", \"vehicle_key\", \"start_ts\", \"stop_ts\"],\n", + " }\n", + "\n", + "\n", + "_OPS = {\n", + " \">\": operator.gt,\n", + " \"<\": operator.lt,\n", + " \">=\": operator.ge,\n", + " \"<=\": operator.le,\n", + " \"==\": operator.eq,\n", + "}\n", + "\n", + "# Whitelist-only binary operators for expression trees. Every node maps to a\n", + "# fixed, known Python operator applied to TSAL objects -- there is no\n", + "# eval()/exec() of user-provided strings anywhere.\n", + "_EXPR_OPS = {\n", + " \"add\": operator.add, \"sub\": operator.sub, \"mul\": operator.mul, \"div\": operator.truediv,\n", + " \"gt\": operator.gt, \"lt\": operator.lt, \"ge\": operator.ge, \"le\": operator.le,\n", + " \"eq\": operator.eq, \"ne\": operator.ne, \"and\": operator.and_, \"or\": operator.or_,\n", + "}\n", + "\n", + "# Whitelist-only unary TSAL methods reachable from an expression tree's\n", + "# \"method\" node. Same safety property as _EXPR_OPS: only names literally in\n", + "# this set are ever passed to getattr(), so there's no way to reach an\n", + "# unintended method (e.g. a dunder) even in principle.\n", + "_EXPR_METHODS = {\n", + " # SampleSeries (raw/derived numeric signals)\n", + " \"resample\", \"cumtrapz\", \"diff\", \"where\",\n", + " \"rising_edges\", \"falling_edges\", \"rising_edge\", \"falling_edge\",\n", + " \"intervals_between_falling_edges\", \"rolling_average\",\n", + " # Intervals (comparison-derived conditions, e.g. channel > threshold)\n", + " \"start_points\", \"end_points\",\n", + "}\n", + "\n", + "_MAX_EXPR_DEPTH = 10\n", + "\n", + "\n", + "def _resolve_arg(arg, db, depth):\n", + " \"\"\"A method node's args/kwargs may be a literal or a nested expression node.\"\"\"\n", + " if isinstance(arg, dict) and ({\"channel\", \"const\", \"op\", \"method\"} & arg.keys()):\n", + " return build_expr(arg, db, depth)\n", + " return arg\n", + "\n", + "\n", + "def build_expr(node: dict, db, _depth: int = 0):\n", + " \"\"\"Recursively build a TSAL expression from a constrained JSON tree --\n", + " see the markdown above for node shapes and the safety rationale.\"\"\"\n", + " if _depth > _MAX_EXPR_DEPTH:\n", + " raise ValueError(f\"Expression tree exceeds max depth of {_MAX_EXPR_DEPTH}\")\n", + " if \"channel\" in node:\n", + " return db.query.channel(channel_name=node[\"channel\"], **node.get(\"tags\", {}))\n", + " if \"const\" in node:\n", + " return node[\"const\"]\n", + " if \"method\" in node:\n", + " method = node[\"method\"]\n", + " if method not in _EXPR_METHODS:\n", + " raise ValueError(f\"Unsupported method {method!r}; must be one of {sorted(_EXPR_METHODS)}\")\n", + " if \"operand\" not in node:\n", + " raise ValueError(f\"Method node must have 'operand': {node}\")\n", + " operand = build_expr(node[\"operand\"], db, _depth + 1)\n", + " args = [_resolve_arg(a, db, _depth + 1) for a in node.get(\"args\", [])]\n", + " kwargs = {k: _resolve_arg(v, db, _depth + 1) for k, v in node.get(\"kwargs\", {}).items()}\n", + " return getattr(operand, method)(*args, **kwargs)\n", + " if \"op\" not in node:\n", + " raise ValueError(f\"Expression node must have 'channel', 'const', 'method', or 'op': {node}\")\n", + " op = node[\"op\"]\n", + " if op not in _EXPR_OPS:\n", + " raise ValueError(f\"Unsupported op {op!r}; must be one of {sorted(_EXPR_OPS)}\")\n", + " left = build_expr(node[\"left\"], db, _depth + 1)\n", + " right = build_expr(node[\"right\"], db, _depth + 1)\n", + " return _EXPR_OPS[op](left, right)\n", + "\n", + "\n", + "def build_event(event_spec: dict | None, db):\n", + " \"\"\"Build an Event from a constrained spec dict -- see the markdown above\n", + " for the container/basic/points_in_time shapes.\"\"\"\n", + " if event_spec is None:\n", + " return ContainerEvent(name=\"adhoc_event\", desc=\"Full measurement\")\n", + "\n", + " etype = event_spec.get(\"type\", \"basic\")\n", + "\n", + " if etype == \"container\":\n", + " return ContainerEvent(name=\"adhoc_event\", desc=\"Full measurement\")\n", + "\n", + " if etype == \"basic\":\n", + " if \"condition_expr\" in event_spec:\n", + " cond = build_expr(event_spec[\"condition_expr\"], db)\n", + " desc = \"custom condition\"\n", + " else:\n", + " if \"signal_expr\" in event_spec:\n", + " channel = build_expr(event_spec[\"signal_expr\"], db)\n", + " elif \"channel_name\" in event_spec:\n", + " channel = db.query.channel(\n", + " channel_name=event_spec[\"channel_name\"], **event_spec.get(\"tags\", {})\n", + " )\n", + " else:\n", + " raise ValueError(\n", + " \"basic event needs condition_expr, or (channel_name/signal_expr \"\n", + " \"+ condition_op + condition_value)\"\n", + " )\n", + " op = event_spec.get(\"condition_op\")\n", + " value = event_spec.get(\"condition_value\")\n", + " if op not in _OPS or value is None:\n", + " raise ValueError(f\"condition_op must be one of {list(_OPS)} and condition_value must be set\")\n", + " cond = _OPS[op](channel, value)\n", + " desc = f\"{event_spec.get('channel_name', 'signal')} {op} {value}\"\n", + " return BasicEvent(name=\"adhoc_event\", expr=cond, desc=desc)\n", + "\n", + " if etype == \"points_in_time\":\n", + " if \"signal_expr\" not in event_spec:\n", + " raise ValueError(\n", + " \"points_in_time event needs signal_expr (must evaluate to PointsInTime, \"\n", + " \"e.g. a start_points/end_points method node)\"\n", + " )\n", + " expr = build_expr(event_spec[\"signal_expr\"], db)\n", + " return PointsInTimeEvent(name=\"adhoc_event\", expr=expr, desc=\"points in time\")\n", + "\n", + " raise ValueError(f\"Unsupported event type {etype!r}; must be one of: container, basic, points_in_time\")\n", + "\n", + "\n", + "def _resolve_weight(weight: dict | None) -> None:\n", + " \"\"\"Only \"duration\" (the default) is supported -- see the \"Known\n", + " limitation\" note above for why distance/custom are gated off.\"\"\"\n", + " if weight is None:\n", + " return\n", + " wtype = weight.get(\"type\", \"duration\")\n", + " if wtype == \"duration\":\n", + " return\n", + " if wtype in (\"distance\", \"custom\"):\n", + " raise NotImplementedError(\n", + " f\"weight type {wtype!r} is not supported yet: verification found \"\n", + " \"HistogramDistance/HistogramCustomWeights produce numerically incorrect \"\n", + " \"results when the weight signal is derived via resample()+cumtrapz() -- \"\n", + " \"see the markdown above. Duration weighting (the default) is fully verified.\"\n", + " )\n", + " raise ValueError(f\"weight type must be 'duration' (distance/custom not yet supported), got {wtype!r}\")\n", + "\n", + "\n", + "def run_adhoc_histogram(\n", + " bins: list[float],\n", + " channel_name: str | None = None,\n", + " tags: dict[str, str] | None = None,\n", + " signal_expr: dict | None = None,\n", + " event: dict | None = None,\n", + " weight: dict | None = None,\n", + " bins_unit: str | None = None,\n", + " values_unit: str = \"s\",\n", + ") -> list[dict]:\n", + " \"\"\"\n", + " Compute a duration-weighted histogram, scoped to an event. Provide either\n", + " (channel_name [+ tags]) or signal_expr for the histogrammed value. See\n", + " the markdown above for the event/weight/expression-tree formats.\n", + " Returns one row per bin: bin_name, lower_bound, duration_s.\n", + " \"\"\"\n", + " report = Report(name=f\"adhoc_{uuid.uuid4().hex[:8]}\", spark=spark, config=_adhoc_config(), workspace_client=WorkspaceClient())\n", + " db = report.get_db()\n", + "\n", + " if signal_expr is not None:\n", + " channel = build_expr(signal_expr, db)\n", + " elif channel_name is not None:\n", + " channel = db.query.channel(channel_name=channel_name, **(tags or {}))\n", + " else:\n", + " raise ValueError(\"Either channel_name or signal_expr must be provided\")\n", + "\n", + " ev = build_event(event, db)\n", + " report.add_event(ev)\n", + " _resolve_weight(weight)\n", + "\n", + " page = Page(page_number=1)\n", + " page.add_aggregation(HistogramDuration(\n", + " name=\"adhoc_histogram\", base_expr=channel, bins=[float(b) for b in bins],\n", + " event=ev, channel_name=channel_name or \"virtual_signal\",\n", + " bins_unit=bins_unit or \"\", values_unit=values_unit,\n", + " ))\n", + " report.add_page(page)\n", + " report.determine_report()\n", + "\n", + " hist_dfs = report.aggregation_dfs[\"HISTOGRAM\"]\n", + " hist_df = hist_dfs.get(\"changed\") if hist_dfs.get(\"changed\") is not None else hist_dfs[\"unchanged\"]\n", + " result_df = (\n", + " hist_df.groupBy(\"bin_name\", \"lower_bound\")\n", + " .agg(F.sum(\"hist_value\").alias(\"duration_us\")).orderBy(\"lower_bound\").toPandas()\n", + " )\n", + " result_df[\"duration_s\"] = result_df[\"duration_us\"] / 1e6\n", + " return result_df[[\"bin_name\", \"lower_bound\", \"duration_s\"]].to_dict(orient=\"records\")\n", + "\n", + "\n", + "def run_adhoc_histogram2d(\n", + " x_bins: list[float],\n", + " y_bins: list[float],\n", + " x_channel_name: str | None = None,\n", + " y_channel_name: str | None = None,\n", + " tags: dict[str, str] | None = None,\n", + " x_signal_expr: dict | None = None,\n", + " y_signal_expr: dict | None = None,\n", + " event: dict | None = None,\n", + " weight: dict | None = None,\n", + " x_bins_unit: str | None = None,\n", + " y_bins_unit: str | None = None,\n", + " values_unit: str = \"s\",\n", + ") -> list[dict]:\n", + " \"\"\"\n", + " 2D duration-weighted heatmap of two signals (x vs y), scoped to an event.\n", + " Same expression-tree/event/weight mechanism as run_adhoc_histogram.\n", + " Returns one row per (x_bin, y_bin).\n", + " \"\"\"\n", + " report = Report(name=f\"adhoc2d_{uuid.uuid4().hex[:8]}\", spark=spark, config=_adhoc_config(), workspace_client=WorkspaceClient())\n", + " db = report.get_db()\n", + "\n", + " if x_signal_expr is not None:\n", + " x_channel = build_expr(x_signal_expr, db)\n", + " elif x_channel_name is not None:\n", + " x_channel = db.query.channel(channel_name=x_channel_name, **(tags or {}))\n", + " else:\n", + " raise ValueError(\"Either x_channel_name or x_signal_expr must be provided\")\n", + "\n", + " if y_signal_expr is not None:\n", + " y_channel = build_expr(y_signal_expr, db)\n", + " elif y_channel_name is not None:\n", + " y_channel = db.query.channel(channel_name=y_channel_name, **(tags or {}))\n", + " else:\n", + " raise ValueError(\"Either y_channel_name or y_signal_expr must be provided\")\n", + "\n", + " ev = build_event(event, db)\n", + " report.add_event(ev)\n", + " _resolve_weight(weight)\n", + "\n", + " page = Page(page_number=1)\n", + " page.add_aggregation(Histogram2DDuration(\n", + " name=\"adhoc_histogram2d\", x_expr=x_channel, y_expr=y_channel,\n", + " x_bins=[float(b) for b in x_bins], y_bins=[float(b) for b in y_bins], event=ev,\n", + " x_channel_name=x_channel_name or \"x_signal\", y_channel_name=y_channel_name or \"y_signal\",\n", + " x_bins_unit=x_bins_unit, y_bins_unit=y_bins_unit, values_unit=values_unit,\n", + " ))\n", + " report.add_page(page)\n", + " report.determine_report()\n", + "\n", + " hist_dfs = report.aggregation_dfs[\"HISTOGRAM2D\"]\n", + " hist_df = hist_dfs.get(\"changed\") if hist_dfs.get(\"changed\") is not None else hist_dfs[\"unchanged\"]\n", + " result_df = (\n", + " hist_df.groupBy(\"x_bin_name\", \"y_bin_name\", \"x_lower_bound\", \"y_lower_bound\")\n", + " .agg(F.sum(\"hist_value\").alias(\"duration_us\")).orderBy(\"x_lower_bound\", \"y_lower_bound\").toPandas()\n", + " )\n", + " result_df[\"duration_s\"] = result_df[\"duration_us\"] / 1e6\n", + " return result_df[[\"x_bin_name\", \"y_bin_name\", \"x_lower_bound\", \"y_lower_bound\", \"duration_s\"]].to_dict(orient=\"records\")\n", + "\n", + "\n", + "_STATS = {\"min\", \"max\", \"mean\", \"median\"}\n", + "\n", + "\n", + "def run_adhoc_stats(\n", + " signals: list[dict],\n", + " statistics: list[str] | None = None,\n", + " event: dict | None = None,\n", + ") -> list[dict]:\n", + " \"\"\"\n", + " Compute summary statistics for one or more signals at once, scoped to an\n", + " event -- this is the \"multiple channels\" case, via StatsAggregator.\n", + " signals: [{\"label\": \"...\", \"channel_name\": \"...\", \"tags\": {...}}] or\n", + " [{\"label\": \"...\", \"signal_expr\": }]. statistics: subset of\n", + " min/max/mean/median (default: all). Statistics are computed per\n", + " container, not collapsed across containers -- averaging an\n", + " already-averaged value across sessions without knowing sample counts\n", + " wouldn't be statistically valid. Returns one row per\n", + " (container_id, label, statistic).\n", + " \"\"\"\n", + " if not signals:\n", + " raise ValueError(\"signals must have at least one entry\")\n", + " stats = statistics or sorted(_STATS)\n", + " bad = set(stats) - _STATS\n", + " if bad:\n", + " raise ValueError(f\"Unsupported statistics {bad}; must be a subset of {sorted(_STATS)}\")\n", + "\n", + " report = Report(name=f\"adhoc_stats_{uuid.uuid4().hex[:8]}\", spark=spark, config=_adhoc_config(), workspace_client=WorkspaceClient())\n", + " db = report.get_db()\n", + "\n", + " ev = build_event(event, db)\n", + " report.add_event(ev)\n", + "\n", + " exprs, labels = [], []\n", + " for sig in signals:\n", + " if \"signal_expr\" in sig:\n", + " exprs.append(build_expr(sig[\"signal_expr\"], db))\n", + " elif \"channel_name\" in sig:\n", + " exprs.append(db.query.channel(channel_name=sig[\"channel_name\"], **sig.get(\"tags\", {})))\n", + " else:\n", + " raise ValueError(f\"signal {sig!r} needs channel_name or signal_expr\")\n", + " labels.append(sig[\"label\"])\n", + "\n", + " page = Page(page_number=1)\n", + " page.add_aggregation(StatsAggregator(\n", + " name=\"adhoc_stats\", input_expressions=exprs, channel_names=labels,\n", + " statistics=list(stats), event=ev,\n", + " ))\n", + " report.add_page(page)\n", + " report.determine_report()\n", + "\n", + " stats_dfs = report.aggregation_dfs[\"STATS_AGGREGATOR\"]\n", + " stats_df = stats_dfs.get(\"changed\") if stats_dfs.get(\"changed\") is not None else stats_dfs[\"unchanged\"]\n", + " result_df = (\n", + " stats_df.select(\"container_id\", \"channel_name\", \"aggregation_label\", \"statistic_value\")\n", + " .orderBy(\"container_id\", \"channel_name\", \"aggregation_label\")\n", + " .toPandas()\n", + " .rename(columns={\"channel_name\": \"label\", \"aggregation_label\": \"statistic\", \"statistic_value\": \"value\"})\n", + " )\n", + " return result_df.to_dict(orient=\"records\")\n", + "\n", + "\n", + "def run_adhoc_point_values(signals: list[dict], event: dict) -> list[dict]:\n", + " \"\"\"\n", + " Sample one or more signals at each instant of a points-in-time event --\n", + " the one case duration-weighted tools can't cover, since they need time\n", + " intervals, not discrete instants. signals: same shape as\n", + " run_adhoc_stats. event: REQUIRED, must be a points_in_time event.\n", + " Returns one row per (container_id, event_instance_id, label).\n", + " \"\"\"\n", + " if not signals:\n", + " raise ValueError(\"signals must have at least one entry\")\n", + " if not event or event.get(\"type\") != \"points_in_time\":\n", + " raise ValueError('event must be a points_in_time event: {\"type\": \"points_in_time\", \"signal_expr\": ...}')\n", + "\n", + " report = Report(name=f\"adhoc_pv_{uuid.uuid4().hex[:8]}\", spark=spark, config=_adhoc_config(), workspace_client=WorkspaceClient())\n", + " db = report.get_db()\n", + "\n", + " ev = build_event(event, db)\n", + " report.add_event(ev)\n", + "\n", + " exprs, labels = [], []\n", + " for sig in signals:\n", + " if \"signal_expr\" in sig:\n", + " exprs.append(build_expr(sig[\"signal_expr\"], db))\n", + " elif \"channel_name\" in sig:\n", + " exprs.append(db.query.channel(channel_name=sig[\"channel_name\"], **sig.get(\"tags\", {})))\n", + " else:\n", + " raise ValueError(f\"signal {sig!r} needs channel_name or signal_expr\")\n", + " labels.append(sig[\"label\"])\n", + "\n", + " page = Page(page_number=1)\n", + " page.add_aggregation(PointValueAggregator(\n", + " name=\"adhoc_point_values\", input_expressions=exprs, channel_names=labels, event=ev,\n", + " ))\n", + " report.add_page(page)\n", + " report.determine_report()\n", + "\n", + " pv_dfs = report.aggregation_dfs[\"POINT_VALUE_AGGREGATOR\"]\n", + " pv_df = pv_dfs.get(\"changed\") if pv_dfs.get(\"changed\") is not None else pv_dfs[\"unchanged\"]\n", + " result_df = (\n", + " pv_df.select(\"container_id\", \"event_instance_id\", \"channel_name\", \"statistic_value\")\n", + " .orderBy(\"container_id\", \"event_instance_id\", \"channel_name\")\n", + " .toPandas()\n", + " .rename(columns={\"channel_name\": \"label\", \"statistic_value\": \"value\"})\n", + " )\n", + " return result_df.to_dict(orient=\"records\")" + ], + "id": "2cd5a30ba3e14729" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Direct tests \u2014 prove the execution logic works against real data before\n", + "wiring it into MCP tools: the original plain-channel question, a virtual\n", + "signal, a 2D heatmap, multi-channel stats, and points-in-time sampling." + ], + "id": "bfdabfdf76ec4688" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result = run_adhoc_histogram(\n", + " channel_name=\"Engine RPM\",\n", + " bins=[0, 1000, 2000, 3000, 4000, 5000],\n", + " event={\"type\": \"basic\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " \"condition_op\": \">\", \"condition_value\": 2000},\n", + " tags={\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + ")\n", + "display(pd.DataFrame(result))" + ], + "id": "3541d455ff4649c3" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "avg_temp_expr = {\n", + " \"op\": \"div\",\n", + " \"left\": {\"op\": \"add\",\n", + " \"left\": {\"channel\": \"Ambient Air Temperature\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " \"right\": {\"channel\": \"Intake Air Temperature\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}}},\n", + " \"right\": {\"const\": 2},\n", + "}\n", + "result_virtual = run_adhoc_histogram(\n", + " signal_expr=avg_temp_expr,\n", + " event={\"type\": \"basic\", \"condition_expr\": {\"op\": \"gt\", \"left\": avg_temp_expr, \"right\": {\"const\": 0}}},\n", + " bins=[-10, 0, 10, 20, 30],\n", + ")\n", + "display(pd.DataFrame(result_virtual))" + ], + "id": "2548844177114428" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result_2d = run_adhoc_histogram2d(\n", + " x_channel_name=\"Engine RPM\",\n", + " y_channel_name=\"Vehicle Speed Sensor\",\n", + " x_bins=[2000, 2500, 3000, 3500, 4000, 4500, 5000],\n", + " y_bins=[0, 50, 100, 150, 200],\n", + " event={\"type\": \"basic\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " \"condition_op\": \">\", \"condition_value\": 2000},\n", + " tags={\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + ")\n", + "display(pd.DataFrame(result_2d))" + ], + "id": "264948788c704d2c" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result_stats = run_adhoc_stats(\n", + " signals=[\n", + " {\"label\": \"Engine RPM\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " {\"label\": \"Vehicle Speed\", \"channel_name\": \"Vehicle Speed Sensor\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " ],\n", + " statistics=[\"min\", \"max\", \"mean\"],\n", + ")\n", + "display(pd.DataFrame(result_stats))" + ], + "id": "10c2bd2e48ef4a55" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "result_points = run_adhoc_point_values(\n", + " signals=[\n", + " {\"label\": \"Vehicle Speed\", \"channel_name\": \"Vehicle Speed Sensor\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " {\"label\": \"Engine RPM\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " ],\n", + " event={\n", + " \"type\": \"points_in_time\",\n", + " \"signal_expr\": {\n", + " \"method\": \"start_points\",\n", + " \"operand\": {\"op\": \"gt\", \"left\": {\"channel\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}}, \"right\": {\"const\": 2000}},\n", + " },\n", + " },\n", + ")\n", + "display(pd.DataFrame(result_points))" + ], + "id": "3b2846d13fe94cb9" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 4. Exposing these as MCP tools\n", + "\n", + "Thin wrappers around the already-tested functions above \u2014 the MCP layer\n", + "adds no new logic, just a typed, discoverable interface an agent can call.\n", + "`Literal` on `condition_op` and `dict` on the expression-tree parameters\n", + "constrain what the tool schema will accept: the agent fills in values or\n", + "composes a tree from a fixed op vocabulary, it never generates code." + ], + "id": "18344cdcc90a4532" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Literal\n", + "from mcp.server.fastmcp import FastMCP\n", + "\n", + "mcp = FastMCP(\"impulse-agent\")\n", + "\n", + "\n", + "# list_channels/list_containers keep their Python function names as *_tool\n", + "# internally to avoid shadowing the plain list_channels()/list_containers()\n", + "# helpers above -- @mcp.tool(name=...) decouples the exposed MCP name from\n", + "# the Python function name.\n", + "\n", + "@mcp.tool(name=\"list_channels\")\n", + "def list_channels_tool() -> list[dict]:\n", + " \"\"\"List available measurement channels and their tags (e.g. brand, model,\n", + " unit -- whatever tag vocabulary this dataset uses). Read-only and\n", + " instant. Call this to ground yourself before preview_histogram/\n", + " preview_histogram_2d/preview_stats/preview_point_values (so channel_name/\n", + " tags are real values, not guesses).\"\"\"\n", + " return list_channels()\n", + "\n", + "\n", + "@mcp.tool(name=\"list_containers\")\n", + "def list_containers_tool() -> list[dict]:\n", + " \"\"\"List available measurement containers (recording sessions) and their\n", + " tags (e.g. vehicle, duration, condition). Read-only and instant.\"\"\"\n", + " return list_containers()\n", + "\n", + "\n", + "@mcp.tool()\n", + "def preview_histogram(\n", + " bins: list[float],\n", + " channel_name: str | None = None,\n", + " tags: dict[str, str] | None = None,\n", + " signal_expr: dict | None = None,\n", + " event: dict | None = None,\n", + " weight: dict | None = None,\n", + " bins_unit: str | None = None,\n", + " values_unit: str = \"s\",\n", + ") -> list[dict]:\n", + " \"\"\"Compute a duration-weighted histogram RIGHT NOW and return the actual\n", + " numeric result -- an instant, read-only preview, not a persisted report\n", + " aggregation. Provide either (channel_name [+ tags]) or signal_expr for a\n", + " derived virtual signal. event scopes the time range (omit for the whole\n", + " recording); weight currently only supports duration (the default) -- see\n", + " the notebook markdown above for the full expression tree / event /\n", + " known-limitation documentation. Returns one row per bin with bin_name,\n", + " lower_bound, and duration_s.\"\"\"\n", + " return run_adhoc_histogram(\n", + " bins=bins, channel_name=channel_name, tags=tags, signal_expr=signal_expr,\n", + " event=event, weight=weight, bins_unit=bins_unit, values_unit=values_unit,\n", + " )\n", + "\n", + "\n", + "@mcp.tool()\n", + "def preview_histogram_2d(\n", + " x_bins: list[float],\n", + " y_bins: list[float],\n", + " x_channel_name: str | None = None,\n", + " y_channel_name: str | None = None,\n", + " tags: dict[str, str] | None = None,\n", + " x_signal_expr: dict | None = None,\n", + " y_signal_expr: dict | None = None,\n", + " event: dict | None = None,\n", + " weight: dict | None = None,\n", + " x_bins_unit: str | None = None,\n", + " y_bins_unit: str | None = None,\n", + " values_unit: str = \"s\",\n", + ") -> list[dict]:\n", + " \"\"\"Compute a 2D duration-weighted heatmap of two signals (x vs y) RIGHT\n", + " NOW and return the actual numeric result -- an instant, read-only\n", + " preview, not a persisted report aggregation. Same expression-tree/event/\n", + " weight mechanism as preview_histogram. Returns one row per (x_bin, y_bin)\n", + " with x_bin_name, y_bin_name, x_lower_bound, y_lower_bound, duration_s.\"\"\"\n", + " return run_adhoc_histogram2d(\n", + " x_bins=x_bins, y_bins=y_bins, x_channel_name=x_channel_name, y_channel_name=y_channel_name,\n", + " tags=tags, x_signal_expr=x_signal_expr, y_signal_expr=y_signal_expr,\n", + " event=event, weight=weight, x_bins_unit=x_bins_unit, y_bins_unit=y_bins_unit, values_unit=values_unit,\n", + " )\n", + "\n", + "\n", + "@mcp.tool()\n", + "def preview_stats(\n", + " signals: list[dict],\n", + " statistics: list[Literal[\"min\", \"max\", \"mean\", \"median\"]] | None = None,\n", + " event: dict | None = None,\n", + ") -> list[dict]:\n", + " \"\"\"Compute summary statistics for one or more signals RIGHT NOW -- an\n", + " instant, read-only preview. This is the tool for statistics across\n", + " MULTIPLE signals at once. signals: [{\"label\": \"...\", \"channel_name\":\n", + " \"...\", \"tags\": {...}}] or [{\"label\": \"...\", \"signal_expr\": }].\n", + " statistics: subset of min/max/mean/median (default: all). event scopes\n", + " the time range (omit for the whole recording). Returns one row per\n", + " (container_id, label, statistic) -- per container, not collapsed across\n", + " containers, since averaging an already-averaged value without sample\n", + " counts wouldn't be statistically valid.\"\"\"\n", + " return run_adhoc_stats(signals=signals, statistics=statistics, event=event)\n", + "\n", + "\n", + "@mcp.tool()\n", + "def preview_point_values(signals: list[dict], event: dict) -> list[dict]:\n", + " \"\"\"Sample one or more signals at each instant of a points-in-time event\n", + " RIGHT NOW -- an instant, read-only preview. Use for \"what was X at each\n", + " Y\" questions (e.g. speed at each RPM-crossing instant) -- the one case\n", + " preview_histogram/preview_stats can't cover, since those need time\n", + " intervals, not discrete instants. signals: same shape as preview_stats.\n", + " event: REQUIRED, must be {\"type\": \"points_in_time\", \"signal_expr\": }. Returns one row per (container_id,\n", + " event_instance_id, label).\"\"\"\n", + " return run_adhoc_point_values(signals=signals, event=event)" + ], + "id": "2742e830590c4494" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 5. End-to-end MCP round trip\n", + "\n", + "An in-process MCP client talking to the server above \u2014 the same\n", + "request/response path a real agent (Claude, or any MCP-speaking client)\n", + "would use, just without a network hop. This is the actual proof that the\n", + "tool surface works, not mocked data.\n", + "\n", + "In production this server would run inside something with a **warm, live\n", + "Spark session** \u2014 e.g. a Databricks App using Databricks Connect / serverless\n", + "compute \u2014 since Impulse has no CLI or REST mode and a Job's cold-start\n", + "latency (cluster spin-up) isn't acceptable for a conversational agent." + ], + "id": "b3c6548bf0c446f1" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from mcp.shared.memory import create_connected_server_and_client_session\n", + "\n", + "\n", + "async def demo_agent_turn():\n", + " async with create_connected_server_and_client_session(mcp._mcp_server) as client:\n", + " tools = await client.list_tools()\n", + " print(\"Available tools:\", [t.name for t in tools.tools])\n", + "\n", + " channels = await client.call_tool(\"list_channels\", {})\n", + " print(\"\\nlist_channels ->\")\n", + " print(json.dumps(channels.structuredContent[\"result\"], indent=2))\n", + "\n", + " histogram = await client.call_tool(\n", + " \"preview_histogram\",\n", + " {\n", + " \"channel_name\": \"Engine RPM\",\n", + " \"bins\": [0, 1000, 2000, 3000, 4000, 5000],\n", + " \"event\": {\"type\": \"basic\", \"channel_name\": \"Engine RPM\",\n", + " \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " \"condition_op\": \">\", \"condition_value\": 2000},\n", + " \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " },\n", + " )\n", + " print(\"\\npreview_histogram ->\")\n", + " print(json.dumps(histogram.structuredContent[\"result\"], indent=2))\n", + "\n", + " heatmap = await client.call_tool(\n", + " \"preview_histogram_2d\",\n", + " {\n", + " \"x_channel_name\": \"Engine RPM\",\n", + " \"y_channel_name\": \"Vehicle Speed Sensor\",\n", + " \"x_bins\": [2000, 2500, 3000, 3500, 4000, 4500, 5000],\n", + " \"y_bins\": [0, 50, 100, 150, 200],\n", + " \"event\": {\"type\": \"basic\", \"channel_name\": \"Engine RPM\",\n", + " \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " \"condition_op\": \">\", \"condition_value\": 2000},\n", + " \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"},\n", + " },\n", + " )\n", + " print(\"\\npreview_histogram_2d ->\")\n", + " print(json.dumps(heatmap.structuredContent[\"result\"], indent=2))\n", + "\n", + " stats = await client.call_tool(\n", + " \"preview_stats\",\n", + " {\n", + " \"signals\": [\n", + " {\"label\": \"Engine RPM\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " {\"label\": \"Vehicle Speed\", \"channel_name\": \"Vehicle Speed Sensor\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " ],\n", + " \"statistics\": [\"min\", \"max\", \"mean\"],\n", + " },\n", + " )\n", + " print(\"\\npreview_stats ->\")\n", + " print(json.dumps(stats.structuredContent[\"result\"], indent=2))\n", + "\n", + " points = await client.call_tool(\n", + " \"preview_point_values\",\n", + " {\n", + " \"signals\": [\n", + " {\"label\": \"Vehicle Speed\", \"channel_name\": \"Vehicle Speed Sensor\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " {\"label\": \"Engine RPM\", \"channel_name\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}},\n", + " ],\n", + " \"event\": {\n", + " \"type\": \"points_in_time\",\n", + " \"signal_expr\": {\n", + " \"method\": \"start_points\",\n", + " \"operand\": {\"op\": \"gt\", \"left\": {\"channel\": \"Engine RPM\", \"tags\": {\"brand\": \"Seat\", \"model\": \"Leon\"}}, \"right\": {\"const\": 2000}},\n", + " },\n", + " },\n", + " },\n", + " )\n", + " print(\"\\npreview_point_values -> (first 3 rows)\")\n", + " print(json.dumps(points.structuredContent[\"result\"][:3], indent=2))\n", + "\n", + "\n", + "# Databricks notebooks (like Jupyter) already run inside an active asyncio\n", + "# event loop, so asyncio.run() isn't valid here -- use top-level await instead.\n", + "await demo_agent_turn()" + ], + "id": "7bd2a4163952452d" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 6. From notebook to production: deploying as a Databricks App\n", + "\n", + "The tools above prove the logic; production needs them reachable over the\n", + "network with a persistent warm Spark session, not run from a notebook cell.\n", + "\n", + "The full source for a deployable MCP server built from these same tools is\n", + "included at **`demos/agent_mcp_app/`** in this repo. To deploy it:\n", + "\n", + "```bash\n", + "cd demos/agent_mcp_app\n", + "./build_wheel.sh # bundles Impulse's own source for the app to ship to remote workers\n", + "# edit app.yaml: set CATALOG/SCHEMA/TABLE_PREFIX\n", + "databricks apps create mcp-impulse-agent # name must start with mcp- to be discoverable in AI Playground\n", + "databricks sync . /Workspace/Users//mcp-impulse-agent\n", + "databricks apps deploy mcp-impulse-agent --source-code-path /Workspace/Users//mcp-impulse-agent\n", + "```\n", + "\n", + "See that folder's `README.md` for the full setup steps (including the\n", + "Unity Catalog grants the app's service principal needs), the reasoning\n", + "behind its dependency pins and latency tuning, and known limitations\n", + "(no per-user auth yet; distance/custom-weighted histograms are gated off \u2014\n", + "see Section 3 above for why)." + ], + "id": "aa36ba4243d246a6" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 7. Testing the deployed agent in AI Playground\n", + "\n", + "Since the app is named `mcp-`, it's automatically discoverable as a\n", + "custom MCP server:\n", + "1. Workspace sidebar \u2192 **Playground**.\n", + "2. In the **Tools** dropdown, add the custom MCP server for your deployed\n", + " app.\n", + "3. Ask a question and watch which tools the model calls, in what order,\n", + " and with what arguments \u2014 this is the real test of whether the tool\n", + " descriptions and parameter schemas are good enough for a model to use\n", + " unprompted, not just whether the plumbing works.\n", + "\n", + "One question per category, from easiest to hardest for the agent:\n", + "\n", + "| Category | Question | What it tests |\n", + "|---|---|---|\n", + "| Grounding | \"What measurement channels are available in this dataset?\" | Calls `list_channels` with no guessing |\n", + "| Basic histogram | \"What's the distribution of Engine RPM while it's above 2000?\" | Maps a plain-English question onto `preview_histogram`'s typed parameters |\n", + "| Threshold / edge case | \"What's the RPM distribution above 10000?\" | Nothing in the data reaches that high \u2014 checks the agent reports an all-zero result correctly instead of treating it as an error |\n", + "| Disambiguation | \"I want to know about the RPM signal \u2014 what car is it from?\" | Should call `list_channels` to ground itself (discovering brand/model tags) rather than guessing |\n", + "| Compound / reasoning | \"Out of the total recording time, what percentage was spent above 2000 RPM?\" | Needs both `preview_histogram` and `list_containers` (for total duration), then arithmetic on the results |\n", + "| Should-fail / boundary | \"What's the Turbo Boost Pressure distribution?\" | This channel doesn't exist in the dataset \u2014 tests whether the agent checks `list_channels` and reports it's unavailable instead of hallucinating a result |\n", + "| 2D heatmap | \"Show me how RPM and speed relate to each other while RPM is above 2000.\" | Should call `preview_histogram_2d` rather than two separate 1D histograms |\n", + "| Virtual signal | \"What's the distribution of the average of ambient and intake air temperature?\" | Needs the agent to build a `signal_expr` tree combining two channels \u2014 the hardest case, since nothing about the tool description hands it a ready-made answer |\n", + "| Multi-signal stats | \"What are the min, max, and average RPM and speed for each test drive?\" | Should call `preview_stats` with both signals in one call, not two separate calls |\n", + "| Points in time | \"What was the vehicle speed each time RPM crossed above 2000?\" | Needs `preview_point_values` with a `points_in_time` event built from `start_points` on a comparison \u2014 the least discoverable tool, since it requires composing two concepts (event type + expression method) together |\n", + "\n", + "A couple of examples from an actual AI Playground run against the deployed\n", + "app:\n", + "\n", + "**Disambiguation** \u2014 the agent calls `list_channels` to ground itself before\n", + "answering, rather than guessing:\n", + "\n", + "\n", + "\n", + "**2D heatmap** \u2014 the agent calls `preview_histogram_2d` and summarizes the\n", + "RPM/speed relationship from the returned bins:\n", + "\n", + "" + ], + "id": "926e00d2cee94464" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 8. Cleanup\n", + "\n", + "Drop the silver tables loaded by this notebook. Ad-hoc scratch Gold tables\n", + "are already cleaned up inline by `run_adhoc_histogram` after each call." + ], + "id": "39c9c40b984e4396" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%skip\n", + "dbutils.widgets.dropdown(\"drop_created_tables\", \"false\", [\"true\", \"false\"], \"Drop Created Tables\")" + ], + "id": "da8499d2a5ff4812" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%skip\n", + "if dbutils.widgets.get(\"drop_created_tables\") == \"true\":\n", + " for t in SILVER:\n", + " spark.sql(f\"DROP TABLE IF EXISTS {pfx}_{t}\")\n", + " print(f\"Dropped {len(SILVER)} tables\")" + ], + "id": "879d340d77a64708" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/demos/images/playground_q4_disambiguation.png b/demos/images/playground_q4_disambiguation.png new file mode 100755 index 00000000..9af80c5a Binary files /dev/null and b/demos/images/playground_q4_disambiguation.png differ diff --git a/demos/images/playground_q7_2d_heatmap.png b/demos/images/playground_q7_2d_heatmap.png new file mode 100755 index 00000000..2b716551 Binary files /dev/null and b/demos/images/playground_q7_2d_heatmap.png differ