diff --git a/.env.example b/.env.example index 9e7be03..98c4f36 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,7 @@ LANGFUSE_HOST=https://cloud.langfuse.com FEW_SHOT_K=3 # Number of traces to retrieve per query MAX_PROMPT_TOKENS=3500 # Token budget for the prompt SIMILARITY_DEDUP_THRESHOLD=0.95 # Cosine threshold for deduplication (Section III.E.3) +SANDBOX_TIMEOUT_SECONDS=30 # Gatekeeper sandbox execution timeout # === Feedback Loop === FEEDBACK_SCORE_NAME=quality # Langfuse score name to watch diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a4879c1..d4c50cf 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -5,6 +5,7 @@ on: tags: ["v*"] permissions: + contents: read id-token: write jobs: @@ -16,3 +17,6 @@ jobs: - uses: astral-sh/setup-uv@v4 - run: uv build - uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist/ + diff --git a/.gitignore b/.gitignore index 5b14b6f..41e1bf0 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,6 @@ Thumbs.db uv.lock -benchmark_results.json +benchmark_results*.json .chainlit/ chainlit.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 353638d..19cafd7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -59,6 +59,16 @@ git commit -m "your message" To run checks manually at any time: ```bash +# Using make (recommended) +make lint # Lint +make format # Auto-format +make typecheck # Type check (mypy strict) +make test # Run all tests +make ci # All checks (lint + format + typecheck + test) +make validate # Pipeline validation (no API keys needed) +make demo # Offline demo + +# Or directly uv run ruff check src/ tests/ agent/ # Lint only uv run ruff check --fix src/ tests/ agent/ # Lint + auto-fix uv run ruff format src/ tests/ agent/ # Format @@ -69,7 +79,7 @@ uv run pytest # Run all tests ## Running Tests ```bash -# All 104 tests (no external services needed) +# All tests (no external services needed) uv run pytest # With verbose output diff --git a/README.md b/README.md index f62c14f..fd4fe36 100644 --- a/README.md +++ b/README.md @@ -1,563 +1,330 @@ # behavioral-memory -**Validated execution traces as memory for MCP-based agent tool orchestration.** +**Give your agent institutional memory. Drop-in retrieval of validated execution traces for any LLM agent framework.** [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE) [![Python](https://img.shields.io/badge/python-3.11%2B-blue)](https://www.python.org) [![CI](https://github.com/harsh-kr11/behavioral-memory/actions/workflows/ci.yml/badge.svg)](https://github.com/harsh-kr11/behavioral-memory/actions/workflows/ci.yml) -A retrieval-based framework that uses a memory bank of validated execution traces to guide LLM tool orchestration during inference. Instead of relying on static few-shot examples, the system dynamically retrieves semantically similar, validated traces from past successful executions — giving your agent **institutional memory** that improves with every interaction. +Your agent makes the same mistakes repeatedly because it has no memory of what worked before. **behavioral-memory** fixes this — it stores validated execution traces (task → tool chain mappings) and retrieves semantically similar ones at query time, so your agent learns from past successes instead of starting from scratch every time. -Based on the IEEE paper: *"Behavioral Memory for Tool Orchestration: Semantic Retrieval of Validated Execution Traces in MCP-Based Agent Systems"* +Based on: *"Behavioral Memory for Tool Orchestration: Semantic Retrieval of Validated Execution Traces in MCP-Based Agent Systems"* (IEEE, 2025) --- -## Key Results (from the paper) - -On a 30-task benchmark with 7 MCP tools, using Gemini 2.5 Pro: - -| Metric | Zero-Shot | Static Few-Shot | **Proposed** | -|--------|-----------|----------------|-------------| -| Tool Selection (TSA) | 63.3% | 70.0% | **83.3%** | -| Parameter Validity (PV) | 72.2% | 79.6% | **84.0%** | -| Plan Correctness (PCR) | 33.3% | 50.0% | **63.3%** | -| Sequence Accuracy (ESA) | 63.3% | 70.0% | **83.3%** | - -McNemar's test: **p = 0.004** vs zero-shot. - -**Reproduced live run** (gemini-2.5-pro, pgvector, May 2026): - -| Metric | Zero-Shot | Static Few-Shot | **Proposed** | -|--------|-----------|----------------|-------------| -| TSA | 66.7% | 80.0% | **86.7%** | -| PV | 63.8% | 74.7% | **82.2%** | -| PCR | 53.3% | 70.0% | **80.0%** | -| ESA | 66.7% | 80.0% | **86.7%** | - -McNemar's test: **p = 0.039** vs zero-shot (statistically significant). - -> All reproduced metrics fall within the paper's 95% bootstrap confidence intervals. See [Running the Real Benchmark](#running-the-real-benchmark) to reproduce yourself. - ---- - -## Quick Start - -### Option A: No API keys needed (validation + demo) +## Install ```bash -git clone https://github.com/harsh-kr11/behavioral-memory.git -cd behavioral-memory -pip install -e ".[agent,eval,dev]" - -# Validate the entire pipeline (30/30 checks, no external services) -python examples/validate_pipeline.py - -# Quick demo showing behavioral memory impact -behavioral-memory demo +pip install behavioral-memory ``` -### Option B: With a Google API key (real benchmark) +--- -```bash -export GOOGLE_API_KEY=your-key-here -python examples/run_live_benchmark.py # all 30 tasks -python examples/run_live_benchmark.py --limit 5 # quick test with 5 tasks -python examples/run_live_benchmark.py --model gemini-2.0-flash # cheaper model -``` +## Plug Into Your Agent (3 lines) -### Option C: Interactive agent +The library is **framework-agnostic**. You bring your own LLM, your own agent — behavioral-memory handles the memory layer. -```bash -export GOOGLE_API_KEY=your-key-here -python -m agent.app --interactive +### Core API -# Or single query: -python -m agent.app "Build a revenue analysis pipeline" -``` +```python +from behavioral_memory import PlanEngine, ToolRegistry, InMemoryTraceStore ---- +# 1. Choose your LLM (any LangChain-compatible model) +from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings +llm = ChatGoogleGenerativeAI(model="gemini-2.5-pro", temperature=0) +embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001") -## How It Works +# 2. Create a memory store (no database needed) +store = InMemoryTraceStore(embeddings=embeddings) +# 3. Generate plans with behavioral memory +engine = PlanEngine(llm=llm, store=store) +plan = engine.generate(query="Get revenue data and email a report") ``` -User Query - │ - ▼ -┌─────────────────────────────────────────────────────┐ -│ 1. BEHAVIORAL LAYER │ -│ Retrieve top-k similar traces from memory │ -│ (pgvector or in-memory — your choice) │ -│ │ -│ 2. TOOL LAYER │ -│ Fetch available tool schemas via MCP │ -│ │ -│ 3. EXECUTIVE LAYER │ -│ Assemble 3-layer prompt → LLM → JSON plan │ -└──────────────────────────┬──────────────────────────┘ - │ - ▼ - Execution Plan - (ordered tool calls) - │ - ▼ -┌──────────────────────────────────────────────────────┐ -│ GATEKEEPER PIPELINE │ -│ ┌──────────────┬──────────────┬──────────────────┐ │ -│ │ Schema │ Sandboxed │ Semantic │ │ -│ │ Validation │ Execution │ Deduplication │ │ -│ └──────────────┴──────────────┴──────────────────┘ │ -│ Only validated traces enter memory │ -└──────────────────────────────────────────────────────┘ - │ - ▼ - ┌────────────────────┐ - │ Langfuse │ - │ (trace + feedback) │ - └────────────────────┘ -``` - ---- - -## Two Ways to Use -### 1. As a Library (Bring Your Own Agent) +That's it. Your agent now has memory. -Install and plug into your existing agent: - -```bash -pip install behavioral-memory -``` +### With OpenAI ```python -from behavioral_memory import PlanEngine, ToolRegistry, InMemoryTraceStore from langchain_openai import ChatOpenAI, OpenAIEmbeddings llm = ChatOpenAI(model="gpt-4o", temperature=0) -embeddings = OpenAIEmbeddings() - -# No PostgreSQL needed — InMemoryTraceStore works anywhere -store = InMemoryTraceStore(embeddings=embeddings) -registry = ToolRegistry() -engine = PlanEngine(llm=llm, store=store, registry=registry) - -plan = engine.generate(query="Get revenue data and email a report") +store = InMemoryTraceStore(embeddings=OpenAIEmbeddings()) +engine = PlanEngine(llm=llm, store=store) ``` -For production with PostgreSQL + pgvector: +### With Ollama (fully local) ```python -from behavioral_memory import TraceStore +from langchain_ollama import ChatOllama, OllamaEmbeddings -store = TraceStore(embeddings=embeddings, connection_url="postgresql+psycopg://...") +llm = ChatOllama(model="llama3") +store = InMemoryTraceStore(embeddings=OllamaEmbeddings(model="nomic-embed-text")) +engine = PlanEngine(llm=llm, store=store) ``` -### 2. Run the Reference Agent (LangGraph 1.x) - -```bash -git clone https://github.com/harsh-kr11/behavioral-memory.git -cd behavioral-memory -pip install -e ".[agent]" - -export GOOGLE_API_KEY=your-key +### Production: PostgreSQL + pgvector -# Interactive mode -python -m agent.app --interactive +```python +from behavioral_memory import TraceStore # pip install behavioral-memory[postgres] -# Single query -python -m agent.app "Build a revenue analysis pipeline" +store = TraceStore( + embeddings=embeddings, + connection_url="postgresql+psycopg://user:pass@localhost/behavioral_memory", +) ``` -The interactive agent supports: -- `/compare ` — run with AND without memory, see the difference -- `/memory` — inspect what's in behavioral memory -- `/quit` — exit - --- -## Running the Real Benchmark - -The benchmark sends 30 tasks through 3 strategies (zero-shot, static few-shot, dynamic retrieval), scoring each plan against gold tool chains. +## How It Helps Your Agent -### Prerequisites +Before behavioral memory, your agent sees only the task and tool schemas — it has to figure out orchestration from scratch every time. With behavioral memory, it retrieves validated examples of similar tasks that worked before. -Only a Google API key is required. PostgreSQL is optional — the benchmark defaults to `InMemoryTraceStore`, but for exact paper reproduction use `--postgres`. - -```bash -pip install -e ".[agent,eval]" -export GOOGLE_API_KEY=your-key-here - -# Optional: for pgvector mode (paper reproduction) -pip install -e ".[postgres]" -podman-compose up -d # or: docker compose up -d ``` - -### Run - -```bash -# Full benchmark (30 tasks × 3 strategies = 90 LLM calls) -python examples/run_live_benchmark.py - -# Quick test (5 tasks × 3 strategies = 15 LLM calls) -python examples/run_live_benchmark.py --limit 5 - -# Use a cheaper/faster model -python examples/run_live_benchmark.py --model gemini-2.0-flash - -# With PostgreSQL+pgvector (reproduces paper numbers exactly) -podman-compose up -d # or: docker compose up -d -python examples/run_live_benchmark.py --postgres - -# With Langfuse logging -export LANGFUSE_SECRET_KEY=sk-lf-... -export LANGFUSE_PUBLIC_KEY=pk-lf-... -python examples/run_live_benchmark.py +Your Agent's Query: "Build a revenue analysis pipeline" + │ + ┌────────────┴────────────┐ + │ BEHAVIORAL MEMORY │ + │ │ + │ 1. Retrieve top-k │ ← finds 3 similar validated traces + │ similar traces │ from past successful executions + │ │ + │ 2. Merge with tool │ ← current MCP tool schemas + │ schemas │ + │ │ + │ 3. Generate plan │ ← LLM sees examples + schemas + query + └────────────┬────────────┘ + │ + ▼ + Better execution plan + (right tools, right params, right order) ``` -### What you'll see +### Seed your memory with domain knowledge -``` -Benchmark Results (N=30, model=gemini-2.5-pro) -┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓ -┃ Metric ┃ Zero-Shot ┃ Static Few-Shot ┃ Dynamic (Proposed) ┃ -┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩ -│ TSA │ 63.3% [53%, 73%] │ 70.0% [56%, 83%] │ 83.3% [70%, 93%] │ -│ PV │ 72.2% │ 79.6% │ 84.0% │ -│ PCR │ 33.3% [16%, 50%] │ 50.0% [33%, 66%] │ 63.3% [46%, 80%] │ -│ ESA │ 63.3% [46%, 80%] │ 70.0% [53%, 86%] │ 83.3% [70%, 93%] │ -└────────┴──────────────────┴─────────────────────┴──────────────────────────┘ +```python +from behavioral_memory import ExecutionTrace, ToolCall + +trace = ExecutionTrace( + task_description="Calculate quarterly revenue", + tool_chain=[ + ToolCall(step_id="s1", tool_name="query_database", + parameters={"query": "SELECT SUM(quantity * unit_price) FROM order_items"}), + ToolCall(step_id="s2", tool_name="generate_report", + parameters={"source_step": "s1", "format": "markdown_table"}), + ], + source="seed", +) +store.add(trace) ``` -Results include per-task breakdowns, difficulty-tier analysis, and McNemar's test. +### Register your own tool schemas -### Reproducing Paper Numbers Exactly +The `PlanEngine` needs to know what tools your agent has: -The paper used PostgreSQL+pgvector for trace storage. The in-memory store gives equivalent TSA/ESA results but lower PV/PCR due to differences in nearest-neighbor retrieval fidelity. To reproduce the exact paper numbers: - -```bash -# 1. Start PostgreSQL+pgvector -podman-compose up -d # or: docker compose up -d - -# 2. Install postgres extras -pip install -e ".[postgres,agent,eval]" +```python +from behavioral_memory import ToolSchema, ToolRegistry + +schema = ToolSchema( + name="search_docs", + description="Search internal documentation", + parameters_schema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, +) -# 3. Run with the paper's model and store -python examples/run_live_benchmark.py --postgres --model gemini-2.5-pro +registry = ToolRegistry() +registry.register(schema) +engine = PlanEngine(llm=llm, store=store, registry=registry) ``` -| Setup | TSA | PV | PCR | ESA | McNemar p | -|-------|-----|-----|-----|-----|-----------| -| Paper | 83.3% | 84.0% | 63.3% | 83.3% | 0.004 | -| `--postgres` (live) | 86.7% | 82.2% | 80.0% | 86.7% | 0.039 | +Or load schemas dynamically from an MCP server: -> All results fall within the paper's 95% bootstrap confidence intervals. McNemar's test confirms statistical significance (p < 0.05). - ---- - -## Pipeline Validation (No API Keys) - -Validates every component works correctly using mock services: +```python +from behavioral_memory.tools.mcp_client import fetch_mcp_schemas -```bash -python examples/validate_pipeline.py +schemas = await fetch_mcp_schemas("http://localhost:3000/sse") +registry.register_many(schemas) ``` -This verifies: -- 12 seed traces load and pass schema validation -- 30 ground truth tasks have correct structure -- InMemoryTraceStore embeds, stores, and retrieves traces -- PlanEngine generates plans (zero-shot, static, dynamic) -- BenchmarkRunner scores and compares strategies -- Gatekeeper pipeline accepts/rejects traces -- Langfuse tracer handles offline mode gracefully - -All **30 checks** pass with zero external dependencies. - ---- - -## Installation - -### Prerequisites +### Validate before storing (Gatekeeper Pipeline) -- Python 3.11+ -- (Optional) PostgreSQL with [pgvector](https://github.com/pgvector/pgvector) for production deployments +Don't let bad traces into memory. The gatekeeper runs three checks before accepting a trace: -### Install with uv (recommended) +```python +from behavioral_memory import GatekeeperPipeline -```bash -uv add behavioral-memory # core framework (no PostgreSQL needed) -uv add behavioral-memory[agent] # + reference LangGraph agent -uv add behavioral-memory[eval] # + evaluation/statistics (scipy) -uv add behavioral-memory[postgres] # + PostgreSQL/pgvector store -uv add behavioral-memory[all] # everything +gatekeeper = GatekeeperPipeline(store=store, registry=registry) +result = gatekeeper.submit(trace) # schema check → sandbox → dedup → store +print(result.accepted) # True if all gates passed ``` -### Install with pip +### Learn from production (Langfuse Feedback Loop) -```bash -pip install behavioral-memory -pip install behavioral-memory[agent,eval] -pip install behavioral-memory[postgres] # only if using PostgreSQL -``` +Traces logged to Langfuse can be reviewed by domain experts. Positively scored traces automatically flow back into memory through the gatekeeper: -### Environment Setup - -```bash -# Interactive setup (guides you through each variable) -behavioral-memory setup +```python +from behavioral_memory import FeedbackPoller, AnnotationHandler -# Or manual -cp .env.example .env +poller = FeedbackPoller(settings=settings) +handler = AnnotationHandler(poller=poller, gatekeeper=gatekeeper) +handler.run_loop() # continuously polls → validates → stores ``` -| Variable | Required | Description | -|----------|----------|-------------| -| `GOOGLE_API_KEY` | For LLM calls | Gemini API key (or use any LangChain-compatible LLM) | -| `VECTOR_STORE_URL` | For PostgreSQL mode | `postgresql+psycopg://localhost/behavioral_memory` | -| `LANGFUSE_SECRET_KEY` | For observability | Langfuse secret key | -| `LANGFUSE_PUBLIC_KEY` | For observability | Langfuse public key | +### Without LangChain (plain Python) ---- +If you don't use LangChain, you can use the lower-level primitives directly: -## Architecture +```python +from behavioral_memory.planner.prompt import SYSTEM_PROMPT, build_prompt +from behavioral_memory.planner.postprocess import postprocess_plan -### Project Structure +# Build the prompt yourself +prompt = build_prompt(query="Get revenue data", traces=my_traces, tool_schemas=my_schemas) -``` -behavioral-memory/ -├── src/behavioral_memory/ # The pip-installable library -│ ├── core/ # Schemas, config, exceptions -│ ├── memory/ # Behavioral Layer (TraceStore, InMemoryTraceStore, dedup) -│ ├── tools/ # Tool Layer (MCP client, registry, mock tools) -│ ├── planner/ # Executive Layer (PlanEngine, prompt, postprocess) -│ ├── gatekeeper/ # Gatekeeper (schema validator, sandbox, dedup gate) -│ ├── observability/ # Langfuse (tracer, feedback poller, annotation) -│ └── evaluation/ # Benchmark (30 tasks, metrics, statistics) -├── agent/ # Reference LangGraph 1.x agent -│ ├── graph.py # StateGraph definition -│ ├── state.py # Agent state -│ └── nodes/ # Graph nodes (retrieve, plan, execute, observe) -├── tests/ # 104 tests (unit + integration + e2e) -│ ├── unit/ # 61 unit tests -│ ├── integration/ # 3 integration tests -│ └── e2e/ # 40 end-to-end tests -├── examples/ -│ ├── validate_pipeline.py # Full pipeline validation (no API keys) -│ ├── run_live_benchmark.py # Real benchmark (needs API key) -│ ├── gatekeeper_ablation.py # Gatekeeper ablation study (Section IV.D.5) -│ └── run_benchmark.py # Benchmark with PostgreSQL -├── Makefile # Common dev tasks (make lint, make test, etc.) -└── .github/workflows/ # CI/CD (lint, typecheck, test on 3.11/3.12/3.13) +# Call your own LLM +raw_output = your_llm.chat(system=SYSTEM_PROMPT, user=prompt) + +# Parse the JSON plan +steps = postprocess_plan(raw_output) # returns list[ToolCall] ``` -### Store Options +--- -| Store | When to Use | Requires | Paper Reproduction | -|-------|------------|----------|-------------------| -| `InMemoryTraceStore` | Development, demos, CI, quick benchmarks | Nothing (numpy only) | TSA/ESA match; PV/PCR lower | -| `TraceStore` (pgvector) | Production, paper reproduction, persistent memory | PostgreSQL + pgvector (`podman-compose up -d`) | Exact paper numbers | +## Persistence and Limitations -### The Framework is Model-Agnostic +| Store | Persistence | Multi-user | Best for | +|-------|------------|------------|----------| +| `InMemoryTraceStore` | Process memory only | No | Dev, CI, demos | +| `TraceStore` (pgvector) | PostgreSQL, survives restarts | Shared DB, single collection | Production | -| Provider | LLM | Embeddings | -|----------|-----|------------| -| Google | `ChatGoogleGenerativeAI` | `GoogleGenerativeAIEmbeddings` | -| OpenAI | `ChatOpenAI` | `OpenAIEmbeddings` | -| Anthropic | `ChatAnthropic` | (use OpenAI or Voyage) | -| Local | `ChatOllama` | `OllamaEmbeddings` | +**Current limitations:** +- All traces share one collection (default: `validated_traces`). No per-user or per-session isolation. +- Langfuse is **optional** — the core framework (planning, retrieval, gatekeeper) works without it. +- The reference agent at `agent/` is a planning demo with stub tool execution — bring your own tool runtime. --- -## How the Agent Learns (Feedback Loop) +## Key Results -The architecture implements a continuous learning cycle via Langfuse (Section III.F): +On a 30-task benchmark with 7 MCP tools (Gemini 2.5 Pro, temperature 0): -``` -User Query → Agent generates plan → Logged to Langfuse - ↓ - SME reviews in Langfuse dashboard - Assigns quality score (≥1.0 = positive) - ↓ - FeedbackPoller detects positive scores - ↓ - GatekeeperPipeline.submit(trace) - ├── Gate 1: Schema validation - ├── Gate 2: Sandboxed execution - └── Gate 3: Semantic deduplication - ↓ - If all gates pass → stored in memory - ↓ - Future queries retrieve this trace - → Agent produces better plans -``` +| Metric | Zero-Shot | Static Few-Shot | **With Behavioral Memory** | +|--------|-----------|----------------|---------------------------| +| Tool Selection (TSA) | 63.3% | 70.0% | **83.3%** | +| Parameter Validity (PV) | 72.2% | 79.6% | **84.0%** | +| Plan Correctness (PCR) | 33.3% | 50.0% | **63.3%** | +| Sequence Accuracy (ESA) | 63.3% | 70.0% | **83.3%** | -**Key insight:** The gatekeeper ensures only high-quality, non-duplicate, structurally valid traces enter memory. This is what separates our approach from systems like Reflexion that store unstructured reflections without validation. +McNemar's test: **p = 0.004** vs zero-shot. Plan correctness nearly doubled. -> **Note:** The paper's benchmark used a fixed memory of 12 seed traces to isolate the impact of retrieval. The feedback loop is implemented but was not exercised during evaluation (see Section V.C). Longitudinal testing with a growing memory is identified as the most important next step. +
+Reproduced live run (May 2026) -```python -from behavioral_memory import FeedbackPoller, GatekeeperPipeline, AnnotationHandler +| Metric | Paper | Live Run (pgvector) | +|--------|-------|---------------------| +| TSA | 83.3% | 86.7% | +| PV | 84.0% | 82.2% | +| PCR | 63.3% | 80.0% | +| ESA | 83.3% | 86.7% | +| McNemar p | 0.004 | 0.039 | -poller = FeedbackPoller(settings=settings) -gatekeeper = GatekeeperPipeline(store=store, registry=registry) -handler = AnnotationHandler(poller=poller, gatekeeper=gatekeeper) +All results within 95% bootstrap confidence intervals. -# Single pass: poll Langfuse → validate → store accepted traces -stats = handler.run_once() -print(f"Found {stats.traces_found}, accepted {stats.accepted}") - -# Continuous background loop -handler.run_loop() -``` +
--- -## Testing - -### Run all tests (104 tests, no external services needed) - -```bash -pip install -e ".[dev]" -pytest tests/ -v -``` - -### Test breakdown - -| Suite | Tests | What it covers | -|-------|-------|---------------| -| `tests/unit/` | 61 | Schemas, metrics, postprocessing, prompt assembly, token budget, in-memory store | -| `tests/integration/` | 3 | Schema validator + sandbox with real traces | -| `tests/e2e/` | 40 | Full pipeline: seed traces → prompt → mock LLM → metrics → gatekeeper | - -### Pipeline validation +## Architecture -```bash -python examples/validate_pipeline.py # 30 checks, 0 external deps -``` +Three layers (from the paper): -### Linting and type checking +| Layer | What it does | Key class | +|-------|-------------|-----------| +| **Behavioral** | Store and retrieve validated execution traces via cosine similarity | `InMemoryTraceStore` / `TraceStore` | +| **Tool** | Load tool schemas dynamically via MCP | `ToolRegistry` / `MCPClient` | +| **Executive** | Assemble prompt (traces + schemas + query), call LLM, parse plan | `PlanEngine` | -```bash -make lint # or: ruff check src/ tests/ agent/ examples/ server.py -make format # or: ruff format src/ tests/ agent/ examples/ server.py -make typecheck # or: mypy src/ -``` +**Gatekeeper Pipeline** guards memory quality with three gates: +1. **Schema validation** — tools exist, params valid, deps logical +2. **Sandboxed execution** — runtime check with timeout +3. **Semantic deduplication** — cosine > 0.95 rejected --- -## Gatekeeper Ablation Study (Section IV.D.5) - -Tests the critical role of the gatekeeper by injecting 8 deliberately poisoned traces -(wrong conventions, broken dependencies, incorrect tools) into memory: +## Reproduce the Paper ```bash -python examples/gatekeeper_ablation.py --verbose -``` +git clone https://github.com/harsh-kr11/behavioral-memory.git +cd behavioral-memory +pip install -e ".[agent,eval]" +export GOOGLE_API_KEY=your-key -Three conditions are compared: -1. **Baseline** — only valid seed traces (gatekeeper ON) -2. **Poisoned** — bad traces injected (gatekeeper OFF) -3. **Recovered** — gatekeeper re-enabled, bad traces filtered out +# Run the 30-task benchmark +python examples/run_live_benchmark.py -The script shows how poisoned traces degrade plan quality (PCR drops) and how the -gatekeeper catches and rejects them. Recovery restores baseline performance. +# Quick test (5 tasks) +python examples/run_live_benchmark.py --limit 5 ---- +# Exact paper reproduction (with pgvector) +pip install -e ".[postgres]" +docker compose up -d # or: podman-compose up -d +python examples/run_live_benchmark.py --postgres -## Development +# Gatekeeper ablation study (Section IV.D.5) +python examples/gatekeeper_ablation.py --verbose -```bash -# Using the Makefile (recommended) -make dev # Install all dev dependencies + pre-commit hooks -make lint # Run ruff linter -make format # Auto-format code -make typecheck # Run mypy -make test # Run all 104 tests -make ci # Run all CI checks locally -make benchmark # Run live benchmark with in-memory store -make benchmark-pg # Run live benchmark with pgvector (paper reproduction) -make ablation # Run gatekeeper ablation study -make validate # Pipeline validation (no API keys) -make demo # Offline demo +# Validate pipeline offline (no API keys) +python examples/validate_pipeline.py ``` ---- - -## Evaluation Metrics (Section IV.C) - -| Metric | Description | -|--------|-------------| -| **TSA** | Tool Selection Accuracy — correct tool multiset | -| **PV** | Parameter Validity — fraction of key params correct | -| **PCR** | Plan Correctness Rate — correct tools AND >= 80% PV | -| **ESA** | Execution Sequence Accuracy — correct tool ordering | +A reference LangGraph agent is included at `agent/` for demo purposes. --- -## CLI Tools +## Development ```bash -behavioral-memory setup # Interactive .env setup -behavioral-memory demo # Offline demo of behavioral memory -behavioral-memory benchmark info # Dataset summary -behavioral-memory benchmark ground-truth # View all 30 tasks -behavioral-memory benchmark seed-traces # View 12 seed traces -behavioral-memory benchmark tools # View 7 tool definitions +pip install -e ".[dev,eval]" +make test # 104 tests +make lint # ruff check +make typecheck # mypy (strict) +make ci # all checks ``` --- ## Configuration -All settings via environment variables or `.env`: +All via environment variables or `.env`: | Variable | Default | Description | |----------|---------|-------------| -| `VECTOR_STORE_URL` | `postgresql+psycopg://localhost/behavioral_memory` | pgvector connection | -| `VECTOR_STORE_COLLECTION` | `validated_traces` | Collection name | | `FEW_SHOT_K` | `3` | Traces to retrieve per query | | `MAX_PROMPT_TOKENS` | `3500` | Token budget for prompt | -| `SIMILARITY_DEDUP_THRESHOLD` | `0.95` | Cosine threshold for dedup | -| `SANDBOX_TIMEOUT_SECONDS` | `30` | Sandbox execution timeout | -| `FEEDBACK_SCORE_NAME` | `quality` | Langfuse score name | -| `FEEDBACK_POSITIVE_THRESHOLD` | `1.0` | Min score for positive | -| `FEEDBACK_POLL_INTERVAL` | `60` | Seconds between polls | - ---- - -## Tech Stack - -| Component | Technology | -|-----------|-----------| -| Vector Store | PostgreSQL + pgvector (production) / In-memory (development) | -| Embeddings | Any LangChain Embeddings (default: Gemini) | -| LLM | Any LangChain ChatModel (default: Gemini 2.5 Pro) | -| Agent Framework | LangGraph 1.x (reference agent) | -| Observability | Langfuse | -| Config | Pydantic Settings | -| Tokenization | tiktoken | -| CLI | Typer + Rich | -| Testing | pytest (104 tests) | -| Linting | ruff + pre-commit hooks | -| Type Checking | mypy (strict) | -| Package Management | uv | +| `SIMILARITY_DEDUP_THRESHOLD` | `0.95` | Dedup cosine threshold | +| `SANDBOX_TIMEOUT_SECONDS` | `30` | Gatekeeper sandbox timeout | +| `VECTOR_STORE_URL` | — | PostgreSQL connection (only for `TraceStore`) | +| `LANGFUSE_SECRET_KEY` | — | Langfuse secret (optional) | +| `LANGFUSE_PUBLIC_KEY` | — | Langfuse public key (optional) | --- ## Citation -If you use this software in your research, please cite our paper: - ```bibtex @inproceedings{khan2025behavioral, - title={Behavioral Memory for Tool Orchestration: Semantic Retrieval of Validated Execution Traces in MCP-Based Agent Systems}, + title={Behavioral Memory for Tool Orchestration: Semantic Retrieval of + Validated Execution Traces in MCP-Based Agent Systems}, author={Khan, Mehvash and Kumar, Harsh and Jangir, Rahul}, booktitle={IEEE Conference Proceedings}, year={2025} } ``` ---- - ## License -Apache License 2.0. See [LICENSE](LICENSE) for details. +Apache 2.0 — See [LICENSE](LICENSE). diff --git a/agent/tools/mcp_tools.py b/agent/tools/mcp_tools.py index 8740b0a..43e79d5 100644 --- a/agent/tools/mcp_tools.py +++ b/agent/tools/mcp_tools.py @@ -1,8 +1,9 @@ -"""MCP tool wrappers for LangGraph's ToolNode. +"""Stub tool executors — example code for wiring real tool execution. -These wrap the mock tools for use in the reference agent's -execution node. In production, these would dispatch to real -MCP servers. +The reference agent uses plan-only mode (no real tool execution). +These stubs show the pattern for dispatching to real MCP servers +or local tool implementations. They are not wired into the default +agent graph but can be used as a starting point for production agents. """ from __future__ import annotations diff --git a/pyproject.toml b/pyproject.toml index bc74579..285ff69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "langfuse>=2.0", "tiktoken>=0.7", "mcp>=1.0", + "typer>=0.12", "rich>=13.0", "numpy>=1.26", ] @@ -59,7 +60,6 @@ agent = [ "langgraph>=1.2", "langgraph-prebuilt>=1.1", "langchain-google-genai>=2.0", - "typer>=0.12", ] eval = [ "scipy>=1.12", diff --git a/src/behavioral_memory/evaluation/metrics.py b/src/behavioral_memory/evaluation/metrics.py index d458262..d8e639e 100644 --- a/src/behavioral_memory/evaluation/metrics.py +++ b/src/behavioral_memory/evaluation/metrics.py @@ -11,21 +11,26 @@ from collections import Counter from typing import Any -# Parameters that reflect orchestration decisions (the paper's focus). +# Orchestration parameters — the paper's primary focus (Section IV.C). # These control HOW tools connect and what structural choices are made. -_ORCHESTRATION_PARAMS = { - "source_step", - "format", - "channel", - "target", - "mode", - "operation", - "interval", - "notify_on_failure", - "attach_step", - "method", - "how", -} +# They use exact match (after normalization) in _param_matches, as does +# any param key not listed in _CONTENT_PARAMS or _IDENTIFIER_PARAMS. +# This set documents which keys are orchestration-relevant for clarity. +ORCHESTRATION_PARAMS = frozenset( + { + "source_step", + "format", + "channel", + "target", + "mode", + "operation", + "interval", + "notify_on_failure", + "attach_step", + "method", + "how", + } +) # Identifier params: orchestration-relevant but naming conventions vary. # Evaluated with lenient matching (key-term overlap). @@ -109,9 +114,12 @@ def compute_metrics(predicted_chain: list[dict[str, Any]], gold_chain: list[dict def _param_matches(predicted: object, gold: object, key: str = "") -> bool: """Compare a single parameter value. - Orchestration params use exact match (after normalization). - Content params (SQL, prose) use lenient semantic comparison - because the paper evaluates orchestration, not content authoring. + Routing logic: + - Content params (SQL, prose): lenient semantic comparison + - Identifier params (names, recipients): lenient key-term overlap + - Orchestration params + everything else: exact match after normalization + + The paper evaluates orchestration decisions, not content authoring. """ if predicted is None: return False @@ -122,6 +130,10 @@ def _param_matches(predicted: object, gold: object, key: str = "") -> bool: if key in _IDENTIFIER_PARAMS: return _content_param_matches(predicted, gold) + # Orchestration params and any unlisted params: exact match. + # _ORCHESTRATION_PARAMS documents which keys fall here (source_step, + # format, channel, target, mode, etc.) but all non-content, + # non-identifier params use exact match regardless. if isinstance(gold, str) and isinstance(predicted, str): return _normalize_str(predicted) == _normalize_str(gold) @@ -173,8 +185,9 @@ def _looks_like_sql(s: str) -> bool: def _sql_structural_match(pred: str, gold: str) -> bool: - """Check SQL structural equivalence: same tables and same aggregate functions. + """Check SQL structural equivalence via shared table names. + Extracts tables from FROM/JOIN clauses and checks for overlap. Aliases, column order, ORDER BY, and formatting are ignored — the orchestration question is "did it query the right data source?" """ diff --git a/src/behavioral_memory/memory/dedup.py b/src/behavioral_memory/memory/dedup.py index e904948..9468b4a 100644 --- a/src/behavioral_memory/memory/dedup.py +++ b/src/behavioral_memory/memory/dedup.py @@ -34,7 +34,8 @@ def is_duplicate(self, trace: ExecutionTrace) -> tuple[bool, float]: """Check if a trace is too similar to an existing one. Returns (is_duplicate, similarity_score). - Works with both TraceStore (PGVector) and InMemoryTraceStore. + Both TraceStore and InMemoryTraceStore return cosine similarity + (0-1, higher = more similar) from similarity_score(). """ score = self._store.similarity_score(trace.task_description) is_dup = score >= self.threshold diff --git a/src/behavioral_memory/memory/store.py b/src/behavioral_memory/memory/store.py index ff9869c..480a70c 100644 --- a/src/behavioral_memory/memory/store.py +++ b/src/behavioral_memory/memory/store.py @@ -129,12 +129,17 @@ def add_bulk(self, traces: list[ExecutionTrace]) -> int: raise MemoryStoreError(f"Bulk add failed: {e}") from e def similarity_score(self, query: str) -> float: - """Return the highest similarity score for a query against the store.""" + """Return the highest cosine similarity (0-1) for a query. + + PGVector returns cosine *distance* (0 = identical, 2 = opposite). + We convert to similarity (1 - distance) so the Deduplicator's + threshold (default 0.95) works consistently across both stores. + """ results = self.vectorstore.similarity_search_with_score(query, k=1) if not results: return 0.0 - _, score = results[0] - return float(score) + _, distance = results[0] + return max(0.0, 1.0 - float(distance)) def count(self) -> int: """Approximate count of traces in the store.""" diff --git a/tests/unit/test_metrics.py b/tests/unit/test_metrics.py index 4a7c0bd..e859d07 100644 --- a/tests/unit/test_metrics.py +++ b/tests/unit/test_metrics.py @@ -3,6 +3,13 @@ from __future__ import annotations from behavioral_memory.evaluation.metrics import ( + _content_param_matches, + _looks_like_sql, + _normalize_str, + _param_matches, + _sql_structural_match, + _structure_match, + _text_overlap_match, compute_metrics, execution_sequence_accuracy, parameter_validity, @@ -65,3 +72,132 @@ def test_full_computation(self): assert result["pcr"] is True assert result["esa"] is True assert result["pv"] == 1.0 + + +# ---------- Lenient PV matching (content / identifier / orchestration) ---------- + + +class TestParamMatchesRouting: + """Verify _param_matches routes to the right comparison by key.""" + + def test_content_param_uses_lenient_match(self): + assert _param_matches("SELECT * FROM orders", "SELECT * FROM orders WHERE id=1", key="query") + + def test_orchestration_param_uses_exact_match(self): + assert _param_matches("csv", "csv", key="format") + assert not _param_matches("json", "csv", key="format") + + def test_identifier_param_uses_lenient_match(self): + assert _param_matches("#data-alerts", "#data-alerts", key="recipient") + assert _param_matches("ops-team", "ops team alert", key="recipient") + + def test_none_predicted_always_false(self): + assert not _param_matches(None, "any", key="format") + assert not _param_matches(None, "SELECT 1", key="query") + + def test_unlisted_param_uses_exact_match(self): + assert _param_matches("30", "30", key="timeout") + assert not _param_matches("60", "30", key="timeout") + + +class TestSqlStructuralMatch: + """SQL matching should check table overlap, not exact string equality.""" + + def test_same_tables_different_aliases(self): + gold = "select sum(quantity * unit_price) as revenue from order_items" + pred = "select sum(oi.quantity * oi.unit_price) as rev from order_items oi" + assert _sql_structural_match(_normalize_str(pred), _normalize_str(gold)) + + def test_shared_table_with_extra_join(self): + gold = "select * from order_items" + pred = "select * from order_items join orders on order_items.order_id = orders.id" + assert _sql_structural_match(_normalize_str(pred), _normalize_str(gold)) + + def test_completely_different_tables_fails(self): + gold = "select * from order_items" + pred = "select * from customers" + assert not _sql_structural_match(_normalize_str(pred), _normalize_str(gold)) + + def test_empty_gold_tables_returns_true(self): + assert _sql_structural_match("select 1", "select count(*)") + + def test_gold_table_not_in_pred_fails(self): + gold = "select * from order_items join products on true" + pred = "select * from orders" + assert not _sql_structural_match(_normalize_str(pred), _normalize_str(gold)) + + +class TestTextOverlapMatch: + """Text matching should check domain-term overlap, ignoring stop words.""" + + def test_same_domain_terms(self): + assert _text_overlap_match("deployment completed successfully", "deployment update") + + def test_no_overlapping_terms(self): + assert not _text_overlap_match("hello world", "quarterly revenue") + + def test_stop_words_only_gold_is_vacuously_true(self): + # Gold with only stop words has no key terms → vacuously true + assert _text_overlap_match("hello world", "the is a an") + + def test_stop_words_filtered_from_matching(self): + # "quarterly" is a key term in gold; "monthly" doesn't match it + assert not _text_overlap_match("monthly the is", "quarterly revenue") + + def test_empty_gold_returns_true(self): + assert _text_overlap_match("anything", "the a is") + + +class TestStructureMatch: + """Dict/list structural matching.""" + + def test_dict_with_shared_keys(self): + assert _structure_match({"a": 1, "b": 2}, {"a": 10, "c": 3}) + + def test_dict_no_shared_keys(self): + assert not _structure_match({"x": 1}, {"y": 2}) + + def test_list_nonempty_if_gold_nonempty(self): + assert _structure_match([1], [10, 20]) + + def test_empty_pred_list_fails(self): + assert not _structure_match([], [1, 2]) + + def test_empty_gold_list_passes(self): + assert _structure_match([], []) + + def test_type_mismatch_fails(self): + assert not _structure_match({"a": 1}, [1, 2]) + + +class TestLooksLikeSql: + def test_sql_detected(self): + assert _looks_like_sql("SELECT * FROM customers WHERE id = 1") + + def test_non_sql_not_detected(self): + assert not _looks_like_sql("Deployment completed successfully") + + def test_needs_two_keywords(self): + assert not _looks_like_sql("select all items") + + +class TestContentParamMatches: + """End-to-end content param matching.""" + + def test_exact_match_short_circuits(self): + assert _content_param_matches("hello", "hello") + + def test_sql_uses_structural_match(self): + assert _content_param_matches( + "SELECT SUM(qty) FROM order_items GROUP BY product_id", + "SELECT SUM(quantity * unit_price) FROM order_items", + ) + + def test_prose_uses_text_overlap(self): + assert _content_param_matches("Revenue report for Q1", "Quarterly revenue summary") + + def test_none_returns_false(self): + assert not _content_param_matches(None, "anything") + + def test_dict_structural(self): + assert _content_param_matches({"new_column": "x"}, {"new_column": "y", "extra": "z"})