diff --git a/Makefile b/Makefile index 4d6dd47..9d5e15b 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help install dev lint format typecheck test test-unit test-e2e demo benchmark benchmark-pg benchmark-multi validate clean +.PHONY: help install dev lint format typecheck test test-unit test-e2e demo showcase benchmark benchmark-pg benchmark-multi validate clean help: ## Show this help @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-15s\033[0m %s\n", $$1, $$2}' @@ -31,6 +31,9 @@ test-e2e: ## Run end-to-end tests only demo: ## Run offline demo (no API keys needed) uv run behavioral-memory demo +showcase: ## Run unified demo showcase for talks (requires GOOGLE_API_KEY) + uv run python demo/showcase.py + benchmark: ## Run live benchmark with in-memory store (requires GOOGLE_API_KEY) uv run python examples/run_live_benchmark.py diff --git a/demo/SETUP_AND_RUN.md b/demo/SETUP_AND_RUN.md new file mode 100644 index 0000000..5596c51 --- /dev/null +++ b/demo/SETUP_AND_RUN.md @@ -0,0 +1,127 @@ +# Demo Setup & Run Guide + +## Quick Start + +```bash +git pull origin main +uv sync --extra agent --extra eval +uv run python demo/showcase.py +``` + +The demo initializes in ~2 seconds, then walks through 4 acts. Press Enter between acts to advance. + +## Run Options + +```bash +# Full demo (all 4 acts with pauses between them) +uv run python demo/showcase.py + +# Skip pauses (continuous output) +uv run python demo/showcase.py --no-pause + +# Run specific acts only +uv run python demo/showcase.py --act 2 # strategy comparison only +uv run python demo/showcase.py --act 1 --act 3 # memory + gatekeeper (instant, no LLM) +uv run python demo/showcase.py --act 2 --act 3 # comparison + gatekeeper + +# Custom query for Act 2 +uv run python demo/showcase.py --act 2 --query "Build a daily revenue pipeline and store for the dashboard" + +# Different model +uv run python demo/showcase.py --model gemini-2.5-pro + +# Via Makefile +make showcase +``` + +## What Each Act Does + +| Act | Name | LLM Calls | Time | What It Shows | +|-----|------|-----------|------|---------------| +| 1 | Memory Inspector | 0 | Instant | All 12 seed traces in the store + embedding neighborhood | +| 2 | Strategy Comparison | 3 | ~15-20s | Same query through zero-shot / static / dynamic with diff | +| 3 | Gatekeeper Challenge | 0 | Instant | Poisoned traces fed through 3 validation gates | +| 4 | Custom Query REPL | 3 per query | ~10s each | Interactive — type any query, type `quit` to exit | + +## Sample Queries for Act 4 + +These are designed to show clear differences between zero-shot and dynamic retrieval. Each one triggers specific domain conventions that the seed traces teach. + +### Best demo queries (show largest differences) + +``` +Build a daily revenue pipeline and store results for the dashboard +``` +Zero-shot wraps everything in a single `schedule_task`. Dynamic builds the full `query_database -> store_results` pipeline with cache target. + +``` +Schedule a weekly fulfillment rate report for the ops team +``` +Zero-shot produces a single `schedule_task`. Dynamic builds `query_database -> generate_report -> schedule_task` with the correct fulfillment formula. + +``` +Get net order values excluding discounts and store for the dashboard +``` +Dynamic retrieves the net order value trace and applies `total_amount - discount` with `transform_data`, stores to cache. + +``` +Archive all valid orders from last quarter as CSV +``` +Zero-shot may only exclude cancelled. Dynamic excludes both cancelled AND returned, uses append mode. + +### Revenue queries (tests quantity * unit_price convention) + +``` +Get monthly revenue breakdown and send to the finance team +``` + +``` +Calculate total revenue by product category and archive the results +``` + +### Alert queries (tests Slack #data-alerts convention) + +``` +Monitor order cancellation rates and alert if above threshold +``` + +``` +Check fulfillment metrics and notify the ops team +``` + +### Pipeline queries (tests query -> transform -> store pattern) + +``` +Build a customer lifetime value pipeline and cache for the dashboard +``` + +``` +Create a product performance ranking and store for analytics +``` + +### Report queries (tests markdown_table + email convention) + +``` +Generate a weekly summary of returned orders and email to management +``` + +``` +Find top 10 customers by spending and create a report +``` + +### Scheduling queries (tests daily + notify_on_failure convention) + +``` +Set up automated daily tracking of new customer signups +``` + +``` +Schedule a daily report of order volumes by region +``` + +## Tips + +- Type **one query at a time** in Act 4 and press Enter. Do not paste multiple lines at once. +- Type `quit`, `exit`, or `q` to end Act 4. +- If a query produces identical plans across strategies, try one of the pipeline or scheduling queries above — those show the biggest differences. +- The talk narration script is at `demo/TALK_SCRIPT.md`. diff --git a/demo/TALK_SCRIPT.md b/demo/TALK_SCRIPT.md new file mode 100644 index 0000000..8c4d5b0 --- /dev/null +++ b/demo/TALK_SCRIPT.md @@ -0,0 +1,184 @@ +# Demo Talk Script — Behavioral Memory for Tool Orchestration + +**Run command:** `uv run python demo/showcase.py` + +**Total runtime:** ~30 seconds of compute, rest is your narration. + +--- + +## OPENING (before running anything) + +> We're looking at a common problem with AI agents that use tools. Whether it's database queries, API calls, sending notifications — agents that orchestrate multiple tools tend to make the same mistakes over and over. Every time a new task comes in, the agent starts from zero. It sees a list of available tools and their parameters, and it guesses the best plan. +> +> The problem is that tool schemas tell you *what* a tool accepts, but not *how* your organization uses it. Which SQL column is revenue? Does "completed order" mean status equals 'completed', or does it mean shipped and delivered? Do alerts go to email or Slack? These are domain conventions — institutional knowledge that lives in people's heads. +> +> What we built is **behavioral memory** — a vector store of validated execution traces. Every time an agent successfully completes a task, we store the task description paired with the exact tool chain that worked. When a new task comes in, we embed the query, search for semantically similar past successes, and inject them into the prompt as reference examples. The LLM follows proven patterns instead of guessing. +> +> Let me show you what this looks like. + +*Run the demo:* `uv run python demo/showcase.py` + +The initialization takes about 1-2 seconds. It loads a Gemini model, initializes the embedding model, registers 7 MCP tool schemas, and seeds the memory with 12 validated traces. + +--- + +## ACT 1 — Memory Inspector + +*The table of 12 traces appears on screen.* + +> This is the behavioral memory store. These 12 traces are the agent's institutional knowledge. Each row is a validated execution trace — a task description paired with the ordered sequence of tool calls that accomplished it. +> +> Look at trace number 1: "Get Q1 revenue data and send a report to stakeholders." The tool chain is `query_database` then `generate_report` then `send_notification`. That's three steps. And the source column says "seed" — these traces were validated by domain experts and seeded into memory as ground truth. +> +> Now look at the embedding neighborhood below the table. We took the first trace and searched the vector store for its nearest neighbors. The most similar trace is "Monthly valid orders trend and email the report" at 0.747 cosine similarity. These are the kinds of traces that would get retrieved when a similar query comes in. +> +> What matters here is that each trace encodes domain conventions that are not in the tool schemas. For example, revenue is computed from `order_items` using `quantity * unit_price` — not from the `total_amount` column in the orders table. A "completed" order means status is 'shipped' or 'delivered', not a literal status called 'completed'. Alerts go to Slack channel `#data-alerts`, not email. Reports use `markdown_table` format. Dashboard data gets stored to cache, not to a database table. +> +> None of this is discoverable from the tool definitions alone. The agent has to learn it from examples. + +*Press Enter to continue.* + +--- + +## ACT 2 — Side-by-Side Strategy Comparison + +*The Act 2 panel appears with the query.* + +> Now we run the core experiment. We take a single query — "Calculate average basket size per customer segment and alert the data team" — and run it through three strategies. +> +> The first strategy is **zero-shot**: the model sees only the query and the tool schemas. No examples, no memory. This is how most agents work today. + +*Zero-shot results appear with pipeline timing.* + +> Notice the pipeline logging on the left — we can see exactly what's happening. The prompt is about 1,050 tokens with zero reference examples. The model takes a few seconds and produces a plan. +> +> Look at what zero-shot chose: `query_database`, then `generate_report`, then `send_notification` via email. It treated this as a reporting task. + +> The second strategy is **static few-shot**: the model gets the same three fixed examples for every query, regardless of what you ask. This is the traditional approach to few-shot prompting. + +*Static few-shot results appear.* + +> Similar result — `query_database`, `generate_report`, `send_notification` via email. The fixed examples didn't help here because they weren't relevant to basket sizes. + +> Now the third strategy — **dynamic retrieval** — our approach. Watch the pipeline. + +*Dynamic retrieval results scroll with timing.* + +> This is where behavioral memory kicks in. The system embeds the query, searches the vector store, and finds 6 candidate traces. The top match at 0.81 similarity is "Get basket sizes per order and alert about unusual patterns" — that's a direct hit. It also retrieves the fulfillment metrics trace and the pipeline trace. +> +> Three traces are selected within the token budget — 1,709 out of 3,500 allowed tokens. These get injected into the prompt as reference examples. +> +> Now look at the plan: `query_database`, then `transform_data`, then `send_notification` via **Slack to `#data-team`**. + +*Point at the comparison table.* + +> The comparison table makes it clear. Zero-shot and static both chose `generate_report` at step 2. Dynamic chose `transform_data`. Why? Because the retrieved trace for basket sizes taught the model that basket size means item count — a quantity aggregation, not a report. And the alert convention says use Slack, not email. + +*Point at the structural diff.* + +> The diff confirms it: `generate_report` was removed, `transform_data` was added. Same input query, materially different plan — because the agent had access to relevant past experience. +> +> This is the thesis of the paper: when you retrieve semantically similar validated traces and inject them into the prompt, the LLM follows proven domain conventions instead of guessing. + +*Press Enter to continue.* + +--- + +## ACT 3 — Gatekeeper Challenge + +*The Act 3 panel appears.* + +> The obvious next question is: what stops bad data from entering memory? If someone stores a trace with wrong conventions, doesn't the whole system degrade? +> +> That's what the gatekeeper pipeline addresses. Before any trace enters behavioral memory, it passes through three validation gates. + +*Candidate #1 appears — broken dependency.* + +> Candidate 1 is a trace for "Get product data and generate report." The tool chain is just `generate_report` — but look at the parameters: it references `source_step: s0_nonexistent`. That step doesn't exist in the trace. +> +> Gate 1, Schema Validation, checks that all tools exist and required parameters are present. This passes — `generate_report` is a real tool with valid params. +> +> Gate 2, Sandbox Execution, does a dry-run data-flow check. It walks through the steps and verifies that every `source_step` reference points to a step that has already produced output. Step `s1` references `s0_nonexistent` which doesn't exist — **Gate 2 fails**. The trace is rejected. It never enters memory. + +*Candidate #2 appears — wrong convention.* + +> Candidate 2 is more subtle. "Get quarterly revenue and send report." The tool chain is `query_database` then `generate_report`. Structurally, this is perfectly valid — real tools, valid parameters, correct data flow. +> +> Gate 1 passes. Gate 2 passes. Gate 3, Semantic Dedup, checks if this trace is too similar to an existing one. The nearest trace in memory has a similarity score of 0.843, which is below the 0.95 threshold — so it's not a duplicate. **All three gates pass. The trace is admitted.** +> +> But here's the thing: this trace uses `total_amount` for revenue instead of `quantity * unit_price`, and it uses CSV format instead of `markdown_table`. It's structurally valid but semantically wrong. + +*The Limitation panel appears.* + +> We're transparent about this. The gatekeeper validates **structure** — does the plan make sense as a sequence of tool calls? It does not validate **domain semantics** — is this the right business logic? That's a fundamentally harder problem. +> +> This is why the seed traces matter. They establish the correct conventions in memory. The gatekeeper prevents broken plans from entering. The seed traces ensure the conventions are right from the start. And in production, the feedback loop — where human-approved executions flow back through the gatekeeper into memory — continuously reinforces correct patterns. + +*Press Enter to continue.* + +--- + +## ACT 4 — Custom Query (if time permits) + +*The REPL prompt appears.* + +> Now I can take any query. Type a task description, and we'll run it through all three strategies in real time. You'll see the retrieval scores, the pipeline timing, and the plan comparison. + +**Good queries to demonstrate live:** + +- `Build a daily revenue pipeline and store results for the dashboard` + - *Shows: zero-shot just wraps everything in schedule_task; dynamic retrieves the pipeline pattern and builds query -> transform -> store with cache target* + +- `Schedule a weekly fulfillment rate report for the ops team` + - *Shows: zero-shot collapses everything into a single schedule_task call; dynamic retrieves the fulfillment metrics trace and builds the full workflow first, then wraps it in schedule_task* + +- `Get net order values excluding discounts and store for the dashboard` + - *Shows: zero-shot may skip the discount subtraction; dynamic retrieves the net order value trace and applies total_amount - discount* + +- `Archive all valid orders from last quarter as CSV` + - *Shows: zero-shot may only exclude cancelled; dynamic retrieves the valid orders trace and excludes both cancelled AND returned, uses append mode* + +*Type `quit` when done.* + +--- + +## CLOSING (after demo ends) + +> To summarize what we just saw: +> +> First, behavioral memory gives agents institutional knowledge. Instead of guessing which tools to use and how to use them, the agent retrieves validated patterns from past successes. +> +> Second, semantic retrieval means the right examples surface for the right query. A question about basket sizes retrieves the basket size trace, not a random revenue trace. This is the key advantage over static few-shot prompting. +> +> Third, the gatekeeper pipeline maintains memory quality. Structurally broken traces are caught and rejected before they can contaminate retrieval results. +> +> And fourth, the entire system is model-agnostic and framework-agnostic. The behavioral memory layer sits between your agent and the LLM. It works with any LangChain-compatible model and any MCP tool ecosystem. You bring your tools, you bring your model — behavioral memory provides the institutional knowledge layer. +> +> The benchmark results across 30 evaluation tasks show that dynamic retrieval improves plan correctness by 15-25 percentage points over zero-shot, with statistically significant results confirmed by McNemar's test. The full results are in the paper, Section IV. + +--- + +## Quick Reference — What to Point At + +| Moment | What to highlight | +|--------|------------------| +| Act 1 table | "Each row is a validated trace — task + tool chain" | +| Act 1 neighborhood | "Vector similarity — these traces are nearby in embedding space" | +| Act 2 pipeline log | "Real-time: embed query, retrieve traces, assemble prompt, call model" | +| Act 2 similarity bars | "0.81 similarity — the system found the right basket size trace" | +| Act 2 comparison table | "Same query, different plan — memory changed the tool selection" | +| Act 2 diff | "Red = removed, Green = added — structural difference in one glance" | +| Act 3 candidate #1 | "Broken data flow — Gate 2 catches it, trace rejected" | +| Act 3 candidate #2 | "Structurally valid but semantically wrong — passes all gates" | +| Act 3 limitation panel | "We're honest about what the gatekeeper does and doesn't catch" | + +--- + +## Troubleshooting During the Talk + +| Problem | Fix | +|---------|-----| +| Slow LLM response | "The model is thinking — you can see the pipeline stages updating in real time" | +| Plans look identical across strategies | Run with a different query (basket size, pipeline, scheduling queries show the biggest differences) | +| Act 4 query concatenation | Type one query at a time, press Enter, wait for results before typing the next | +| API key error | Make sure `.env` has `GOOGLE_API_KEY=your-key` and restart the script | diff --git a/demo/showcase.py b/demo/showcase.py new file mode 100644 index 0000000..d853ebd --- /dev/null +++ b/demo/showcase.py @@ -0,0 +1,631 @@ +"""Unified demo runner for conference talks and presentations. + +Initializes the full behavioral-memory pipeline once, then runs +interactive demo "acts" that showcase each layer of the system: + + Act 1 — Memory Inspector: Browse the trace store + Act 2 — Side-by-Side Compare: Zero-shot vs static vs dynamic retrieval + Act 3 — Gatekeeper Challenge: Feed poisoned traces through validation + Act 4 — Custom Query: Interactive REPL for audience questions + +Usage: + python demo/showcase.py # run all acts sequentially + python demo/showcase.py --act 2 # jump to a specific act + python demo/showcase.py --act 2 --query "…" # custom query for compare + python demo/showcase.py --model gemini-2.5-flash + +Requires: GOOGLE_API_KEY (set in .env or environment) +""" + +from __future__ import annotations + +import argparse +import difflib +import os +import sys +import time +from typing import Any + +from rich.console import Console +from rich.panel import Panel +from rich.table import Table + +from behavioral_memory.core.config import Settings +from behavioral_memory.core.schemas import ExecutionTrace, GatekeeperResult, Plan, ToolCall, ToolSchema +from behavioral_memory.evaluation.seed_traces import get_seed_traces +from behavioral_memory.gatekeeper.pipeline import GatekeeperPipeline +from behavioral_memory.memory.in_memory_store import InMemoryTraceStore +from behavioral_memory.memory.token_budget import count_tokens, select_traces_within_budget +from behavioral_memory.planner.engine import PlanEngine +from behavioral_memory.planner.prompt import SYSTEM_PROMPT, build_prompt +from behavioral_memory.tools.mock_tools import get_tool_schemas +from behavioral_memory.tools.registry import ToolRegistry + +console = Console() + +DEFAULT_QUERY = "Calculate average basket size per customer segment and alert the data team" + +POISONED_TRACES_RAW: list[dict[str, Any]] = [ + { + "task": "Get product data and generate report", + "chain": [ + {"step_id": "s1", "tool": "generate_report", "params": {"source_step": "s0_nonexistent", "format": "markdown_table", "title": "Products"}}, + ], + "poison_type": "broken_dependency", + "explanation": "References s0_nonexistent which doesn't exist; skips the query step entirely", + }, + { + "task": "Get quarterly revenue and send report", + "chain": [ + {"step_id": "s1", "tool": "query_database", "params": {"query": "SELECT total_amount FROM orders;"}}, + {"step_id": "s2", "tool": "generate_report", "params": {"source_step": "s1", "format": "csv", "title": "Revenue"}}, + ], + "poison_type": "wrong_convention", + "explanation": "Uses total_amount instead of quantity*unit_price for revenue; uses csv instead of markdown_table", + }, + { + "task": "Get completed orders and cache for dashboard", + "chain": [ + {"step_id": "s1", "tool": "query_database", "params": {"query": "SELECT * FROM orders WHERE status = 'completed';"}}, + {"step_id": "s2", "tool": "store_results", "params": {"source_step": "s1", "target": "database_table", "target_name": "dashboard_data"}}, + ], + "poison_type": "wrong_convention", + "explanation": "Uses status='completed' instead of IN('shipped','delivered'); stores to database_table instead of cache", + }, +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _build_poisoned_trace(raw: dict[str, Any]) -> ExecutionTrace: + chain = [ToolCall(step_id=s["step_id"], tool_name=s["tool"], parameters=s["params"]) for s in raw["chain"]] + return ExecutionTrace( + task_description=raw["task"], + tool_chain=chain, + validated=False, + source="execution", + metadata={"poison_type": raw["poison_type"], "explanation": raw["explanation"]}, + ) + + +def _tool_chain_str(steps: list[ToolCall]) -> str: + return " -> ".join(s.tool_name for s in steps) + + +def _pause(message: str = "Press Enter to continue...") -> None: + console.print(f"\n[dim]{message}[/dim]") + try: + input() + except (EOFError, KeyboardInterrupt): + pass + + +class PipelineTimer: + """Lightweight wrapper that prints timestamped pipeline events.""" + + def __init__(self) -> None: + self._t0 = time.perf_counter() + + def reset(self) -> None: + self._t0 = time.perf_counter() + + def _elapsed(self) -> float: + return time.perf_counter() - self._t0 + + def log(self, msg: str) -> None: + console.print(f" [dim][{self._elapsed():5.1f}s][/dim] {msg}") + + +def _timed_generate( + timer: PipelineTimer, + engine: PlanEngine, + store: InMemoryTraceStore, + query: str, + schemas: list[ToolSchema], + strategy: str, + static_traces: list[ExecutionTrace] | None = None, +) -> tuple[Plan, list[tuple[ExecutionTrace, float]]]: + """Run plan generation with pipeline event logging. + + Returns (plan, retrieved_traces_with_scores). + """ + timer.reset() + retrieved_with_scores: list[tuple[ExecutionTrace, float]] = [] + + if strategy == "zero_shot": + timer.log("Strategy: [yellow]zero-shot[/yellow] (no memory)") + prompt = build_prompt(query=query, traces=[], tool_schemas=schemas) + token_count = count_tokens(SYSTEM_PROMPT) + count_tokens(prompt) + timer.log(f"Assembling prompt ({token_count:,} tokens, 0 examples)") + timer.log("Calling model...") + plan = engine.generate_zero_shot(query, schemas) + timer.log(f"Plan parsed ({len(plan.steps)} steps)") + return plan, [] + + if strategy == "static": + traces = static_traces or [] + timer.log(f"Strategy: [blue]static few-shot[/blue] ({len(traces)} fixed examples)") + prompt = build_prompt(query=query, traces=traces, tool_schemas=schemas) + token_count = count_tokens(SYSTEM_PROMPT) + count_tokens(prompt) + timer.log(f"Assembling prompt ({token_count:,} tokens)") + timer.log("Calling model...") + plan = engine.generate_static_few_shot(query, schemas, traces) + timer.log(f"Plan parsed ({len(plan.steps)} steps)") + return plan, [] + + # Dynamic retrieval + timer.log("Strategy: [green]dynamic retrieval[/green] (behavioral memory)") + timer.log("Embedding query...") + retrieved_with_scores = store.search(query, k=6) + scores_str = ", ".join(f"{s:.2f}" for _, s in retrieved_with_scores[:5]) + timer.log(f"Retrieved {len(retrieved_with_scores)} candidates (scores: {scores_str})") + + settings = engine._settings + selected = select_traces_within_budget( + store=store, query=query, tool_schemas=schemas, settings=settings, + ) + timer.log(f"Selected {len(selected)} traces within token budget") + + prompt = build_prompt(query=query, traces=selected, tool_schemas=schemas) + token_count = count_tokens(SYSTEM_PROMPT) + count_tokens(prompt) + timer.log(f"Assembling prompt ({token_count:,} / {settings.max_prompt_tokens:,} tokens)") + timer.log("Calling model...") + plan = engine.generate(query=query, tool_schemas=schemas, traces=selected) + timer.log(f"Plan parsed ({len(plan.steps)} steps)") + return plan, retrieved_with_scores + + +# --------------------------------------------------------------------------- +# Act 1 — Memory Inspector +# --------------------------------------------------------------------------- + +def act_memory_inspector(store: InMemoryTraceStore) -> None: + console.print(Panel.fit( + "[bold]Act 1: Memory Inspector[/bold]\n\n" + "What's inside behavioral memory? Every trace the agent can\n" + "retrieve at query time — task descriptions, validated tool\n" + "chains, provenance, and similarity neighborhoods.", + title="ACT 1", + )) + + traces = store._traces + table = Table(title=f"Behavioral Memory Store ({len(traces)} traces)") + table.add_column("#", style="dim", width=3) + table.add_column("Task Description", max_width=52) + table.add_column("Steps", justify="center", width=5) + table.add_column("Tool Chain", style="cyan") + table.add_column("Source", justify="center", width=8) + + for i, trace in enumerate(traces, 1): + table.add_row( + str(i), + trace.task_description[:50] + ("..." if len(trace.task_description) > 50 else ""), + str(len(trace.tool_chain)), + _tool_chain_str(trace.tool_chain), + f"[green]{trace.source}[/green]" if trace.validated else trace.source, + ) + + console.print(table) + + # Show embedding neighborhood for one trace + if len(traces) >= 2: + probe = traces[0] + console.print(f"\n[bold]Embedding neighborhood[/bold] for: [italic]\"{probe.task_description[:60]}\"[/italic]") + neighbors = store.search(probe.task_description, k=4) + for rank, (t, score) in enumerate(neighbors): + if t.task_description == probe.task_description: + continue + bar_len = int(score * 30) + bar = "[green]" + "█" * bar_len + "[/green]" + "░" * (30 - bar_len) + console.print(f" {bar} {score:.3f} {t.task_description[:55]}") + + +# --------------------------------------------------------------------------- +# Act 2 — Side-by-Side Compare +# --------------------------------------------------------------------------- + +def _print_plan_steps(steps: list[ToolCall], color: str) -> None: + for step in steps: + params_short = ", ".join(f"{k}={repr(v)[:30]}" for k, v in list(step.parameters.items())[:3]) + console.print(f" [{color}]{step.step_id}[/{color}]: {step.tool_name}({params_short})") + + +def _print_diff(zs_steps: list[ToolCall], dyn_steps: list[ToolCall]) -> None: + """Print a structural diff of tool chains: zero-shot vs dynamic.""" + zs_tools = [s.tool_name for s in zs_steps] + dyn_tools = [s.tool_name for s in dyn_steps] + + diff = list(difflib.unified_diff( + zs_tools, dyn_tools, + fromfile="zero-shot", tofile="dynamic", + lineterm="", + )) + + if not diff: + console.print(" [dim]Plans are identical — no diff to show.[/dim]") + return + + console.print() + for line in diff: + if line.startswith("---"): + console.print(f" [bold red]{line}[/bold red]") + elif line.startswith("+++"): + console.print(f" [bold green]{line}[/bold green]") + elif line.startswith("@@"): + console.print(f" [dim]{line}[/dim]") + elif line.startswith("-"): + console.print(f" [red]{line} ← removed/changed[/red]") + elif line.startswith("+"): + console.print(f" [green]{line} ← added/changed[/green]") + else: + console.print(f" [dim]{line}[/dim]") + + +def act_compare( + engine: PlanEngine, + store: InMemoryTraceStore, + schemas: list[ToolSchema], + seed_traces: list[ExecutionTrace], + query: str, +) -> None: + console.print(Panel.fit( + "[bold]Act 2: Side-by-Side Strategy Comparison[/bold]\n\n" + "The same query is sent through three strategies:\n" + " [yellow]Zero-shot[/yellow] — model sees only tool schemas, no examples\n" + " [blue]Static few-shot[/blue] — 3 fixed examples regardless of query\n" + " [green]Dynamic retrieval[/green] — semantically similar traces from memory\n\n" + f"Query: [italic]{query}[/italic]", + title="ACT 2", + )) + + timer = PipelineTimer() + + # --- Zero-shot --- + console.print("\n[bold yellow]━━━ ZERO-SHOT (no memory) ━━━[/bold yellow]") + zs_plan, _ = _timed_generate(timer, engine, store, query, schemas, "zero_shot") + _print_plan_steps(zs_plan.steps, "yellow") + + # --- Static few-shot --- + console.print("\n[bold blue]━━━ STATIC FEW-SHOT (3 fixed examples) ━━━[/bold blue]") + static_traces = seed_traces[:3] + sf_plan, _ = _timed_generate(timer, engine, store, query, schemas, "static", static_traces) + _print_plan_steps(sf_plan.steps, "blue") + + # --- Dynamic retrieval --- + console.print("\n[bold green]━━━ DYNAMIC RETRIEVAL (behavioral memory) ━━━[/bold green]") + dyn_plan, retrieved = _timed_generate(timer, engine, store, query, schemas, "dynamic") + _print_plan_steps(dyn_plan.steps, "green") + + # --- Retrieved traces panel --- + if retrieved: + console.print("\n[bold]Retrieved traces (ranked by cosine similarity):[/bold]") + for rank, (trace, score) in enumerate(retrieved[:5], 1): + bar_len = int(score * 25) + bar = "█" * bar_len + "░" * (25 - bar_len) + console.print(f" [green]{bar}[/green] {score:.3f} #{rank}") + console.print(f" [italic]{trace.task_description[:70]}[/italic]") + console.print(f" [dim]{_tool_chain_str(trace.tool_chain)}[/dim]") + + # --- Comparison table --- + console.print() + max_steps = max(len(zs_plan.steps), len(sf_plan.steps), len(dyn_plan.steps)) + comp_table = Table(title="Plan Comparison (tool sequence)") + comp_table.add_column("Step", style="dim", width=6) + comp_table.add_column("Zero-Shot", style="yellow") + comp_table.add_column("Static Few-Shot", style="blue") + comp_table.add_column("Dynamic (Proposed)", style="bold green") + + for i in range(max_steps): + zs_tool = zs_plan.steps[i].tool_name if i < len(zs_plan.steps) else "—" + sf_tool = sf_plan.steps[i].tool_name if i < len(sf_plan.steps) else "—" + dy_tool = dyn_plan.steps[i].tool_name if i < len(dyn_plan.steps) else "—" + comp_table.add_row(f"#{i+1}", zs_tool, sf_tool, dy_tool) + + console.print(comp_table) + + # --- Structural diff --- + console.print("\n[bold]Structural diff (zero-shot vs dynamic):[/bold]") + _print_diff(zs_plan.steps, dyn_plan.steps) + + +# --------------------------------------------------------------------------- +# Act 3 — Gatekeeper Challenge +# --------------------------------------------------------------------------- + +def _run_gate_by_gate( + trace: ExecutionTrace, + gatekeeper: GatekeeperPipeline, +) -> GatekeeperResult: + """Run a trace through the gatekeeper, printing each gate's result.""" + # Gate 1: Schema + schema_valid, schema_failures = gatekeeper._schema_validator.validate(trace) + if schema_valid: + console.print(" Gate 1 (Schema Validation): [green]PASS[/green]") + else: + console.print(" Gate 1 (Schema Validation): [red]FAIL[/red]") + for f in schema_failures: + console.print(f" [red]> {f}[/red]") + console.print(" Gate 2 (Sandbox Execution): [dim]SKIPPED[/dim]") + console.print(" Gate 3 (Semantic Dedup): [dim]SKIPPED[/dim]") + console.print(" [bold red]VERDICT: REJECTED[/bold red]") + return GatekeeperResult( + accepted=False, schema_valid=False, + rejection_reason="Schema validation failed", failures=schema_failures, + ) + + # Gate 2: Sandbox + sandbox_passed, sandbox_detail = gatekeeper._sandbox.execute(trace) + if sandbox_passed: + console.print(" Gate 2 (Sandbox Execution): [green]PASS[/green]") + else: + console.print(" Gate 2 (Sandbox Execution): [red]FAIL[/red]") + console.print(f" [red]> {sandbox_detail}[/red]") + console.print(" Gate 3 (Semantic Dedup): [dim]SKIPPED[/dim]") + console.print(" [bold red]VERDICT: REJECTED[/bold red]") + return GatekeeperResult( + accepted=False, schema_valid=True, sandbox_passed=False, + rejection_reason=f"Sandbox check failed: {sandbox_detail}", + failures=[sandbox_detail], + ) + + # Gate 3: Dedup + is_unique, score = gatekeeper._dedup_gate.check(trace) + if is_unique: + console.print(f" Gate 3 (Semantic Dedup): [green]PASS[/green] (nearest: {score:.3f}, threshold: 0.95)") + else: + console.print(f" Gate 3 (Semantic Dedup): [red]FAIL[/red] (score: {score:.3f} >= 0.95)") + console.print(" [bold red]VERDICT: REJECTED (duplicate)[/bold red]") + return GatekeeperResult( + accepted=False, schema_valid=True, sandbox_passed=True, is_duplicate=True, + rejection_reason=f"Semantic duplicate (score={score:.3f})", + failures=[f"Duplicate score {score:.3f} >= threshold"], + ) + + console.print(" [bold green]VERDICT: ADMITTED[/bold green]") + return GatekeeperResult( + accepted=True, schema_valid=True, sandbox_passed=True, is_duplicate=False, + ) + + +def act_gatekeeper( + gatekeeper: GatekeeperPipeline, +) -> None: + console.print(Panel.fit( + "[bold]Act 3: Gatekeeper Challenge[/bold]\n\n" + "The gatekeeper pipeline runs three validation gates before\n" + "any trace enters memory:\n" + " 1. Schema Validation — tools exist, params valid, deps logical\n" + " 2. Sandbox Execution — dry-run data-flow check\n" + " 3. Semantic Dedup — cosine similarity < 0.95 to existing traces\n\n" + "We feed it deliberately poisoned traces and watch what happens.", + title="ACT 3", + )) + + for i, raw in enumerate(POISONED_TRACES_RAW, 1): + trace = _build_poisoned_trace(raw) + poison_label = raw["poison_type"].replace("_", " ") + + console.print(f"\n[bold]Candidate #{i}[/bold] [dim]({poison_label})[/dim]") + console.print(f" Task: [italic]{trace.task_description}[/italic]") + console.print(f" Chain: [cyan]{_tool_chain_str(trace.tool_chain)}[/cyan]") + console.print(f" Defect: [dim]{raw['explanation']}[/dim]") + console.print() + + _run_gate_by_gate(trace, gatekeeper) + + # --- The honest moment --- + console.print(Panel.fit( + "[bold]The honest limitation (Section IV-F)[/bold]\n\n" + "Notice that the 'wrong convention' traces PASS all three gates.\n" + "They use real tools, valid parameters, and correct data flow.\n" + "The gatekeeper checks [italic]structure[/italic], not [italic]domain semantics[/italic].\n\n" + "A trace that computes revenue from total_amount instead of\n" + "quantity*unit_price is structurally valid but semantically wrong.\n" + "This is exactly the paper's honest finding — the gatekeeper\n" + "catches broken traces, not subtly incorrect ones.\n\n" + "That's why the [green]seed traces[/green] matter: they teach the\n" + "correct conventions that the LLM learns to follow.", + title="Limitation", + )) + + +# --------------------------------------------------------------------------- +# Act 4 — Custom Query REPL +# --------------------------------------------------------------------------- + +def act_custom_query( + engine: PlanEngine, + store: InMemoryTraceStore, + schemas: list[ToolSchema], + seed_traces: list[ExecutionTrace], +) -> None: + console.print(Panel.fit( + "[bold]Act 4: Custom Query (interactive)[/bold]\n\n" + "Type any task and see how all three strategies handle it.\n" + "Type [bold]quit[/bold] or [bold]exit[/bold] to end.", + title="ACT 4", + )) + + timer = PipelineTimer() + static_traces = seed_traces[:3] + + while True: + console.print() + try: + query = console.input("[bold]Query > [/bold]").strip() + except (EOFError, KeyboardInterrupt): + break + + if not query or query.lower() in ("quit", "exit", "q"): + break + + console.print(f"\n[dim]Running 3 strategies for: \"{query}\"[/dim]\n") + + # Zero-shot + console.print("[bold yellow]━━━ ZERO-SHOT ━━━[/bold yellow]") + try: + zs_plan, _ = _timed_generate(timer, engine, store, query, schemas, "zero_shot") + _print_plan_steps(zs_plan.steps, "yellow") + except Exception as e: + console.print(f" [red]Error: {e}[/red]") + zs_plan = None + + # Static + console.print("\n[bold blue]━━━ STATIC FEW-SHOT ━━━[/bold blue]") + try: + sf_plan, _ = _timed_generate(timer, engine, store, query, schemas, "static", static_traces) + _print_plan_steps(sf_plan.steps, "blue") + except Exception as e: + console.print(f" [red]Error: {e}[/red]") + sf_plan = None + + # Dynamic + console.print("\n[bold green]━━━ DYNAMIC RETRIEVAL ━━━[/bold green]") + try: + dyn_plan, retrieved = _timed_generate(timer, engine, store, query, schemas, "dynamic") + _print_plan_steps(dyn_plan.steps, "green") + if retrieved: + console.print("\n [dim]Top retrieved traces:[/dim]") + for rank, (t, score) in enumerate(retrieved[:3], 1): + console.print(f" {score:.3f} {t.task_description[:60]}") + except Exception as e: + console.print(f" [red]Error: {e}[/red]") + dyn_plan = None + + # Quick comparison table + if zs_plan and sf_plan and dyn_plan: + max_steps = max(len(zs_plan.steps), len(sf_plan.steps), len(dyn_plan.steps)) + console.print() + t = Table(title="Comparison") + t.add_column("Step", style="dim", width=5) + t.add_column("Zero-Shot", style="yellow") + t.add_column("Static", style="blue") + t.add_column("Dynamic", style="bold green") + for j in range(max_steps): + zs = zs_plan.steps[j].tool_name if j < len(zs_plan.steps) else "—" + sf = sf_plan.steps[j].tool_name if j < len(sf_plan.steps) else "—" + dy = dyn_plan.steps[j].tool_name if j < len(dyn_plan.steps) else "—" + t.add_row(f"#{j+1}", zs, sf, dy) + console.print(t) + + console.print("\n[bold]Diff (zero-shot vs dynamic):[/bold]") + _print_diff(zs_plan.steps, dyn_plan.steps) + + console.print("[dim]Exiting interactive mode.[/dim]") + + +# --------------------------------------------------------------------------- +# Main — initialization & menu +# --------------------------------------------------------------------------- + +def _init_runtime(model_name: str | None = None) -> dict[str, Any]: + """One-time initialization of LLM, embeddings, store, and tools.""" + from dotenv import load_dotenv + from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings + + load_dotenv() + + model = model_name or os.getenv("GEMINI_MODEL", "gemini-2.5-flash") + console.print(f"[dim]Initializing LLM: {model}[/dim]") + llm = ChatGoogleGenerativeAI(model=model, temperature=0) + + console.print("[dim]Initializing embeddings: gemini-embedding-001[/dim]") + embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001") + + settings = Settings() + store = InMemoryTraceStore(embeddings=embeddings, settings=settings) + + console.print("[dim]Loading 7 tool schemas...[/dim]") + schemas = get_tool_schemas() + registry = ToolRegistry() + registry.register_many(schemas) + + console.print("[dim]Seeding 12 validated traces...[/dim]") + seed_traces = get_seed_traces() + store.add_bulk(seed_traces) + console.print(f"[dim]Store ready: {store.count()} traces embedded[/dim]") + + engine = PlanEngine(llm=llm, store=store, registry=registry, settings=settings) + gatekeeper = GatekeeperPipeline(store=store, registry=registry, settings=settings) + + return { + "llm": llm, + "store": store, + "schemas": schemas, + "registry": registry, + "seed_traces": seed_traces, + "engine": engine, + "gatekeeper": gatekeeper, + "settings": settings, + } + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Behavioral Memory — Unified Demo Showcase", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Examples:\n" + " python demo/showcase.py # all acts\n" + " python demo/showcase.py --act 2 # just compare\n" + ' python demo/showcase.py --act 2 --query "..." # custom query\n' + " python demo/showcase.py --act 1 --act 3 # memory + gatekeeper\n" + ), + ) + parser.add_argument( + "--act", type=int, action="append", choices=[1, 2, 3, 4], + help="Run specific act(s). Omit to run all sequentially.", + ) + parser.add_argument("--query", type=str, default=DEFAULT_QUERY, help="Query for Act 2 compare") + parser.add_argument("--model", type=str, default=None, help="LLM model name override") + parser.add_argument("--no-pause", action="store_true", help="Skip pauses between acts") + args = parser.parse_args() + + acts_to_run = sorted(args.act) if args.act else [1, 2, 3, 4] + + console.print(Panel.fit( + "[bold]Behavioral Memory — Demo Showcase[/bold]\n\n" + "\"Semantic Retrieval of Validated Execution Traces\n" + " in MCP-Based Agent Systems\" (IEEE, 2025)\n\n" + f"Acts: {', '.join(str(a) for a in acts_to_run)}", + title="DEMO", + )) + + # --- One-time initialization --- + t0 = time.perf_counter() + try: + runtime = _init_runtime(args.model) + except Exception as e: + console.print(f"\n[bold red]Initialization failed:[/bold red] {e}") + console.print("[dim]Make sure GOOGLE_API_KEY is set in your environment or .env file.[/dim]") + sys.exit(1) + init_time = time.perf_counter() - t0 + console.print(f"[dim]Initialization complete in {init_time:.1f}s[/dim]\n") + + # --- Run acts --- + for act_num in acts_to_run: + if act_num == 1: + act_memory_inspector(runtime["store"]) + elif act_num == 2: + act_compare( + runtime["engine"], runtime["store"], runtime["schemas"], + runtime["seed_traces"], args.query, + ) + elif act_num == 3: + act_gatekeeper(runtime["gatekeeper"]) + elif act_num == 4: + act_custom_query( + runtime["engine"], runtime["store"], runtime["schemas"], + runtime["seed_traces"], + ) + + if not args.no_pause and act_num != acts_to_run[-1] and act_num != 4: + _pause() + + console.print(Panel.fit("[bold]Demo complete.[/bold]", title="END")) + + +if __name__ == "__main__": + main()