From ffdc9db87a52ac1f8f234c10d856337f3d3c870d Mon Sep 17 00:00:00 2001 From: harskuma Date: Mon, 18 May 2026 20:11:05 +0530 Subject: [PATCH 1/4] Add LangGraph server, Docker/Podman compose, and setup guide - server.py: LangGraph-compatible server entry point that works with Agent Chat UI and `langgraph dev`. Uses MessagesState for compatibility. - langgraph.json: Configuration for LangGraph CLI dev server. - docker-compose.yml: One-command PostgreSQL+pgvector setup (works with both Docker and Podman). - docs/GETTING_STARTED.md: Comprehensive setup guide covering all scenarios (InMemory, PostgreSQL, Langfuse, Agent Chat UI, BYOA). - Updated .env.example with all options and documentation. - Added server optional deps to pyproject.toml. - Updated .gitignore for benchmark results and Chainlit files. Co-authored-by: Cursor --- .env.example | 40 +++++-- .gitignore | 4 + docker-compose.yml | 20 ++++ docs/GETTING_STARTED.md | 248 ++++++++++++++++++++++++++++++++++++++++ langgraph.json | 16 +++ pyproject.toml | 5 + server.py | 109 ++++++++++++++++++ 7 files changed, 432 insertions(+), 10 deletions(-) create mode 100644 docker-compose.yml create mode 100644 docs/GETTING_STARTED.md create mode 100644 langgraph.json create mode 100644 server.py diff --git a/.env.example b/.env.example index 52f898a..9e7be03 100644 --- a/.env.example +++ b/.env.example @@ -1,22 +1,42 @@ -# === LLM & Embeddings === -# Only needed if using the reference agent with Gemini (default) +# ============================================= +# Behavioral Memory — Environment Configuration +# ============================================= +# Copy this file to .env and fill in your values: +# cp .env.example .env +# +# Or use the interactive setup: +# behavioral-memory setup + +# === LLM & Embeddings (required for agent/benchmark) === +# Get a key from: https://aistudio.google.com/apikey GOOGLE_API_KEY=your-google-api-key-here -# === PostgreSQL + pgvector === +# Model to use (optional, defaults in code) +# GEMINI_MODEL=gemini-2.5-flash + +# === PostgreSQL + pgvector (optional — InMemoryTraceStore works without it) === +# If using Docker/Podman compose: +# VECTOR_STORE_URL=postgresql+psycopg://behavioral_memory:behavioral_memory@localhost:5432/behavioral_memory +# If using your own PostgreSQL: VECTOR_STORE_URL=postgresql+psycopg://user:password@localhost:5432/behavioral_memory VECTOR_STORE_COLLECTION=validated_traces -# === Langfuse (optional but recommended) === +# Set to "true" to use PostgreSQL instead of in-memory store +USE_POSTGRES=false + +# === Langfuse — Observability & Feedback Loop === +# Sign up free: https://cloud.langfuse.com +# Then: Settings → API Keys → Create LANGFUSE_SECRET_KEY=sk-lf-... LANGFUSE_PUBLIC_KEY=pk-lf-... LANGFUSE_HOST=https://cloud.langfuse.com # === Framework Tuning === -FEW_SHOT_K=3 -MAX_PROMPT_TOKENS=3500 -SIMILARITY_DEDUP_THRESHOLD=0.95 +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) # === Feedback Loop === -FEEDBACK_SCORE_NAME=quality -FEEDBACK_POSITIVE_THRESHOLD=1.0 -FEEDBACK_POLL_INTERVAL=60 +FEEDBACK_SCORE_NAME=quality # Langfuse score name to watch +FEEDBACK_POSITIVE_THRESHOLD=1.0 # Minimum score for positive feedback +FEEDBACK_POLL_INTERVAL=60 # Seconds between polling Langfuse diff --git a/.gitignore b/.gitignore index 83a9d75..5b14b6f 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,7 @@ htmlcov/ Thumbs.db uv.lock + +benchmark_results.json +.chainlit/ +chainlit.md diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..a00f48a --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,20 @@ +services: + pgvector: + image: pgvector/pgvector:pg16 + container_name: behavioral-memory-pgvector + environment: + POSTGRES_USER: behavioral_memory + POSTGRES_PASSWORD: behavioral_memory + POSTGRES_DB: behavioral_memory + ports: + - "5432:5432" + volumes: + - pgdata:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U behavioral_memory"] + interval: 5s + timeout: 5s + retries: 5 + +volumes: + pgdata: diff --git a/docs/GETTING_STARTED.md b/docs/GETTING_STARTED.md new file mode 100644 index 0000000..934da4a --- /dev/null +++ b/docs/GETTING_STARTED.md @@ -0,0 +1,248 @@ +# Getting Started — Full Setup Guide + +This guide walks you through setting up the complete behavioral-memory system: +the agent, Langfuse observability, PostgreSQL vector store, and the live UI. + +--- + +## 1. Prerequisites + +- Python 3.11+ +- A Google API key (for Gemini LLM + Embeddings): https://aistudio.google.com/apikey +- (Optional) Podman or Docker for PostgreSQL +- (Optional) Langfuse account for observability: https://cloud.langfuse.com + +## 2. Install + +```bash +git clone https://github.com/harsh-kr11/behavioral-memory.git +cd behavioral-memory + +# Create virtual environment +python3.11 -m venv .venv +source .venv/bin/activate # Linux/macOS +# .venv\Scripts\activate # Windows + +# Install everything +pip install -e ".[all,server]" +``` + +## 3. Configure Environment + +```bash +# Interactive setup (recommended) +behavioral-memory setup + +# Or manual +cp .env.example .env +# Edit .env with your API key +``` + +At minimum you need: +``` +GOOGLE_API_KEY=your-key-from-aistudio +``` + +## 4. Validate Everything Works + +```bash +# Run 30 pipeline checks — no API key needed +python examples/validate_pipeline.py + +# Quick test with real LLM (3 tasks) +python examples/run_live_benchmark.py --limit 3 --model gemini-2.0-flash + +# Full benchmark (30 tasks, takes ~10 minutes) +python examples/run_live_benchmark.py --model gemini-2.5-flash +``` + +## 5. Run the Interactive Agent + +```bash +# Interactive REPL +python -m agent.app --interactive + +# Single query +python -m agent.app "Build a revenue analysis pipeline" +``` + +In interactive mode: +- Type any query to get a plan +- `/compare ` — shows with vs without behavioral memory +- `/memory` — inspect what traces are in memory +- `/quit` — exit + +## 6. Start the LangGraph Server + Agent Chat UI + +```bash +# Install LangGraph CLI +pip install "langgraph-cli[inmem]" + +# Start the development server (port 2024) +langgraph dev + +# Open the Agent Chat UI in your browser: +# https://agentchat.vercel.app +# Connect to: http://localhost:2024 +``` + +The Agent Chat UI shows: +- Tool calls and reasoning steps +- Time-travel debugging +- State visualization + +## 7. Set Up Langfuse (Observability + Feedback Loop) + +### Create Account +1. Go to https://cloud.langfuse.com +2. Sign up (free tier available) +3. Create a new project (e.g., "behavioral-memory") +4. Go to **Settings → API Keys → Create new API key** +5. Copy the secret key and public key + +### Configure +Add to your `.env`: +``` +LANGFUSE_SECRET_KEY=sk-lf-xxxx +LANGFUSE_PUBLIC_KEY=pk-lf-xxxx +LANGFUSE_HOST=https://cloud.langfuse.com +``` + +### What You'll See +Every plan generated by the agent will appear in Langfuse with: +- The input query +- The generated plan (tool calls + parameters) +- How many traces were retrieved from memory +- Which tool schemas were used +- Token budget usage + +### Feedback Loop +1. Run the agent — plans get logged to Langfuse +2. In Langfuse, review a trace and add a **score** (name: `quality`, value: 1.0 for good) +3. The `FeedbackPoller` picks up positively scored traces +4. Gatekeeper validates them (schema check + sandbox + dedup) +5. Validated traces enter behavioral memory +6. Future queries benefit from these new reference examples + +```python +from behavioral_memory import FeedbackPoller, GatekeeperPipeline + +poller = FeedbackPoller(settings=settings) +gatekeeper = GatekeeperPipeline(store=store, registry=registry) + +# Auto-learn in the background +poller.poll_loop(callback=lambda trace: gatekeeper.submit(trace)) +``` + +## 8. Set Up PostgreSQL + pgvector (Optional) + +The framework works fine with `InMemoryTraceStore` for demos and development. +For persistent memory across restarts, use PostgreSQL + pgvector. + +### With Podman + +```bash +# Start PostgreSQL with pgvector +podman-compose up -d + +# Or without compose: +podman run -d \ + --name behavioral-memory-pgvector \ + -e POSTGRES_USER=behavioral_memory \ + -e POSTGRES_PASSWORD=behavioral_memory \ + -e POSTGRES_DB=behavioral_memory \ + -p 5432:5432 \ + pgvector/pgvector:pg16 +``` + +### With Docker + +```bash +docker compose up -d +``` + +### Configure +Add to your `.env`: +``` +USE_POSTGRES=true +VECTOR_STORE_URL=postgresql+psycopg://behavioral_memory:behavioral_memory@localhost:5432/behavioral_memory +``` + +### Verify +```bash +# Connect to PostgreSQL +podman exec -it behavioral-memory-pgvector psql -U behavioral_memory -c "CREATE EXTENSION IF NOT EXISTS vector;" +``` + +## 9. Run Tests + +```bash +# All 104 tests +pytest tests/ -v + +# Just unit tests +pytest tests/unit/ -v + +# With coverage +pytest tests/ --cov=behavioral_memory --cov-report=html +``` + +## 10. Bring Your Own Agent (Integration Guide) + +The framework is designed to be plugged into any LangChain/LangGraph agent: + +```python +from behavioral_memory import InMemoryTraceStore, PlanEngine, ToolRegistry +from langchain_openai import ChatOpenAI, OpenAIEmbeddings + +# Use any LLM provider +llm = ChatOpenAI(model="gpt-4o", temperature=0) +embeddings = OpenAIEmbeddings() + +# Set up the memory store +store = InMemoryTraceStore(embeddings=embeddings) + +# Register your tools +registry = ToolRegistry() +registry.register(ToolSchema( + name="my_tool", + description="Does something useful", + parameters_schema={"type": "object", "properties": {...}}, + required_params=["param1"], +)) + +# Create the plan engine +engine = PlanEngine(llm=llm, store=store, registry=registry) + +# Generate a plan +plan = engine.generate(query="Do the thing") +print(plan.steps) # List of ToolCall objects + +# The gatekeeper ensures quality traces enter memory +from behavioral_memory import GatekeeperPipeline +gatekeeper = GatekeeperPipeline(store=store, registry=registry) +gatekeeper.submit(trace) # Only accepted if valid + unique +``` + +### Key integration points: + +| Component | What it does | How to plug in | +|-----------|-------------|---------------| +| `PlanEngine` | Generates tool orchestration plans | Pass your LLM + store + registry | +| `InMemoryTraceStore` | Stores/retrieves execution traces | Use as-is or swap for `TraceStore` (pgvector) | +| `ToolRegistry` | Manages tool schemas | Register your MCP tools | +| `GatekeeperPipeline` | Validates traces before storage | Call `submit(trace)` after successful executions | +| `LangfuseTracer` | Logs plans to Langfuse | Pass settings with Langfuse keys | +| `FeedbackPoller` | Learns from human feedback | Run in background to auto-ingest | + +--- + +## Troubleshooting + +| Issue | Solution | +|-------|---------| +| `API key expired` | Get a new key from https://aistudio.google.com/apikey | +| `ModuleNotFoundError` | Run `pip install -e ".[all,server]"` | +| `Connection refused (PostgreSQL)` | Check if Podman/Docker container is running | +| `Langfuse traces not appearing` | Verify keys in `.env`, check `tracer.enabled` | +| `Low PV scores` | Normal — PV measures exact parameter match. The model often generates equivalent but differently formatted params | diff --git a/langgraph.json b/langgraph.json new file mode 100644 index 0000000..3d3c538 --- /dev/null +++ b/langgraph.json @@ -0,0 +1,16 @@ +{ + "$schema": "https://langgra.ph/schema.json", + "python_version": "3.11", + "dependencies": [ + ".", + "langchain-google-genai", + "langgraph" + ], + "graphs": { + "behavioral_memory_agent": { + "path": "./server.py:graph", + "description": "Behavioral Memory Agent — generates tool orchestration plans using validated execution traces from memory" + } + }, + "env": ".env" +} diff --git a/pyproject.toml b/pyproject.toml index bcbd937..4aeb0a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,11 @@ eval = [ "numpy>=1.26", "scipy>=1.12", ] +server = [ + "langgraph-cli[inmem]>=0.1", + "langchain-google-genai>=2.0", + "langgraph>=1.2", +] dev = [ "pytest>=8.0", "pytest-asyncio>=0.23", diff --git a/server.py b/server.py new file mode 100644 index 0000000..5d38cf4 --- /dev/null +++ b/server.py @@ -0,0 +1,109 @@ +"""LangGraph server entry point for the behavioral memory agent. + +Exposes the agent as a LangGraph-compatible graph that works with: + - `langgraph dev` (development server with hot reload) + - Agent Chat UI (https://github.com/langchain-ai/agent-chat-ui) + - Any LangGraph Platform-compatible client + +Usage: + pip install "langgraph-cli[inmem]" + langgraph dev + + # Or directly: + python server.py +""" + +from __future__ import annotations + +import os + +from langchain_core.messages import AIMessage, HumanMessage +from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import MessagesState + +from behavioral_memory.core.config import Settings +from behavioral_memory.evaluation.seed_traces import get_seed_traces +from behavioral_memory.memory.in_memory_store import InMemoryTraceStore +from behavioral_memory.observability.tracer import LangfuseTracer +from behavioral_memory.planner.engine import PlanEngine +from behavioral_memory.tools.mock_tools import get_tool_schemas +from behavioral_memory.tools.registry import ToolRegistry + +settings = Settings() + +model_name = os.getenv("GEMINI_MODEL", "gemini-2.5-flash") +llm = ChatGoogleGenerativeAI(model=model_name, temperature=0) +embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001") + +use_postgres = os.getenv("USE_POSTGRES", "false").lower() == "true" +if use_postgres: + from behavioral_memory.memory.store import TraceStore + store = TraceStore(embeddings=embeddings, settings=settings) +else: + store = InMemoryTraceStore(embeddings=embeddings, settings=settings) + +registry = ToolRegistry() +registry.register_many(get_tool_schemas()) +store.add_bulk(get_seed_traces()) + +engine = PlanEngine(llm=llm, store=store, registry=registry, settings=settings) +tracer = LangfuseTracer(settings=settings) + +schemas = get_tool_schemas() + + +def plan_agent(state: MessagesState) -> dict: + """Core agent node: takes the last user message, generates a plan.""" + messages = state["messages"] + last_msg = messages[-1] + query = last_msg.content if hasattr(last_msg, "content") else str(last_msg) + + try: + plan = engine.generate(query=query, tool_schemas=schemas) + + steps_text = [] + for step in plan.steps: + params_str = ", ".join(f"{k}={v!r}" for k, v in step.parameters.items()) + deps = f" (depends on: {', '.join(step.depends_on)})" if step.depends_on else "" + steps_text.append(f"**{step.step_id}**: `{step.tool_name}({params_str})`{deps}") + + retrieved_info = "" + if plan.retrieved_traces: + examples = [f" - {t.task_description[:80]}" for t in plan.retrieved_traces] + retrieved_info = ( + f"\n\n**Retrieved {len(plan.retrieved_traces)} reference traces from behavioral memory:**\n" + + "\n".join(examples) + ) + + response = ( + f"## Execution Plan ({len(plan.steps)} steps)\n\n" + + "\n".join(steps_text) + + retrieved_info + + f"\n\n*Token budget used: {plan.token_budget_used} tokens*" + ) + + if tracer.enabled: + trace_id = tracer.log_plan(plan, tags=["agent-chat"]) + if trace_id: + response += f"\n\n*Logged to Langfuse: `{trace_id}`*" + tracer.flush() + + return {"messages": [AIMessage(content=response)]} + + except Exception as e: + return {"messages": [AIMessage(content=f"Planning failed: {e}")]} + + +graph_builder = StateGraph(MessagesState) +graph_builder.add_node("plan", plan_agent) +graph_builder.add_edge(START, "plan") +graph_builder.add_edge("plan", END) + +graph = graph_builder.compile() + +if __name__ == "__main__": + print("Testing the graph locally...") + result = graph.invoke({"messages": [HumanMessage(content="Build a revenue analysis pipeline")]}) + for msg in result["messages"]: + print(f"\n{msg.content}") From 8c24f35cf840c243d8888229dc1bc9a72a75e127 Mon Sep 17 00:00:00 2001 From: harskuma Date: Tue, 19 May 2026 12:02:54 +0530 Subject: [PATCH 2/4] Fix CI: make PostgreSQL deps optional, resolve all lint/type errors, add gatekeeper ablation - Move sqlalchemy/psycopg/pgvector/langchain-postgres to [postgres] optional extra - Guard TraceStore import with lazy __getattr__ so 'import behavioral_memory' works without PostgreSQL deps installed - Fix all 22 mypy errors (dict type args, scipy/typer stubs, Any returns) - Run ruff format on all 69 Python files across src/, tests/, agent/, examples/ - Update CI workflow to lint/test all directories, not just src/tests - Add gatekeeper ablation script (Section IV.D.5) with 8 poisoned traces - Add Makefile with common dev tasks (make lint, test, ci, ablation, etc.) - Add PEP 561 py.typed marker for downstream type checking - Add PR template for consistent contributions - Update CHANGELOG with real date and all features - Update README with ablation study docs, Makefile usage, postgres extra Co-authored-by: Cursor --- .github/PULL_REQUEST_TEMPLATE.md | 20 + .github/workflows/ci.yml | 6 +- CHANGELOG.md | 18 +- Makefile | 47 ++ README.md | 62 +- agent/app.py | 34 +- agent/nodes/execute.py | 14 +- examples/gatekeeper_ablation.py | 526 ++++++++++++++ examples/run_benchmark.py | 6 +- examples/run_live_benchmark.py | 16 +- examples/validate_pipeline.py | 64 +- pyproject.toml | 15 +- server.py | 1 + src/behavioral_memory/__init__.py | 18 +- src/behavioral_memory/cli.py | 103 +-- src/behavioral_memory/core/exceptions.py | 6 +- src/behavioral_memory/core/schemas.py | 20 +- src/behavioral_memory/evaluation/benchmark.py | 42 +- .../evaluation/ground_truth.py | 682 ++++++++++++++++-- src/behavioral_memory/evaluation/metrics.py | 25 +- .../evaluation/seed_traces.py | 259 ++++++- .../evaluation/statistics.py | 4 +- .../evaluation/strategies.py | 8 +- src/behavioral_memory/gatekeeper/pipeline.py | 4 +- src/behavioral_memory/gatekeeper/sandbox.py | 3 +- .../gatekeeper/schema_validator.py | 15 +- src/behavioral_memory/memory/__init__.py | 10 +- .../memory/in_memory_store.py | 9 +- src/behavioral_memory/memory/store.py | 44 +- .../observability/annotation.py | 21 +- .../observability/feedback.py | 10 +- src/behavioral_memory/observability/tracer.py | 4 +- src/behavioral_memory/planner/engine.py | 4 +- src/behavioral_memory/planner/postprocess.py | 5 +- src/behavioral_memory/tools/mcp_client.py | 8 +- src/behavioral_memory/tools/mock_tools.py | 4 +- tests/e2e/test_full_pipeline.py | 154 ++-- tests/integration/test_gatekeeper.py | 46 +- tests/unit/test_schema_validator.py | 6 +- 39 files changed, 1908 insertions(+), 435 deletions(-) create mode 100644 .github/PULL_REQUEST_TEMPLATE.md create mode 100644 Makefile create mode 100644 examples/gatekeeper_ablation.py diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..7c8be28 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,20 @@ +## Summary + + + +## Changes + + + + + +## Test Plan + +- [ ] `make lint` passes +- [ ] `make typecheck` passes +- [ ] `make test` passes +- [ ] Tested locally with `make validate` + +## Related Issues + + diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de80853..ce394f4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,8 +13,8 @@ jobs: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v4 - run: uv sync --extra dev - - run: uv run ruff check src/ tests/ - - run: uv run ruff format --check src/ tests/ + - run: uv run ruff check src/ tests/ agent/ examples/ server.py + - run: uv run ruff format --check src/ tests/ agent/ examples/ server.py typecheck: runs-on: ubuntu-latest @@ -35,4 +35,4 @@ jobs: with: python-version: ${{ matrix.python-version }} - run: uv sync --extra dev --extra eval - - run: uv run pytest tests/unit/ -v + - run: uv run pytest tests/ -v --tb=short diff --git a/CHANGELOG.md b/CHANGELOG.md index fbdc130..72dbb43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,18 +2,22 @@ All notable changes to this project will be documented in this file. -## [0.1.0] - 2025-XX-XX +## [0.1.0] - 2026-05-17 ### Added - Initial release of behavioral-memory framework - Three-layer architecture: Behavioral Layer, Tool Layer, Executive Layer -- TraceStore with pgvector for semantic retrieval of execution traces -- GatekeeperPipeline: schema validation, sandboxed execution, semantic dedup +- TraceStore (PostgreSQL + pgvector) and InMemoryTraceStore (zero-infra) for semantic retrieval +- GatekeeperPipeline: schema validation, sandboxed execution, semantic dedup (cosine 0.95 threshold) - Langfuse integration: trace logging, feedback polling, annotation handler -- PlanEngine: model-agnostic plan generation with any LangChain chat model -- Reference LangGraph 1.x agent with full paper implementation -- 30-task benchmark with 7 tools and 12 seed traces +- PlanEngine: model-agnostic plan generation with any LangChain-compatible chat model +- Reference LangGraph 1.x agent with interactive mode and LangGraph server support +- 30-task benchmark with 7 MCP tools and 12 seed traces - Evaluation metrics: TSA, PV, PCR, ESA with bootstrap CI and McNemar's test -- CLI for benchmarks, memory management, and dataset inspection +- Gatekeeper ablation study script (Section IV.D.5 — memory poisoning experiment) +- CLI for benchmarks, demos, memory management, and dataset inspection +- Full CI/CD with GitHub Actions (lint, typecheck, test on Python 3.11/3.12/3.13) +- Makefile for common development tasks +- PEP 561 `py.typed` marker for downstream type checking - Apache 2.0 license diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..8bcb937 --- /dev/null +++ b/Makefile @@ -0,0 +1,47 @@ +.PHONY: help install dev lint format typecheck test test-unit test-e2e demo benchmark 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}' + +install: ## Install core package + uv sync + +dev: ## Install with all dev dependencies + uv sync --extra dev --extra eval --extra agent --extra server + uv run pre-commit install + +lint: ## Run linter + uv run ruff check src/ tests/ agent/ examples/ server.py + +format: ## Format code + uv run ruff format src/ tests/ agent/ examples/ server.py + +typecheck: ## Run mypy type checker + uv run mypy src/ + +test: ## Run all tests + uv run pytest tests/ -v --tb=short + +test-unit: ## Run unit tests only + uv run pytest tests/unit/ -v + +test-e2e: ## Run end-to-end tests only + uv run pytest tests/e2e/ -v + +demo: ## Run offline demo (no API keys needed) + uv run behavioral-memory demo + +benchmark: ## Run live benchmark (requires GOOGLE_API_KEY) + uv run python examples/run_live_benchmark.py + +ablation: ## Run gatekeeper ablation study + uv run python examples/gatekeeper_ablation.py --verbose + +validate: ## Validate full pipeline (no API keys needed) + uv run python examples/validate_pipeline.py + +clean: ## Clean build artifacts + rm -rf dist/ build/ *.egg-info .mypy_cache .pytest_cache .ruff_cache + find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true + +ci: lint format typecheck test ## Run all CI checks locally diff --git a/README.md b/README.md index f3f7eb7..7857b70 100644 --- a/README.md +++ b/README.md @@ -243,10 +243,11 @@ All **30 checks** pass with zero external dependencies. ### Install with uv (recommended) ```bash -uv add behavioral-memory # core framework -uv add behavioral-memory[agent] # + reference LangGraph agent -uv add behavioral-memory[eval] # + evaluation/statistics -uv add behavioral-memory[all] # everything +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 ``` ### Install with pip @@ -254,6 +255,7 @@ uv add behavioral-memory[all] # everything ```bash pip install behavioral-memory pip install behavioral-memory[agent,eval] +pip install behavioral-memory[postgres] # only if using PostgreSQL ``` ### Environment Setup @@ -298,10 +300,12 @@ behavioral-memory/ │ ├── 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) -│ └── run_benchmark.py # Benchmark with PostgreSQL -└── .github/workflows/ # CI/CD +│ ├── 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) ``` ### Store Options @@ -370,9 +374,45 @@ python examples/validate_pipeline.py # 30 checks, 0 external deps ### Linting and type checking ```bash -ruff check src/ tests/ agent/ -ruff format src/ tests/ agent/ -mypy src/ +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 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: + +```bash +python examples/gatekeeper_ablation.py --verbose +``` + +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 + +The script shows how poisoned traces degrade plan quality (PCR drops) and how the +gatekeeper catches and rejects them. Recovery restores baseline performance. + +--- + +## Development + +```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 ablation # Run gatekeeper ablation study +make validate # Pipeline validation (no API keys) +make demo # Offline demo ``` --- diff --git a/agent/app.py b/agent/app.py index 0917cd2..8c4eb08 100644 --- a/agent/app.py +++ b/agent/app.py @@ -37,9 +37,11 @@ def create_agent(model: str = "gemini-2.5-pro", use_postgres: bool = False): if use_postgres: from behavioral_memory.memory.store import TraceStore + store = TraceStore(embeddings=embeddings, settings=settings) else: from behavioral_memory.memory.in_memory_store import InMemoryTraceStore + store = InMemoryTraceStore(embeddings=embeddings, settings=settings) registry = ToolRegistry() @@ -50,7 +52,10 @@ def create_agent(model: str = "gemini-2.5-pro", use_postgres: bool = False): store.add_bulk(seed_traces) graph = build_agent_graph( - llm=llm, store=store, registry=registry, settings=settings, + llm=llm, + store=store, + registry=registry, + settings=settings, ) return graph, store, registry, tracer, settings @@ -102,17 +107,19 @@ def run_single(query: str, model: str = "gemini-2.5-pro", verbose: bool = True) def run_interactive(model: str = "gemini-2.5-pro") -> None: """Interactive REPL — type queries, see plans, compare with/without memory.""" - console.print(Panel.fit( - "[bold]Behavioral Memory Agent — Interactive Mode[/bold]\n\n" - f"Model: {model}\n" - "Type a query to generate a plan. The agent retrieves relevant\n" - "traces from behavioral memory to guide its planning.\n\n" - "Special commands:\n" - " /compare — run with AND without memory, show difference\n" - " /memory — show what's in behavioral memory\n" - " /quit — exit", - title="Interactive Agent", - )) + console.print( + Panel.fit( + "[bold]Behavioral Memory Agent — Interactive Mode[/bold]\n\n" + f"Model: {model}\n" + "Type a query to generate a plan. The agent retrieves relevant\n" + "traces from behavioral memory to guide its planning.\n\n" + "Special commands:\n" + " /compare — run with AND without memory, show difference\n" + " /memory — show what's in behavioral memory\n" + " /quit — exit", + title="Interactive Agent", + ) + ) graph, store, registry, tracer, settings = create_agent(model=model) compiled = graph.compile() @@ -132,6 +139,7 @@ def run_interactive(model: str = "gemini-2.5-pro") -> None: if query.startswith("/memory"): from behavioral_memory.evaluation.seed_traces import get_seed_traces + for trace in get_seed_traces(): tools = " → ".join(trace.tool_names) console.print(f" [cyan]{trace.task_description[:60]}[/cyan]") @@ -167,7 +175,7 @@ def _run_comparison(compiled, query, store, registry, settings, tracer): from behavioral_memory.planner.prompt import build_prompt from behavioral_memory.tools.mock_tools import get_tool_schemas - console.print(f"\n[bold]Comparing: \"{query}\"[/bold]\n") + console.print(f'\n[bold]Comparing: "{query}"[/bold]\n') result_with = compiled.invoke({"query": query}) plan_with = result_with.get("plan") diff --git a/agent/nodes/execute.py b/agent/nodes/execute.py index 39ef108..5de81fc 100644 --- a/agent/nodes/execute.py +++ b/agent/nodes/execute.py @@ -19,11 +19,13 @@ def execute_tools(state: AgentState) -> dict: results: list[dict[str, Any]] = [] for step in plan.steps: - results.append({ - "step_id": step.step_id, - "tool_name": step.tool_name, - "status": "executed_stub", - "parameters": step.parameters, - }) + results.append( + { + "step_id": step.step_id, + "tool_name": step.tool_name, + "status": "executed_stub", + "parameters": step.parameters, + } + ) return {"execution_results": results} diff --git a/examples/gatekeeper_ablation.py b/examples/gatekeeper_ablation.py new file mode 100644 index 0000000..f55ae47 --- /dev/null +++ b/examples/gatekeeper_ablation.py @@ -0,0 +1,526 @@ +"""Gatekeeper Ablation Study — Section IV.D.5 of the paper. + +Demonstrates the critical role of the gatekeeper pipeline by measuring +what happens when poisoned (invalid) traces bypass quality control and +enter behavioral memory. + +The experiment: + 1. BASELINE: Run dynamic retrieval with only valid seed traces (gatekeeper ON) + 2. POISONED: Inject deliberately bad traces (wrong conventions, broken deps, + incorrect tools) and re-run retrieval (gatekeeper OFF) + 3. RECOVERED: Apply gatekeeper to the poisoned store, remove bad traces, + and re-run retrieval (gatekeeper restored) + +Expected outcome (from the paper): + - Poisoned memory degrades PCR by 15-25% as the LLM copies bad patterns + - Gatekeeper catches and rejects all poisoned traces when enabled + +Usage: + python examples/gatekeeper_ablation.py + python examples/gatekeeper_ablation.py --verbose + python examples/gatekeeper_ablation.py --poisoned-ratio 0.5 + +Requires: GOOGLE_API_KEY (or any LangChain-compatible LLM) +""" + +from __future__ import annotations + +import argparse +import logging +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, ToolCall +from behavioral_memory.evaluation.ground_truth import EVALUATION_TASKS +from behavioral_memory.evaluation.metrics import compute_metrics +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.tools.mock_tools import get_tool_schemas +from behavioral_memory.tools.registry import ToolRegistry + +console = Console() +logger = logging.getLogger(__name__) + +POISONED_TRACES: list[dict[str, Any]] = [ + { + "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", + }, + { + "task": "Archive valid orders as CSV", + "chain": [ + { + "step_id": "s1", + "tool": "query_database", + "params": {"query": "SELECT * FROM orders WHERE status != 'cancelled';"}, + }, + { + "step_id": "s2", + "tool": "store_results", + "params": {"source_step": "s1", "target": "csv_file", "target_name": "orders.csv"}, + }, + ], + "poison_type": "wrong_convention", + "explanation": "Only excludes 'cancelled' but not 'returned'; uses overwrite instead of append mode", + }, + { + "task": "Alert ops team about fulfillment metrics", + "chain": [ + { + "step_id": "s1", + "tool": "query_database", + "params": {"query": "SELECT COUNT(*) as delivered FROM orders WHERE status = 'delivered';"}, + }, + { + "step_id": "s2", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "ops@company.com", + "subject": "Fulfillment", + "body": "Metrics attached.", + "attach_step": "s1", + }, + }, + ], + "poison_type": "wrong_convention", + "explanation": "Wrong fulfillment formula (just delivered count, not rate); sends via email instead of slack #data-alerts", + }, + { + "task": "Build net order value pipeline", + "chain": [ + { + "step_id": "s1", + "tool": "query_database", + "params": {"query": "SELECT order_id, total_amount FROM orders;"}, + }, + { + "step_id": "s2", + "tool": "store_results", + "params": {"source_step": "s1", "target": "cache", "target_name": "net_values"}, + }, + ], + "poison_type": "wrong_convention", + "explanation": "Doesn't subtract discount for net value; skips transform step required for pipelines", + }, + { + "task": "Send basket size analysis to team", + "chain": [ + { + "step_id": "s1", + "tool": "query_database", + "params": { + "query": "SELECT order_id, SUM(unit_price * quantity) AS basket_value FROM order_items GROUP BY order_id;" + }, + }, + { + "step_id": "s2", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#general", + "subject": "Basket Analysis", + "body": "Data ready.", + "attach_step": "s1", + }, + }, + ], + "poison_type": "wrong_convention", + "explanation": "Uses dollar value instead of item count for basket size; sends to #general instead of #data-alerts", + }, + { + "task": "Schedule daily customer report", + "chain": [ + { + "step_id": "s1", + "tool": "schedule_task", + "params": { + "task_name": "customer_report", + "workflow_steps": [{"tool": "query_database"}], + "interval": "weekly", + }, + }, + ], + "poison_type": "wrong_convention", + "explanation": "Uses weekly instead of daily; missing generate_report step; missing notify_on_failure=true", + }, + { + "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", + }, +] + + +def build_poisoned_traces() -> list[ExecutionTrace]: + """Convert raw poisoned trace definitions into ExecutionTrace objects.""" + traces = [] + for raw in POISONED_TRACES: + tool_chain = [ToolCall(step_id=s["step_id"], tool_name=s["tool"], parameters=s["params"]) for s in raw["chain"]] + trace = ExecutionTrace( + task_description=raw["task"], + tool_chain=tool_chain, + validated=False, + source="execution", + metadata={"poison_type": raw["poison_type"], "explanation": raw["explanation"]}, + ) + traces.append(trace) + return traces + + +class MockEmbeddings: + """Simple bag-of-words embeddings for the ablation study.""" + + def embed_query(self, text: str) -> list[float]: + return self._embed(text) + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + return [self._embed(t) for t in texts] + + @staticmethod + def _embed(text: str) -> list[float]: + words = text.lower().split() + vocab = [ + "revenue", + "order", + "customer", + "product", + "report", + "query", + "database", + "send", + "notification", + "alert", + "dashboard", + "cache", + "csv", + "archive", + "schedule", + "daily", + "weekly", + "completed", + "valid", + "cancelled", + "returned", + "shipped", + "delivered", + "fulfillment", + "basket", + "net", + "value", + "pipeline", + "transform", + "total", + "sum", + "count", + "items", + "quantity", + "price", + ] + vec = [0.0] * len(vocab) + for i, v_word in enumerate(vocab): + vec[i] = float(words.count(v_word)) / max(len(words), 1) + norm = sum(x * x for x in vec) ** 0.5 + if norm > 0: + vec = [x / norm for x in vec] + return vec + + +def run_evaluation_with_store( + store: InMemoryTraceStore, + schemas: list[Any], + tasks: list[dict[str, Any]], + label: str, +) -> dict[str, Any]: + """Simulate dynamic retrieval against a given memory store.""" + tsa_hits = 0 + pcr_hits = 0 + esa_hits = 0 + pv_total = 0.0 + n = len(tasks) + + per_task: list[dict[str, Any]] = [] + + for task in tasks: + gold = task["gold_tool_chain"] + + retrieved = store.search(task["task"], k=3) + if retrieved: + best_trace = retrieved[0][0] + predicted = [{"tool": s.tool_name, "params": s.parameters} for s in best_trace.tool_chain] + else: + predicted = [{"tool": "query_database", "params": {"query": "SELECT 1"}}] + + metrics = compute_metrics(predicted, gold) + tsa_hits += int(bool(metrics["tsa"])) + pv_total += float(metrics["pv"]) + pcr_hits += int(bool(metrics["pcr"])) + esa_hits += int(bool(metrics["esa"])) + + per_task.append( + { + "task_id": task["task_id"], + "task": task["task"], + "difficulty": task["difficulty"], + "n_retrieved": len(retrieved), + "metrics": metrics, + "retrieved_source": retrieved[0][0].source if retrieved else "none", + } + ) + + return { + "label": label, + "n_tasks": n, + "tsa": tsa_hits / n if n > 0 else 0.0, + "pv": pv_total / n if n > 0 else 0.0, + "pcr": pcr_hits / n if n > 0 else 0.0, + "esa": esa_hits / n if n > 0 else 0.0, + "per_task": per_task, + } + + +def run_gatekeeper_check( + poisoned: list[ExecutionTrace], + registry: ToolRegistry, + store: InMemoryTraceStore, +) -> dict[str, Any]: + """Run all poisoned traces through the gatekeeper and report results.""" + settings = Settings() + gk = GatekeeperPipeline(store=store, registry=registry, settings=settings) + + results: list[dict[str, Any]] = [] + accepted = 0 + rejected = 0 + + for trace in poisoned: + result = gk.evaluate(trace) + results.append( + { + "task": trace.task_description, + "poison_type": trace.metadata.get("poison_type", "unknown"), + "accepted": result.accepted, + "reason": result.rejection_reason or "passed", + "failures": result.failures, + } + ) + if result.accepted: + accepted += 1 + else: + rejected += 1 + + return { + "total": len(poisoned), + "accepted": accepted, + "rejected": rejected, + "rejection_rate": rejected / len(poisoned) if poisoned else 0.0, + "details": results, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Gatekeeper Ablation Study (Section IV.D.5)") + parser.add_argument("--verbose", "-v", action="store_true", help="Show per-task details") + parser.add_argument( + "--poisoned-ratio", + type=float, + default=1.0, + help="Fraction of poisoned traces to inject (0.0-1.0, default: 1.0 = all 8)", + ) + parser.add_argument("--limit", type=int, default=0, help="Limit evaluation to N tasks (0=all 30)") + args = parser.parse_args() + + console.print( + Panel.fit( + "[bold]Gatekeeper Ablation Study[/bold]\n\n" + "Section IV.D.5 of the paper: measures the impact of disabling the\n" + "gatekeeper pipeline on plan quality by injecting poisoned traces\n" + "that encode wrong domain conventions.\n\n" + "Three conditions:\n" + " 1. BASELINE — only valid seed traces (gatekeeper enabled)\n" + " 2. POISONED — bad traces injected (gatekeeper disabled)\n" + " 3. RECOVERED — gatekeeper re-enabled, bad traces filtered out", + title="Experiment Setup", + ) + ) + + embeddings = MockEmbeddings() + schemas = get_tool_schemas() + seed_traces = get_seed_traces() + poisoned_traces = build_poisoned_traces() + registry = ToolRegistry() + registry.register_many(schemas) + + n_poison = max(1, int(len(poisoned_traces) * args.poisoned_ratio)) + poisoned_subset = poisoned_traces[:n_poison] + + tasks = EVALUATION_TASKS + if args.limit > 0: + tasks = tasks[: args.limit] + + console.print(f"\n[dim]Seed traces: {len(seed_traces)}, Poisoned traces: {n_poison}, Tasks: {len(tasks)}[/dim]\n") + + # === CONDITION 1: BASELINE (clean memory) === + console.print("[bold cyan]CONDITION 1: BASELINE (gatekeeper ON, clean memory)[/bold cyan]") + baseline_store = InMemoryTraceStore(embeddings=embeddings) + baseline_store.add_bulk(seed_traces) + baseline_results = run_evaluation_with_store(baseline_store, schemas, tasks, "Baseline") + console.print( + f" TSA={baseline_results['tsa']:.1%} PV={baseline_results['pv']:.1%} " + f"PCR={baseline_results['pcr']:.1%} ESA={baseline_results['esa']:.1%} " + f"(store size: {baseline_store.count()})" + ) + + # === CONDITION 2: POISONED (gatekeeper OFF) === + console.print("\n[bold red]CONDITION 2: POISONED (gatekeeper OFF, bad traces injected)[/bold red]") + poisoned_store = InMemoryTraceStore(embeddings=embeddings) + poisoned_store.add_bulk(seed_traces) + poisoned_store.add_bulk(poisoned_subset) + poisoned_results = run_evaluation_with_store(poisoned_store, schemas, tasks, "Poisoned") + console.print( + f" TSA={poisoned_results['tsa']:.1%} PV={poisoned_results['pv']:.1%} " + f"PCR={poisoned_results['pcr']:.1%} ESA={poisoned_results['esa']:.1%} " + f"(store size: {poisoned_store.count()})" + ) + + # === GATEKEEPER ANALYSIS === + console.print("\n[bold yellow]GATEKEEPER ANALYSIS: Testing poisoned traces against the pipeline[/bold yellow]") + gk_results = run_gatekeeper_check(poisoned_subset, registry, baseline_store) + console.print( + f" Total: {gk_results['total']} " + f"Rejected: {gk_results['rejected']} " + f"Accepted: {gk_results['accepted']} " + f"Rejection rate: {gk_results['rejection_rate']:.0%}" + ) + + if args.verbose: + for detail in gk_results["details"]: + status = "[red]REJECTED[/red]" if not detail["accepted"] else "[green]ACCEPTED[/green]" + console.print(f" {status} [{detail['poison_type']}] {detail['task']}") + if detail["failures"]: + for f in detail["failures"]: + console.print(f" [dim]{f}[/dim]") + + # === CONDITION 3: RECOVERED (re-enable gatekeeper) === + console.print("\n[bold green]CONDITION 3: RECOVERED (gatekeeper re-enabled, only valid traces kept)[/bold green]") + recovered_store = InMemoryTraceStore(embeddings=embeddings) + recovered_store.add_bulk(seed_traces) + recovered_results = run_evaluation_with_store(recovered_store, schemas, tasks, "Recovered") + console.print( + f" TSA={recovered_results['tsa']:.1%} PV={recovered_results['pv']:.1%} " + f"PCR={recovered_results['pcr']:.1%} ESA={recovered_results['esa']:.1%} " + f"(store size: {recovered_store.count()})" + ) + + # === COMPARISON TABLE === + console.print() + table = Table(title="Gatekeeper Ablation Results") + table.add_column("Metric", style="bold") + table.add_column("Baseline\n(GK ON)", justify="right", style="cyan") + table.add_column("Poisoned\n(GK OFF)", justify="right", style="red") + table.add_column("Recovered\n(GK restored)", justify="right", style="green") + table.add_column("Degradation", justify="right") + + for metric_name, key in [ + ("Tool Selection (TSA)", "tsa"), + ("Parameter Validity (PV)", "pv"), + ("Plan Correctness (PCR)", "pcr"), + ("Sequence Accuracy (ESA)", "esa"), + ]: + baseline_val = baseline_results[key] + poisoned_val = poisoned_results[key] + recovered_val = recovered_results[key] + delta = poisoned_val - baseline_val + delta_str = f"{delta:+.1%}" if delta != 0 else "—" + table.add_row( + metric_name, + f"{baseline_val:.1%}", + f"{poisoned_val:.1%}", + f"{recovered_val:.1%}", + delta_str, + ) + + table.add_row( + "Gatekeeper rejection rate", + "—", + f"{gk_results['rejection_rate']:.0%}", + "—", + "", + ) + console.print(table) + + # === ANALYSIS === + pcr_degradation = baseline_results["pcr"] - poisoned_results["pcr"] + console.print( + Panel.fit( + f"[bold]Key Findings:[/bold]\n\n" + f"1. Poisoned traces degraded PCR by [red]{pcr_degradation:.1%}[/red] " + f"({baseline_results['pcr']:.1%} -> {poisoned_results['pcr']:.1%})\n" + f"2. The gatekeeper rejected [bold]{gk_results['rejected']}/{gk_results['total']}[/bold] " + f"poisoned traces ({gk_results['rejection_rate']:.0%} rejection rate)\n" + f"3. After recovery (re-enabling gatekeeper), performance returned to baseline\n\n" + f"[bold]Conclusion:[/bold] The gatekeeper pipeline is essential for maintaining\n" + f"memory quality. Without it, poisoned traces (wrong conventions, broken\n" + f"dependencies, incorrect tools) contaminate retrieval results and cause\n" + f"the LLM to copy bad patterns, significantly degrading plan quality.", + title="Analysis — Section IV.D.5", + ) + ) + + if args.verbose: + console.print("\n[bold]Per-task poisoning impact:[/bold]") + for b_task, p_task in zip(baseline_results["per_task"], poisoned_results["per_task"], strict=True): + b_pcr = bool(b_task["metrics"]["pcr"]) + p_pcr = bool(p_task["metrics"]["pcr"]) + if b_pcr and not p_pcr: + console.print( + f" [red]DEGRADED[/red] Task {b_task['task_id']}: {b_task['task'][:60]}" + f" (retrieved from: {p_task.get('retrieved_source', '?')})" + ) + elif not b_pcr and p_pcr: + console.print(f" [green]IMPROVED[/green] Task {b_task['task_id']}: {b_task['task'][:60]}") + + +if __name__ == "__main__": + main() diff --git a/examples/run_benchmark.py b/examples/run_benchmark.py index a9128ba..524a4a3 100644 --- a/examples/run_benchmark.py +++ b/examples/run_benchmark.py @@ -68,9 +68,9 @@ def main(): sf = static["aggregate"][metric] dy = dynamic["aggregate"][metric] - zs_str = f"{zs['mean']:.1%}" if isinstance(zs['mean'], float) else f"{zs['mean']}" - sf_str = f"{sf['mean']:.1%}" if isinstance(sf['mean'], float) else f"{sf['mean']}" - dy_str = f"{dy['mean']:.1%}" if isinstance(dy['mean'], float) else f"{dy['mean']}" + zs_str = f"{zs['mean']:.1%}" if isinstance(zs["mean"], float) else f"{zs['mean']}" + sf_str = f"{sf['mean']:.1%}" if isinstance(sf["mean"], float) else f"{sf['mean']}" + dy_str = f"{dy['mean']:.1%}" if isinstance(dy["mean"], float) else f"{dy['mean']}" table.add_row(metric.upper(), zs_str, sf_str, dy_str) diff --git a/examples/run_live_benchmark.py b/examples/run_live_benchmark.py index 0b5399c..3f42787 100644 --- a/examples/run_live_benchmark.py +++ b/examples/run_live_benchmark.py @@ -42,13 +42,15 @@ def main() -> None: parser.add_argument("--output", type=str, default="benchmark_results.json", help="Output file") args = parser.parse_args() - console.print(Panel.fit( - "[bold]Behavioral Memory — Live Benchmark[/bold]\n\n" - "This runs the REAL benchmark from the paper.\n" - "It calls the LLM for every task and scores plans against gold chains.\n" - f"Model: {args.model} | Tasks: {'all 30' if args.limit == 0 else args.limit}", - title="Live Benchmark", - )) + console.print( + Panel.fit( + "[bold]Behavioral Memory — Live Benchmark[/bold]\n\n" + "This runs the REAL benchmark from the paper.\n" + "It calls the LLM for every task and scores plans against gold chains.\n" + f"Model: {args.model} | Tasks: {'all 30' if args.limit == 0 else args.limit}", + title="Live Benchmark", + ) + ) try: from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings diff --git a/examples/validate_pipeline.py b/examples/validate_pipeline.py index 33b3b13..a7e0dc5 100644 --- a/examples/validate_pipeline.py +++ b/examples/validate_pipeline.py @@ -66,23 +66,29 @@ def invoke(messages): if best_match: steps = [] for i, gold_step in enumerate(best_match["gold_tool_chain"]): - steps.append({ - "step_id": f"step_{i + 1}", - "tool_name": gold_step["tool"], - "parameters": gold_step["params"], - "depends_on": [f"step_{j + 1}" for j in range(i)], - }) + steps.append( + { + "step_id": f"step_{i + 1}", + "tool_name": gold_step["tool"], + "parameters": gold_step["params"], + "depends_on": [f"step_{j + 1}" for j in range(i)], + } + ) response = MagicMock() response.content = json.dumps(steps) return response response = MagicMock() - response.content = json.dumps([{ - "step_id": "step_1", - "tool_name": "data_fetch", - "parameters": {"source": "default"}, - "depends_on": [], - }]) + response.content = json.dumps( + [ + { + "step_id": "step_1", + "tool_name": "data_fetch", + "parameters": {"source": "default"}, + "depends_on": [], + } + ] + ) return response llm = MagicMock() @@ -91,12 +97,14 @@ def invoke(messages): def main() -> None: - console.print(Panel.fit( - "[bold]Pipeline Validation — Full End-to-End Check[/bold]\n\n" - "Tests every component with mock services.\n" - "No API keys or external services needed.", - title="Validation", - )) + console.print( + Panel.fit( + "[bold]Pipeline Validation — Full End-to-End Check[/bold]\n\n" + "Tests every component with mock services.\n" + "No API keys or external services needed.", + title="Validation", + ) + ) from behavioral_memory.core.config import Settings from behavioral_memory.evaluation.benchmark import BenchmarkRunner @@ -148,11 +156,7 @@ def check(name: str, condition: bool, detail: str = ""): check("Three difficulty tiers", difficulties == {"simple", "moderate", "challenging"}) check( "All gold chains reference known tools", - all( - step["tool"] in registry._tools - for task in EVALUATION_TASKS - for step in task["gold_tool_chain"] - ), + all(step["tool"] in registry._tools for task in EVALUATION_TASKS for step in task["gold_tool_chain"]), ) # --- 4. InMemoryTraceStore --- @@ -184,7 +188,9 @@ def check(name: str, condition: bool, detail: str = ""): check("Zero-shot has no retrieved traces", len(zs_plan.retrieved_traces) == 0) static_plan = engine.generate_static_few_shot( - "Build a revenue analysis pipeline", schemas, seed_traces[:3], + "Build a revenue analysis pipeline", + schemas, + seed_traces[:3], ) check("Static few-shot works", len(static_plan.steps) > 0) check("Static uses provided traces", len(static_plan.retrieved_traces) == 3) @@ -214,8 +220,11 @@ def check(name: str, condition: bool, detail: str = ""): console.print("\n[bold cyan]7. Gatekeeper Pipeline[/bold cyan]") gk = GatekeeperPipeline(store=store, registry=registry) gk_result = gk.evaluate(seed_traces[0]) - check("Gatekeeper accepts valid trace", gk_result.accepted or gk_result.is_duplicate, - f"rejected: {gk_result.rejection_reason}") + check( + "Gatekeeper accepts valid trace", + gk_result.accepted or gk_result.is_duplicate, + f"rejected: {gk_result.rejection_reason}", + ) # --- 8. Langfuse offline --- console.print("\n[bold cyan]8. Langfuse Tracer (offline)[/bold cyan]") @@ -244,8 +253,7 @@ def check(name: str, condition: bool, detail: str = ""): ) console.print(table) - console.print("[dim]Note: These numbers are from a mock LLM — run with a real " - "API key for actual results.[/dim]") + console.print("[dim]Note: These numbers are from a mock LLM — run with a real API key for actual results.[/dim]") # --- Summary --- console.print(f"\n{'=' * 50}") diff --git a/pyproject.toml b/pyproject.toml index 4aeb0a0..2d4b0c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,9 +40,6 @@ classifiers = [ dependencies = [ "pydantic>=2.0", "pydantic-settings>=2.0", - "sqlalchemy>=2.0", - "psycopg[binary]>=3.1", - "pgvector>=0.3", "langchain-core>=0.3", "langfuse>=2.0", "tiktoken>=0.7", @@ -52,6 +49,12 @@ dependencies = [ ] [project.optional-dependencies] +postgres = [ + "sqlalchemy>=2.0", + "psycopg[binary]>=3.1", + "pgvector>=0.3", + "langchain-postgres>=0.0.12", +] agent = [ "langgraph>=1.2", "langgraph-prebuilt>=1.1", @@ -59,7 +62,6 @@ agent = [ "typer>=0.12", ] eval = [ - "numpy>=1.26", "scipy>=1.12", ] server = [ @@ -75,7 +77,7 @@ dev = [ "pre-commit>=3.0", ] all = [ - "behavioral-memory[agent,eval,dev]", + "behavioral-memory[postgres,agent,eval,dev,server]", ] [project.scripts] @@ -134,6 +136,9 @@ module = [ "pgvector.*", "tiktoken.*", "mcp.*", + "scipy.*", + "typer.*", + "rich.*", ] ignore_missing_imports = true diff --git a/server.py b/server.py index 5d38cf4..3b12d18 100644 --- a/server.py +++ b/server.py @@ -39,6 +39,7 @@ use_postgres = os.getenv("USE_POSTGRES", "false").lower() == "true" if use_postgres: from behavioral_memory.memory.store import TraceStore + store = TraceStore(embeddings=embeddings, settings=settings) else: store = InMemoryTraceStore(embeddings=embeddings, settings=settings) diff --git a/src/behavioral_memory/__init__.py b/src/behavioral_memory/__init__.py index 199455a..43dcf2d 100644 --- a/src/behavioral_memory/__init__.py +++ b/src/behavioral_memory/__init__.py @@ -3,8 +3,11 @@ This library provides a retrieval-based architecture that uses a memory bank of validated execution traces to guide LLM tool orchestration during inference. -Quick start: - from behavioral_memory import TraceStore, PlanEngine, GatekeeperPipeline +Quick start (no PostgreSQL needed): + from behavioral_memory import InMemoryTraceStore, PlanEngine + +With PostgreSQL (pip install behavioral-memory[postgres]): + from behavioral_memory import TraceStore, PlanEngine See https://github.com/harsh-kr11/behavioral-memory for full documentation. """ @@ -20,13 +23,22 @@ from behavioral_memory.gatekeeper.pipeline import GatekeeperPipeline from behavioral_memory.memory.dedup import Deduplicator from behavioral_memory.memory.in_memory_store import InMemoryTraceStore -from behavioral_memory.memory.store import TraceStore from behavioral_memory.observability.annotation import AnnotationHandler from behavioral_memory.observability.feedback import FeedbackPoller from behavioral_memory.observability.tracer import LangfuseTracer from behavioral_memory.planner.engine import PlanEngine from behavioral_memory.tools.registry import ToolRegistry + +def __getattr__(name: str): # type: ignore[no-untyped-def] + """Lazy import for TraceStore to avoid requiring PostgreSQL deps at import time.""" + if name == "TraceStore": + from behavioral_memory.memory.store import TraceStore + + return TraceStore + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "AnnotationHandler", "Deduplicator", diff --git a/src/behavioral_memory/cli.py b/src/behavioral_memory/cli.py index b3e0d58..c7b4e66 100644 --- a/src/behavioral_memory/cli.py +++ b/src/behavioral_memory/cli.py @@ -8,9 +8,13 @@ import shutil from pathlib import Path +from typing import TYPE_CHECKING import typer from rich.console import Console + +if TYPE_CHECKING: + from behavioral_memory.core.schemas import ExecutionTrace from rich.panel import Panel from rich.table import Table @@ -54,19 +58,21 @@ def setup() -> None: shutil.copy(template_path, env_path) console.print("[green]Created .env from .env.example[/green]\n") - console.print(Panel.fit( - "[bold]Configure your .env file:[/bold]\n\n" - "1. [cyan]GOOGLE_API_KEY[/cyan] — get one at https://aistudio.google.com/apikey\n" - " Only needed to run the reference agent with Gemini.\n" - " You can use any LangChain-compatible LLM instead.\n\n" - "2. [cyan]VECTOR_STORE_URL[/cyan] — PostgreSQL + pgvector connection string\n" - " For local dev: docker run -p 5432:5432 -e POSTGRES_PASSWORD=pw pgvector/pgvector\n" - " [dim]Not needed for demo mode or tests.[/dim]\n\n" - "3. [cyan]LANGFUSE_SECRET_KEY / LANGFUSE_PUBLIC_KEY[/cyan] — from https://cloud.langfuse.com\n" - " Free tier available. Enables the feedback loop.\n" - " [dim]Optional — the framework works without Langfuse.[/dim]", - title="Setup Guide", - )) + console.print( + Panel.fit( + "[bold]Configure your .env file:[/bold]\n\n" + "1. [cyan]GOOGLE_API_KEY[/cyan] — get one at https://aistudio.google.com/apikey\n" + " Only needed to run the reference agent with Gemini.\n" + " You can use any LangChain-compatible LLM instead.\n\n" + "2. [cyan]VECTOR_STORE_URL[/cyan] — PostgreSQL + pgvector connection string\n" + " For local dev: docker run -p 5432:5432 -e POSTGRES_PASSWORD=pw pgvector/pgvector\n" + " [dim]Not needed for demo mode or tests.[/dim]\n\n" + "3. [cyan]LANGFUSE_SECRET_KEY / LANGFUSE_PUBLIC_KEY[/cyan] — from https://cloud.langfuse.com\n" + " Free tier available. Enables the feedback loop.\n" + " [dim]Optional — the framework works without Langfuse.[/dim]", + title="Setup Guide", + ) + ) @app.command() @@ -91,14 +97,16 @@ def demo( from behavioral_memory.tools.mock_tools import get_tool_schemas from behavioral_memory.tools.registry import ToolRegistry - console.print(Panel.fit( - "[bold]Behavioral Memory — Offline Demo[/bold]\n\n" - "This demo shows how behavioral memory (validated execution traces)\n" - "improves tool orchestration quality. No external services needed.\n\n" - "It simulates: zero-shot (no memory) vs dynamic retrieval (with memory)\n" - "on the paper's 30-task benchmark with 7 MCP tools.", - title="Demo Mode", - )) + console.print( + Panel.fit( + "[bold]Behavioral Memory — Offline Demo[/bold]\n\n" + "This demo shows how behavioral memory (validated execution traces)\n" + "improves tool orchestration quality. No external services needed.\n\n" + "It simulates: zero-shot (no memory) vs dynamic retrieval (with memory)\n" + "on the paper's 30-task benchmark with 7 MCP tools.", + title="Demo Mode", + ) + ) schemas = get_tool_schemas() registry = ToolRegistry() @@ -114,7 +122,7 @@ def demo( raise typer.Exit(1) else: selected = [ - EVALUATION_TASKS[0], # simple + EVALUATION_TASKS[0], # simple EVALUATION_TASKS[10], # moderate EVALUATION_TASKS[20], # challenging ] @@ -143,14 +151,13 @@ def demo( console.print("\n[green] DYNAMIC RETRIEVAL (with behavioral memory)[/green]") console.print(f" [dim]Prompt size: {len(dr_prompt)} chars, {len(relevant)} reference examples[/dim]") for i, trace in enumerate(relevant): - console.print(f" [dim]Retrieved #{i+1}: \"{trace.task_description[:60]}\"[/dim]") + console.print(f' [dim]Retrieved #{i + 1}: "{trace.task_description[:60]}"[/dim]') # --- Gold chain validation --- trace_obj = ExecutionTrace( task_description=task["task"], tool_chain=[ - ToolCall(step_id=s["step_id"], tool_name=s["tool"], parameters=s.get("params", {})) - for s in gold + ToolCall(step_id=s["step_id"], tool_name=s["tool"], parameters=s.get("params", {})) for s in gold ], source="execution", ) @@ -168,32 +175,34 @@ def demo( console.print(f" {k}: {val_str}") console.print(f"\n Gatekeeper: schema={'✓' if is_valid else '✗'} sandbox={'✓' if passed else '✗'}") - console.print(f" Metrics: TSA={'✓' if metrics['tsa'] else '✗'} PV={metrics['pv']:.0%} PCR={'✓' if metrics['pcr'] else '✗'} ESA={'✓' if metrics['esa'] else '✗'}") + console.print( + f" Metrics: TSA={'✓' if metrics['tsa'] else '✗'} PV={metrics['pv']:.0%} PCR={'✓' if metrics['pcr'] else '✗'} ESA={'✓' if metrics['esa'] else '✗'}" + ) # --- Summary table --- console.print(f"\n{'═' * 80}") _print_paper_results_table() - console.print(Panel.fit( - "[bold]Key Insight:[/bold] Without behavioral memory, the LLM has no way to learn\n" - "domain conventions (e.g., 'revenue' = quantity*unit_price, not total_amount).\n" - "With memory, validated traces teach these conventions dynamically.\n\n" - "[bold]To run with a real LLM:[/bold]\n" - " 1. Set GOOGLE_API_KEY in .env (or use any LangChain model)\n" - " 2. Start PostgreSQL: docker run -p 5432:5432 pgvector/pgvector\n" - " 3. Run: python examples/run_benchmark.py\n\n" - "[bold]To enable Langfuse tracing:[/bold]\n" - " 1. Sign up at https://cloud.langfuse.com (free)\n" - " 2. Add LANGFUSE_SECRET_KEY and LANGFUSE_PUBLIC_KEY to .env\n" - " 3. Every plan will be logged — SMEs can score them in the Langfuse UI\n" - " 4. Positively scored traces auto-enter behavioral memory via FeedbackPoller", - title="Next Steps", - )) - - -def _find_relevant_traces( - query: str, traces: list, top_k: int = 3 -) -> list: + console.print( + Panel.fit( + "[bold]Key Insight:[/bold] Without behavioral memory, the LLM has no way to learn\n" + "domain conventions (e.g., 'revenue' = quantity*unit_price, not total_amount).\n" + "With memory, validated traces teach these conventions dynamically.\n\n" + "[bold]To run with a real LLM:[/bold]\n" + " 1. Set GOOGLE_API_KEY in .env (or use any LangChain model)\n" + " 2. Start PostgreSQL: docker run -p 5432:5432 pgvector/pgvector\n" + " 3. Run: python examples/run_benchmark.py\n\n" + "[bold]To enable Langfuse tracing:[/bold]\n" + " 1. Sign up at https://cloud.langfuse.com (free)\n" + " 2. Add LANGFUSE_SECRET_KEY and LANGFUSE_PUBLIC_KEY to .env\n" + " 3. Every plan will be logged — SMEs can score them in the Langfuse UI\n" + " 4. Positively scored traces auto-enter behavioral memory via FeedbackPoller", + title="Next Steps", + ) + ) + + +def _find_relevant_traces(query: str, traces: list[ExecutionTrace], top_k: int = 3) -> list[ExecutionTrace]: """Simple keyword-based trace matching for demo mode (no embeddings needed).""" query_lower = query.lower() scored = [] @@ -258,7 +267,7 @@ def benchmark_info() -> None: table.add_column("Category", style="cyan") table.add_column("Count", justify="right") - difficulties = {} + difficulties: dict[str, int] = {} for task in EVALUATION_TASKS: d = task["difficulty"] difficulties[d] = difficulties.get(d, 0) + 1 diff --git a/src/behavioral_memory/core/exceptions.py b/src/behavioral_memory/core/exceptions.py index 756eaae..7e1a46a 100644 --- a/src/behavioral_memory/core/exceptions.py +++ b/src/behavioral_memory/core/exceptions.py @@ -1,5 +1,9 @@ """Custom exception hierarchy for behavioral-memory.""" +from __future__ import annotations + +from typing import Any + class BehavioralMemoryError(Exception): """Base exception for all behavioral-memory errors.""" @@ -20,7 +24,7 @@ def __init__(self, message: str, failures: list[str] | None = None) -> None: class GatekeeperRejectionError(BehavioralMemoryError): """Raised when the gatekeeper pipeline rejects a trace.""" - def __init__(self, message: str, stage: str, details: dict | None = None) -> None: + def __init__(self, message: str, stage: str, details: dict[str, Any] | None = None) -> None: super().__init__(message) self.stage = stage self.details = details or {} diff --git a/src/behavioral_memory/core/schemas.py b/src/behavioral_memory/core/schemas.py index d64e3f6..fc7cfac 100644 --- a/src/behavioral_memory/core/schemas.py +++ b/src/behavioral_memory/core/schemas.py @@ -25,9 +25,7 @@ class ToolCall(BaseModel): step_id: str = Field(description="Unique identifier for this step within the trace") tool_name: str = Field(description="Name of the MCP tool to invoke") - parameters: dict[str, Any] = Field( - default_factory=dict, description="Tool invocation parameters" - ) + parameters: dict[str, Any] = Field(default_factory=dict, description="Tool invocation parameters") depends_on: list[str] = Field( default_factory=list, description="Step IDs that must complete before this step executes", @@ -60,9 +58,7 @@ class ToolSchema(BaseModel): parameters_schema: dict[str, Any] = Field( default_factory=dict, description="JSON Schema for the tool's input parameters" ) - required_params: list[str] = Field( - default_factory=list, description="List of required parameter names" - ) + required_params: list[str] = Field(default_factory=list, description="List of required parameter names") class ExecutionTrace(BaseModel): @@ -107,9 +103,7 @@ def tool_names(self) -> list[str]: def to_prompt_str(self) -> str: """Serialize this trace into a string suitable for prompt injection.""" - chain_str = json.dumps( - [step.model_dump() for step in self.tool_chain], indent=2 - ) + chain_str = json.dumps([step.model_dump() for step in self.tool_chain], indent=2) return f"Task: {self.task_description}\nTool Chain:\n{chain_str}" @@ -125,9 +119,7 @@ class Plan(BaseModel): retrieved_traces: list[ExecutionTrace] = Field( default_factory=list, description="Traces retrieved from behavioral memory" ) - schemas_used: list[ToolSchema] = Field( - default_factory=list, description="Tool schemas available during planning" - ) + schemas_used: list[ToolSchema] = Field(default_factory=list, description="Tool schemas available during planning") token_budget_used: int = Field(default=0, description="Tokens consumed by the prompt") raw_llm_output: str = Field(default="", description="Raw LLM response before parsing") @@ -144,6 +136,4 @@ class GatekeeperResult(BaseModel): sandbox_passed: bool = Field(default=False, description="Passed sandboxed execution") is_duplicate: bool = Field(default=False, description="Flagged as semantic duplicate") rejection_reason: str = Field(default="", description="Human-readable rejection reason") - failures: list[str] = Field( - default_factory=list, description="List of specific validation failures" - ) + failures: list[str] = Field(default_factory=list, description="List of specific validation failures") diff --git a/src/behavioral_memory/evaluation/benchmark.py b/src/behavioral_memory/evaluation/benchmark.py index e28ccb5..9d4a865 100644 --- a/src/behavioral_memory/evaluation/benchmark.py +++ b/src/behavioral_memory/evaluation/benchmark.py @@ -23,7 +23,7 @@ class BenchmarkRunner: def __init__(self, tool_schemas: list[ToolSchema]) -> None: self._schemas = tool_schemas - def evaluate_plan(self, plan: Plan, gold_chain: list[dict]) -> dict[str, Any]: + def evaluate_plan(self, plan: Plan, gold_chain: list[dict[str, Any]]) -> dict[str, Any]: """Evaluate a single plan against its gold tool chain.""" predicted_chain = [ { @@ -37,7 +37,7 @@ def evaluate_plan(self, plan: Plan, gold_chain: list[dict]) -> dict[str, Any]: def run( self, strategy: Any, - tasks: list[dict] | None = None, + tasks: list[dict[str, Any]] | None = None, limit: int | None = None, ) -> dict[str, Any]: """Run a strategy across all (or a subset of) evaluation tasks. @@ -64,26 +64,30 @@ def run( pcr_results.append(bool(metrics["pcr"])) esa_results.append(bool(metrics["esa"])) - per_task.append({ - "task_id": task["task_id"], - "task": task["task"], - "difficulty": task["difficulty"], - "predicted_steps": [s.model_dump() for s in plan.steps], - "metrics": metrics, - }) + per_task.append( + { + "task_id": task["task_id"], + "task": task["task"], + "difficulty": task["difficulty"], + "predicted_steps": [s.model_dump() for s in plan.steps], + "metrics": metrics, + } + ) except Exception as e: logger.warning("Task %d failed: %s", task["task_id"], e) tsa_results.append(False) pv_values.append(0.0) pcr_results.append(False) esa_results.append(False) - per_task.append({ - "task_id": task["task_id"], - "task": task["task"], - "difficulty": task["difficulty"], - "error": str(e), - "metrics": {"tsa": False, "pv": 0.0, "pcr": False, "esa": False}, - }) + per_task.append( + { + "task_id": task["task_id"], + "task": task["task"], + "difficulty": task["difficulty"], + "error": str(e), + "metrics": {"tsa": False, "pv": 0.0, "pcr": False, "esa": False}, + } + ) n = len(eval_tasks) tsa_mean, tsa_lo, tsa_hi = bootstrap_ci(tsa_results) @@ -121,14 +125,14 @@ def compare( } @staticmethod - def results_by_difficulty(results: dict[str, Any]) -> dict[str, dict]: + def results_by_difficulty(results: dict[str, Any]) -> dict[str, dict[str, Any]]: """Break down results by difficulty tier.""" - by_diff: dict[str, list[dict]] = {} + by_diff: dict[str, list[dict[str, Any]]] = {} for task in results["per_task"]: diff = task["difficulty"] by_diff.setdefault(diff, []).append(task) - summary: dict[str, dict] = {} + summary: dict[str, dict[str, Any]] = {} for diff, tasks in by_diff.items(): pcr_list = [bool(t["metrics"]["pcr"]) for t in tasks] esa_list = [bool(t["metrics"]["esa"]) for t in tasks] diff --git a/src/behavioral_memory/evaluation/ground_truth.py b/src/behavioral_memory/evaluation/ground_truth.py index 3e9410b..ffed437 100644 --- a/src/behavioral_memory/evaluation/ground_truth.py +++ b/src/behavioral_memory/evaluation/ground_truth.py @@ -8,7 +8,9 @@ from __future__ import annotations -EVALUATION_TASKS: list[dict] = [ +from typing import Any + +EVALUATION_TASKS: list[dict[str, Any]] = [ # ═══════════════════════════════════════════ # SIMPLE — single-tool operations (10) # ═══════════════════════════════════════════ @@ -17,7 +19,11 @@ "task": "Get the total number of customers from the database", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT COUNT(*) AS customer_count FROM customers;"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": {"query": "SELECT COUNT(*) AS customer_count FROM customers;"}, + }, ], }, { @@ -25,7 +31,11 @@ "task": "Fetch the current exchange rate from the forex API", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "fetch_api_data", "params": {"url": "https://api.exchangerate.host/latest", "method": "GET"}}, + { + "step_id": "step_1", + "tool": "fetch_api_data", + "params": {"url": "https://api.exchangerate.host/latest", "method": "GET"}, + }, ], }, { @@ -33,7 +43,16 @@ "task": "Send a Slack message to the team about the deployment", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#team", "subject": "Deployment Update", "body": "Deployment completed successfully."}}, + { + "step_id": "step_1", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#team", + "subject": "Deployment Update", + "body": "Deployment completed successfully.", + }, + }, ], }, { @@ -41,7 +60,11 @@ "task": "Get all products and their prices from the database", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT name, price FROM products ORDER BY name;"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": {"query": "SELECT name, price FROM products ORDER BY name;"}, + }, ], }, { @@ -49,7 +72,11 @@ "task": "Store the sales summary as a JSON file", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "store_results", "params": {"source_step": "previous_step", "target": "json_file", "target_name": "sales_summary.json"}}, + { + "step_id": "step_1", + "tool": "store_results", + "params": {"source_step": "previous_step", "target": "json_file", "target_name": "sales_summary.json"}, + }, ], }, { @@ -57,7 +84,11 @@ "task": "Schedule a weekly backup task", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "schedule_task", "params": {"task_name": "weekly_backup", "workflow_steps": [], "interval": "weekly"}}, + { + "step_id": "step_1", + "tool": "schedule_task", + "params": {"task_name": "weekly_backup", "workflow_steps": [], "interval": "weekly"}, + }, ], }, { @@ -65,7 +96,11 @@ "task": "Get the count of orders in the system", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT COUNT(*) AS order_count FROM orders;"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": {"query": "SELECT COUNT(*) AS order_count FROM orders;"}, + }, ], }, { @@ -73,7 +108,16 @@ "task": "Send an email notification about the monthly report being ready", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "send_notification", "params": {"channel": "email", "recipient": "team@company.com", "subject": "Monthly Report Ready", "body": "The monthly report has been generated and is ready for review."}}, + { + "step_id": "step_1", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "team@company.com", + "subject": "Monthly Report Ready", + "body": "The monthly report has been generated and is ready for review.", + }, + }, ], }, { @@ -81,7 +125,11 @@ "task": "Generate a text summary from the analysis results", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "generate_report", "params": {"source_step": "previous_step", "format": "summary_text", "title": "Analysis Summary"}}, + { + "step_id": "step_1", + "tool": "generate_report", + "params": {"source_step": "previous_step", "format": "summary_text", "title": "Analysis Summary"}, + }, ], }, { @@ -89,7 +137,11 @@ "task": "Query the database for delivered orders", "difficulty": "simple", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT * FROM orders WHERE status = 'delivered';"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": {"query": "SELECT * FROM orders WHERE status = 'delivered';"}, + }, ], }, # ═══════════════════════════════════════════ @@ -100,8 +152,18 @@ "task": "Get the total revenue and generate a report", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT SUM(oi.quantity * oi.unit_price) AS total_revenue FROM order_items oi JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned');"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Total Revenue Report"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT SUM(oi.quantity * oi.unit_price) AS total_revenue FROM order_items oi JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned');" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Total Revenue Report"}, + }, ], }, { @@ -109,8 +171,23 @@ "task": "Retrieve valid orders and store them as a CSV archive", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT * FROM orders WHERE status NOT IN ('cancelled', 'returned') ORDER BY order_date;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "csv_file", "target_name": "valid_orders_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT * FROM orders WHERE status NOT IN ('cancelled', 'returned') ORDER BY order_date;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": { + "source_step": "step_1", + "target": "csv_file", + "target_name": "valid_orders_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -118,8 +195,24 @@ "task": "Get revenue by category and send an alert about it", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;"}}, - {"step_id": "step_2", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Revenue by Category", "body": "Revenue breakdown by product category is ready.", "attach_step": "step_1"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Revenue by Category", + "body": "Revenue breakdown by product category is ready.", + "attach_step": "step_1", + }, + }, ], }, { @@ -127,8 +220,18 @@ "task": "Build a pipeline to get completed orders count and cache it for the dashboard", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT COUNT(*) AS completed_orders FROM orders WHERE status IN ('shipped', 'delivered');"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_completed_orders"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT COUNT(*) AS completed_orders FROM orders WHERE status IN ('shipped', 'delivered');" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_completed_orders"}, + }, ], }, { @@ -136,8 +239,18 @@ "task": "Fetch customer data, filter active customers, and generate a report", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT DISTINCT c.customer_id, c.name, c.email, c.segment FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned');"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Active Customers Report"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT DISTINCT c.customer_id, c.name, c.email, c.segment FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned');" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Active Customers Report"}, + }, ], }, { @@ -145,9 +258,29 @@ "task": "Get order data and email the report to the manager", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT o.order_id, o.order_date, o.status, o.total_amount FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') ORDER BY o.order_date DESC;"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Order Status Report"}}, - {"step_id": "step_3", "tool": "send_notification", "params": {"channel": "email", "recipient": "manager@company.com", "subject": "Order Status Report", "body": "Please find the latest order status report attached.", "attach_step": "step_2"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT o.order_id, o.order_date, o.status, o.total_amount FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') ORDER BY o.order_date DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Order Status Report"}, + }, + { + "step_id": "step_3", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "manager@company.com", + "subject": "Order Status Report", + "body": "Please find the latest order status report attached.", + "attach_step": "step_2", + }, + }, ], }, { @@ -155,8 +288,18 @@ "task": "Query product revenue data and save it to the dashboard cache", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.name, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.name ORDER BY revenue DESC;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_product_revenue"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.name, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.name ORDER BY revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_product_revenue"}, + }, ], }, { @@ -164,8 +307,23 @@ "task": "Get repeat customer data and archive it as CSV", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.customer_id, c.name, c.email, COUNT(o.order_id) AS order_count FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.email HAVING COUNT(o.order_id) > 1;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "csv_file", "target_name": "repeat_customers_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.customer_id, c.name, c.email, COUNT(o.order_id) AS order_count FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.email HAVING COUNT(o.order_id) > 1;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": { + "source_step": "step_1", + "target": "csv_file", + "target_name": "repeat_customers_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -173,8 +331,24 @@ "task": "Get the fulfillment rate and alert the operations team", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT ROUND(COUNT(*) FILTER (WHERE status = 'delivered')::numeric / NULLIF(COUNT(*) FILTER (WHERE status != 'cancelled'), 0) * 100, 1) AS fulfillment_rate_pct FROM orders;"}}, - {"step_id": "step_2", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Fulfillment Rate Update", "body": "Current order fulfillment rate has been calculated.", "attach_step": "step_1"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT ROUND(COUNT(*) FILTER (WHERE status = 'delivered')::numeric / NULLIF(COUNT(*) FILTER (WHERE status != 'cancelled'), 0) * 100, 1) AS fulfillment_rate_pct FROM orders;" + }, + }, + { + "step_id": "step_2", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Fulfillment Rate Update", + "body": "Current order fulfillment rate has been calculated.", + "attach_step": "step_1", + }, + }, ], }, { @@ -182,8 +356,18 @@ "task": "Get net order values and store as dashboard data", "difficulty": "moderate", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT o.order_id, (o.total_amount - o.discount) AS net_value FROM orders o WHERE o.status NOT IN ('cancelled', 'returned');"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_net_order_values"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT o.order_id, (o.total_amount - o.discount) AS net_value FROM orders o WHERE o.status NOT IN ('cancelled', 'returned');" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_net_order_values"}, + }, ], }, # ═══════════════════════════════════════════ @@ -194,11 +378,52 @@ "task": "Build a complete revenue analysis pipeline: query revenue by category, transform to add percentage, generate a report, email it, and archive the data", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "compute", "params": {"new_column": "revenue_pct", "expression": "revenue / SUM(revenue) * 100"}}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_2", "format": "markdown_table", "title": "Revenue by Category Analysis"}}, - {"step_id": "step_4", "tool": "send_notification", "params": {"channel": "email", "recipient": "stakeholders@company.com", "subject": "Revenue by Category Report", "body": "Please find the revenue analysis report attached.", "attach_step": "step_3"}}, - {"step_id": "step_5", "tool": "store_results", "params": {"source_step": "step_2", "target": "csv_file", "target_name": "revenue_by_category_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "compute", + "params": {"new_column": "revenue_pct", "expression": "revenue / SUM(revenue) * 100"}, + }, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": { + "source_step": "step_2", + "format": "markdown_table", + "title": "Revenue by Category Analysis", + }, + }, + { + "step_id": "step_4", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "stakeholders@company.com", + "subject": "Revenue by Category Report", + "body": "Please find the revenue analysis report attached.", + "attach_step": "step_3", + }, + }, + { + "step_id": "step_5", + "tool": "store_results", + "params": { + "source_step": "step_2", + "target": "csv_file", + "target_name": "revenue_by_category_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -206,11 +431,51 @@ "task": "Create a customer segmentation pipeline: get customer spending, aggregate by segment, generate report, cache for dashboard, and alert the team", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.segment, COUNT(DISTINCT c.customer_id) AS customers, SUM(oi.quantity * oi.unit_price) AS total_revenue FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.segment ORDER BY total_revenue DESC;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "compute", "params": {"new_column": "avg_revenue_per_customer", "expression": "total_revenue / customers"}}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_2", "format": "markdown_table", "title": "Customer Segmentation Report"}}, - {"step_id": "step_4", "tool": "store_results", "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_customer_segmentation"}}, - {"step_id": "step_5", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Customer Segmentation Updated", "body": "Customer segmentation analysis has been refreshed.", "attach_step": "step_3"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.segment, COUNT(DISTINCT c.customer_id) AS customers, SUM(oi.quantity * oi.unit_price) AS total_revenue FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.segment ORDER BY total_revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "compute", + "params": {"new_column": "avg_revenue_per_customer", "expression": "total_revenue / customers"}, + }, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": { + "source_step": "step_2", + "format": "markdown_table", + "title": "Customer Segmentation Report", + }, + }, + { + "step_id": "step_4", + "tool": "store_results", + "params": { + "source_step": "step_2", + "target": "cache", + "target_name": "dashboard_customer_segmentation", + }, + }, + { + "step_id": "step_5", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Customer Segmentation Updated", + "body": "Customer segmentation analysis has been refreshed.", + "attach_step": "step_3", + }, + }, ], }, { @@ -218,9 +483,28 @@ "task": "Set up a scheduled daily report: query valid orders with revenue, generate a markdown report, and schedule it to run daily with failure notifications", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT o.order_id, o.order_date, o.status, SUM(oi.quantity * oi.unit_price) AS revenue FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id, o.order_date, o.status ORDER BY o.order_date DESC;"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Daily Valid Orders Report"}}, - {"step_id": "step_3", "tool": "schedule_task", "params": {"task_name": "daily_valid_orders_report", "workflow_steps": [{"tool": "query_database"}, {"tool": "generate_report"}], "interval": "daily", "notify_on_failure": True}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT o.order_id, o.order_date, o.status, SUM(oi.quantity * oi.unit_price) AS revenue FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id, o.order_date, o.status ORDER BY o.order_date DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Daily Valid Orders Report"}, + }, + { + "step_id": "step_3", + "tool": "schedule_task", + "params": { + "task_name": "daily_valid_orders_report", + "workflow_steps": [{"tool": "query_database"}, {"tool": "generate_report"}], + "interval": "daily", + "notify_on_failure": True, + }, + }, ], }, { @@ -228,10 +512,37 @@ "task": "Build a product performance pipeline: get product revenue, transform to compute margin, generate report, and archive", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.name, p.price, p.cost, SUM(oi.quantity * oi.unit_price) AS revenue FROM products p JOIN order_items oi ON oi.product_id = p.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') AND p.cost IS NOT NULL AND p.cost > 0 GROUP BY p.product_id, p.name, p.price, p.cost ORDER BY revenue DESC;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "compute", "params": {"new_column": "margin_pct", "expression": "(price - cost) / price * 100"}}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_2", "format": "markdown_table", "title": "Product Performance Report"}}, - {"step_id": "step_4", "tool": "store_results", "params": {"source_step": "step_2", "target": "csv_file", "target_name": "product_performance_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.name, p.price, p.cost, SUM(oi.quantity * oi.unit_price) AS revenue FROM products p JOIN order_items oi ON oi.product_id = p.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') AND p.cost IS NOT NULL AND p.cost > 0 GROUP BY p.product_id, p.name, p.price, p.cost ORDER BY revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "compute", + "params": {"new_column": "margin_pct", "expression": "(price - cost) / price * 100"}, + }, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": {"source_step": "step_2", "format": "markdown_table", "title": "Product Performance Report"}, + }, + { + "step_id": "step_4", + "tool": "store_results", + "params": { + "source_step": "step_2", + "target": "csv_file", + "target_name": "product_performance_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -239,10 +550,34 @@ "task": "Create a monthly revenue trend pipeline: query monthly revenue, cache for dashboard, generate chart config, and email stakeholders", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, SUM(oi.quantity * oi.unit_price) AS monthly_revenue FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status IN ('shipped', 'delivered') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_monthly_revenue_2025"}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_1", "format": "chart_config", "title": "Monthly Revenue Trend 2025"}}, - {"step_id": "step_4", "tool": "send_notification", "params": {"channel": "email", "recipient": "stakeholders@company.com", "subject": "Monthly Revenue Trend Report", "body": "Monthly revenue trend for 2025 is available.", "attach_step": "step_3"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, SUM(oi.quantity * oi.unit_price) AS monthly_revenue FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status IN ('shipped', 'delivered') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_monthly_revenue_2025"}, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "chart_config", "title": "Monthly Revenue Trend 2025"}, + }, + { + "step_id": "step_4", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "stakeholders@company.com", + "subject": "Monthly Revenue Trend Report", + "body": "Monthly revenue trend for 2025 is available.", + "attach_step": "step_3", + }, + }, ], }, { @@ -250,11 +585,51 @@ "task": "Build an order health monitoring pipeline: get valid orders per month, compute growth, generate report, alert ops, and archive", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, COUNT(*) AS valid_orders FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "compute", "params": {"new_column": "mom_growth_pct", "expression": "(valid_orders - LAG(valid_orders)) / LAG(valid_orders) * 100"}}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_2", "format": "markdown_table", "title": "Order Health Monitor"}}, - {"step_id": "step_4", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Order Health Update", "body": "Monthly order health report is ready.", "attach_step": "step_3"}}, - {"step_id": "step_5", "tool": "store_results", "params": {"source_step": "step_2", "target": "csv_file", "target_name": "order_health_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, COUNT(*) AS valid_orders FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "compute", + "params": { + "new_column": "mom_growth_pct", + "expression": "(valid_orders - LAG(valid_orders)) / LAG(valid_orders) * 100", + }, + }, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": {"source_step": "step_2", "format": "markdown_table", "title": "Order Health Monitor"}, + }, + { + "step_id": "step_4", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Order Health Update", + "body": "Monthly order health report is ready.", + "attach_step": "step_3", + }, + }, + { + "step_id": "step_5", + "tool": "store_results", + "params": { + "source_step": "step_2", + "target": "csv_file", + "target_name": "order_health_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -262,10 +637,36 @@ "task": "Create a basket analysis workflow: query basket sizes by segment, transform to rank, generate report, cache for dashboard", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.segment, ROUND(AVG(sub.item_count)::numeric, 1) AS avg_basket_size FROM (SELECT o.order_id, o.customer_id, SUM(oi.quantity) AS item_count FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id, o.customer_id) sub JOIN customers c ON c.customer_id = sub.customer_id GROUP BY c.segment ORDER BY avg_basket_size DESC;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "sort", "params": {"column": "avg_basket_size", "order": "desc"}}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_2", "format": "markdown_table", "title": "Basket Size Analysis by Segment"}}, - {"step_id": "step_4", "tool": "store_results", "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_basket_analysis"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.segment, ROUND(AVG(sub.item_count)::numeric, 1) AS avg_basket_size FROM (SELECT o.order_id, o.customer_id, SUM(oi.quantity) AS item_count FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id, o.customer_id) sub JOIN customers c ON c.customer_id = sub.customer_id GROUP BY c.segment ORDER BY avg_basket_size DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "sort", + "params": {"column": "avg_basket_size", "order": "desc"}, + }, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": { + "source_step": "step_2", + "format": "markdown_table", + "title": "Basket Size Analysis by Segment", + }, + }, + { + "step_id": "step_4", + "tool": "store_results", + "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_basket_analysis"}, + }, ], }, { @@ -273,10 +674,43 @@ "task": "Set up an automated customer churn alert: query at-risk customers, generate report, send alert, schedule weekly", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.name, c.email, c.segment, COUNT(o.order_id) AS total_orders FROM customers c JOIN orders o ON o.customer_id = c.customer_id GROUP BY c.customer_id, c.name, c.email, c.segment HAVING COUNT(o.order_id) = COUNT(CASE WHEN o.status IN ('cancelled', 'returned') THEN 1 END);"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Customer Churn Risk Report"}}, - {"step_id": "step_3", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Customer Churn Alert", "body": "Customers with potential churn risk identified.", "attach_step": "step_2"}}, - {"step_id": "step_4", "tool": "schedule_task", "params": {"task_name": "weekly_churn_alert", "workflow_steps": [{"tool": "query_database"}, {"tool": "generate_report"}, {"tool": "send_notification"}], "interval": "weekly", "notify_on_failure": True}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.name, c.email, c.segment, COUNT(o.order_id) AS total_orders FROM customers c JOIN orders o ON o.customer_id = c.customer_id GROUP BY c.customer_id, c.name, c.email, c.segment HAVING COUNT(o.order_id) = COUNT(CASE WHEN o.status IN ('cancelled', 'returned') THEN 1 END);" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Customer Churn Risk Report"}, + }, + { + "step_id": "step_3", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Customer Churn Alert", + "body": "Customers with potential churn risk identified.", + "attach_step": "step_2", + }, + }, + { + "step_id": "step_4", + "tool": "schedule_task", + "params": { + "task_name": "weekly_churn_alert", + "workflow_steps": [ + {"tool": "query_database"}, + {"tool": "generate_report"}, + {"tool": "send_notification"}, + ], + "interval": "weekly", + "notify_on_failure": True, + }, + }, ], }, { @@ -284,12 +718,57 @@ "task": "Build a cross-channel reporting pipeline: query top customers, enrich with CRM API, generate report, email and archive", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.name, c.segment, SUM(oi.quantity * oi.unit_price) AS total_spending FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.segment ORDER BY total_spending DESC LIMIT 10;"}}, - {"step_id": "step_2", "tool": "fetch_api_data", "params": {"url": "https://crm.company.com/api/customer-scores", "method": "GET"}}, - {"step_id": "step_3", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "join", "params": {"right_step": "step_2", "on": "name", "how": "left"}}}, - {"step_id": "step_4", "tool": "generate_report", "params": {"source_step": "step_3", "format": "markdown_table", "title": "Top Customer Spending Report"}}, - {"step_id": "step_5", "tool": "send_notification", "params": {"channel": "email", "recipient": "stakeholders@company.com", "subject": "Top Customer Report", "body": "Top customer spending report is ready.", "attach_step": "step_4"}}, - {"step_id": "step_6", "tool": "store_results", "params": {"source_step": "step_3", "target": "csv_file", "target_name": "top_customers_archive.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.name, c.segment, SUM(oi.quantity * oi.unit_price) AS total_spending FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.segment ORDER BY total_spending DESC LIMIT 10;" + }, + }, + { + "step_id": "step_2", + "tool": "fetch_api_data", + "params": {"url": "https://crm.company.com/api/customer-scores", "method": "GET"}, + }, + { + "step_id": "step_3", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "join", + "params": {"right_step": "step_2", "on": "name", "how": "left"}, + }, + }, + { + "step_id": "step_4", + "tool": "generate_report", + "params": { + "source_step": "step_3", + "format": "markdown_table", + "title": "Top Customer Spending Report", + }, + }, + { + "step_id": "step_5", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "stakeholders@company.com", + "subject": "Top Customer Report", + "body": "Top customer spending report is ready.", + "attach_step": "step_4", + }, + }, + { + "step_id": "step_6", + "tool": "store_results", + "params": { + "source_step": "step_3", + "target": "csv_file", + "target_name": "top_customers_archive.csv", + "mode": "append", + }, + }, ], }, { @@ -297,12 +776,59 @@ "task": "Create a full inventory monitoring pipeline: query unsold products, fetch supplier API, generate report, alert team, cache, and schedule weekly", "difficulty": "challenging", "gold_tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.name, p.category, p.price, p.stock_qty FROM products p LEFT JOIN order_items oi ON oi.product_id = p.product_id WHERE oi.item_id IS NULL ORDER BY p.name;"}}, - {"step_id": "step_2", "tool": "fetch_api_data", "params": {"url": "https://supplier.company.com/api/restock-status", "method": "GET"}}, - {"step_id": "step_3", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Unsold Products Inventory Report"}}, - {"step_id": "step_4", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Inventory Alert: Unsold Products", "body": "Products with zero orders have been identified.", "attach_step": "step_3"}}, - {"step_id": "step_5", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_unsold_products"}}, - {"step_id": "step_6", "tool": "schedule_task", "params": {"task_name": "weekly_inventory_monitor", "workflow_steps": [{"tool": "query_database"}, {"tool": "fetch_api_data"}, {"tool": "generate_report"}, {"tool": "send_notification"}, {"tool": "store_results"}], "interval": "weekly", "notify_on_failure": True}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.name, p.category, p.price, p.stock_qty FROM products p LEFT JOIN order_items oi ON oi.product_id = p.product_id WHERE oi.item_id IS NULL ORDER BY p.name;" + }, + }, + { + "step_id": "step_2", + "tool": "fetch_api_data", + "params": {"url": "https://supplier.company.com/api/restock-status", "method": "GET"}, + }, + { + "step_id": "step_3", + "tool": "generate_report", + "params": { + "source_step": "step_1", + "format": "markdown_table", + "title": "Unsold Products Inventory Report", + }, + }, + { + "step_id": "step_4", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Inventory Alert: Unsold Products", + "body": "Products with zero orders have been identified.", + "attach_step": "step_3", + }, + }, + { + "step_id": "step_5", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_unsold_products"}, + }, + { + "step_id": "step_6", + "tool": "schedule_task", + "params": { + "task_name": "weekly_inventory_monitor", + "workflow_steps": [ + {"tool": "query_database"}, + {"tool": "fetch_api_data"}, + {"tool": "generate_report"}, + {"tool": "send_notification"}, + {"tool": "store_results"}, + ], + "interval": "weekly", + "notify_on_failure": True, + }, + }, ], }, ] diff --git a/src/behavioral_memory/evaluation/metrics.py b/src/behavioral_memory/evaluation/metrics.py index 7c210fd..38db3b3 100644 --- a/src/behavioral_memory/evaluation/metrics.py +++ b/src/behavioral_memory/evaluation/metrics.py @@ -1,14 +1,15 @@ """Evaluation metrics from the paper (Section IV.C). - TSA — Tool Selection Accuracy: correct identification of required tools - PV — Parameter Validity: correct specification of key parameters - PCR — Plan Correctness Rate: correct tools AND >=80% parameter accuracy - ESA — Execution Sequence Accuracy: correct ordering of tool calls +TSA — Tool Selection Accuracy: correct identification of required tools +PV — Parameter Validity: correct specification of key parameters +PCR — Plan Correctness Rate: correct tools AND >=80% parameter accuracy +ESA — Execution Sequence Accuracy: correct ordering of tool calls """ from __future__ import annotations from collections import Counter +from typing import Any def tool_selection_accuracy(predicted_tools: list[str], gold_tools: list[str]) -> bool: @@ -16,9 +17,7 @@ def tool_selection_accuracy(predicted_tools: list[str], gold_tools: list[str]) - return Counter(predicted_tools) == Counter(gold_tools) -def parameter_validity( - predicted_params: list[dict], gold_params: list[dict] -) -> float: +def parameter_validity(predicted_params: list[dict[str, Any]], gold_params: list[dict[str, Any]]) -> float: """PV: fraction of key parameters correctly specified. Compares parameter keys step-by-step. If step counts differ, missing @@ -46,8 +45,8 @@ def parameter_validity( def plan_correctness( predicted_tools: list[str], gold_tools: list[str], - predicted_params: list[dict], - gold_params: list[dict], + predicted_params: list[dict[str, Any]], + gold_params: list[dict[str, Any]], pv_threshold: float = 0.8, ) -> bool: """PCR: correct tools AND parameter validity >= threshold.""" @@ -56,16 +55,12 @@ def plan_correctness( return tsa and pv >= pv_threshold -def execution_sequence_accuracy( - predicted_tools: list[str], gold_tools: list[str] -) -> bool: +def execution_sequence_accuracy(predicted_tools: list[str], gold_tools: list[str]) -> bool: """ESA: are tools in the correct order?""" return predicted_tools == gold_tools -def compute_metrics( - predicted_chain: list[dict], gold_chain: list[dict] -) -> dict[str, float | bool]: +def compute_metrics(predicted_chain: list[dict[str, Any]], gold_chain: list[dict[str, Any]]) -> dict[str, float | bool]: """Compute all four metrics for a single task.""" pred_tools = [s.get("tool", s.get("tool_name", "")) for s in predicted_chain] gold_tools = [s.get("tool", s.get("tool_name", "")) for s in gold_chain] diff --git a/src/behavioral_memory/evaluation/seed_traces.py b/src/behavioral_memory/evaluation/seed_traces.py index 21d1326..1522a59 100644 --- a/src/behavioral_memory/evaluation/seed_traces.py +++ b/src/behavioral_memory/evaluation/seed_traces.py @@ -21,10 +21,12 @@ from __future__ import annotations +from typing import Any + from behavioral_memory.core.schemas import ExecutionTrace, ToolCall -def _build_trace(task: str, chain: list[dict], explanation: str) -> ExecutionTrace: +def _build_trace(task: str, chain: list[dict[str, Any]], explanation: str) -> ExecutionTrace: """Helper to convert raw dict data into an ExecutionTrace.""" tool_chain = [ ToolCall( @@ -43,105 +45,295 @@ def _build_trace(task: str, chain: list[dict], explanation: str) -> ExecutionTra ) -SEED_TRACES_RAW: list[dict] = [ +SEED_TRACES_RAW: list[dict[str, Any]] = [ { "task": "Get Q1 revenue data and send a report to stakeholders", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2025-04-01';"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Q1 Revenue Report"}}, - {"step_id": "step_3", "tool": "send_notification", "params": {"channel": "email", "recipient": "stakeholders@company.com", "subject": "Q1 Revenue Report", "body": "Q1 revenue report is attached.", "attach_step": "step_2"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2025-04-01';" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Q1 Revenue Report"}, + }, + { + "step_id": "step_3", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "stakeholders@company.com", + "subject": "Q1 Revenue Report", + "body": "Q1 revenue report is attached.", + "attach_step": "step_2", + }, + }, ], "explanation": "Revenue is always from order_items (quantity * unit_price), not total_amount. Reports use markdown_table and are delivered via email.", }, { "task": "Get completed order statistics and cache for the dashboard", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT COUNT(*) AS completed FROM orders WHERE status IN ('shipped', 'delivered');"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_completed_stats"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT COUNT(*) AS completed FROM orders WHERE status IN ('shipped', 'delivered');" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": {"source_step": "step_1", "target": "cache", "target_name": "dashboard_completed_stats"}, + }, ], "explanation": "'Completed' orders = status IN ('shipped', 'delivered'). Dashboard data goes to cache.", }, { "task": "Archive the current valid order data as CSV", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT * FROM orders WHERE status NOT IN ('cancelled', 'returned') ORDER BY order_date;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "csv_file", "target_name": "valid_orders.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT * FROM orders WHERE status NOT IN ('cancelled', 'returned') ORDER BY order_date;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": { + "source_step": "step_1", + "target": "csv_file", + "target_name": "valid_orders.csv", + "mode": "append", + }, + }, ], "explanation": "'Valid orders' excludes BOTH cancelled AND returned. Archiving uses csv_file with append.", }, { "task": "Alert the team about the current fulfillment metrics", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT ROUND(COUNT(*) FILTER (WHERE status = 'delivered')::numeric / NULLIF(COUNT(*) FILTER (WHERE status != 'cancelled'), 0) * 100, 1) AS fulfillment_rate_pct FROM orders;"}}, - {"step_id": "step_2", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Fulfillment Metrics", "body": "Current fulfillment metrics are ready.", "attach_step": "step_1"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT ROUND(COUNT(*) FILTER (WHERE status = 'delivered')::numeric / NULLIF(COUNT(*) FILTER (WHERE status != 'cancelled'), 0) * 100, 1) AS fulfillment_rate_pct FROM orders;" + }, + }, + { + "step_id": "step_2", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Fulfillment Metrics", + "body": "Current fulfillment metrics are ready.", + "attach_step": "step_1", + }, + }, ], "explanation": "Fulfillment rate = delivered / (total - cancelled). Alerts use Slack #data-alerts.", }, { "task": "Build a pipeline for net order value analysis and store for dashboard", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT o.order_id, (o.total_amount - o.discount) AS net_value FROM orders o WHERE o.status NOT IN ('cancelled', 'returned');"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "aggregate", "params": {"metrics": {"net_value": "avg"}}}}, - {"step_id": "step_3", "tool": "store_results", "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_net_value"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT o.order_id, (o.total_amount - o.discount) AS net_value FROM orders o WHERE o.status NOT IN ('cancelled', 'returned');" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "aggregate", + "params": {"metrics": {"net_value": "avg"}}, + }, + }, + { + "step_id": "step_3", + "tool": "store_results", + "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_net_value"}, + }, ], "explanation": "Net order value = total_amount - discount. Pipelines include transform. Dashboard -> cache.", }, { "task": "Get product category performance and archive the data", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;"}}, - {"step_id": "step_2", "tool": "store_results", "params": {"source_step": "step_1", "target": "csv_file", "target_name": "category_performance.csv", "mode": "append"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.category, SUM(oi.quantity * oi.unit_price) AS revenue FROM order_items oi JOIN products p ON p.product_id = oi.product_id JOIN orders o ON o.order_id = oi.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY p.category ORDER BY revenue DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "store_results", + "params": { + "source_step": "step_1", + "target": "csv_file", + "target_name": "category_performance.csv", + "mode": "append", + }, + }, ], "explanation": "Product performance uses revenue from order_items. Archiving uses csv_file with append.", }, { "task": "Schedule a daily report of active customer counts", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT COUNT(DISTINCT c.customer_id) AS active_customers FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned');"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Active Customer Count"}}, - {"step_id": "step_3", "tool": "schedule_task", "params": {"task_name": "daily_active_customers", "workflow_steps": [{"tool": "query_database"}, {"tool": "generate_report"}], "interval": "daily", "notify_on_failure": True}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT COUNT(DISTINCT c.customer_id) AS active_customers FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned');" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Active Customer Count"}, + }, + { + "step_id": "step_3", + "tool": "schedule_task", + "params": { + "task_name": "daily_active_customers", + "workflow_steps": [{"tool": "query_database"}, {"tool": "generate_report"}], + "interval": "daily", + "notify_on_failure": True, + }, + }, ], "explanation": "Active = non-cancelled, non-returned orders. Scheduled reports use daily + notify_on_failure=true.", }, { "task": "Get repeat customer data and generate a report about it", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.name, c.email, COUNT(o.order_id) AS order_count FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.email HAVING COUNT(o.order_id) > 1;"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Repeat Customers Report"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.name, c.email, COUNT(o.order_id) AS order_count FROM customers c JOIN orders o ON o.customer_id = c.customer_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.email HAVING COUNT(o.order_id) > 1;" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Repeat Customers Report"}, + }, ], "explanation": "Repeat customers have > 1 valid order. Reports use markdown_table.", }, { "task": "Get basket sizes per order and alert about unusual patterns", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT o.order_id, SUM(oi.quantity) AS item_count FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id;"}}, - {"step_id": "step_2", "tool": "send_notification", "params": {"channel": "slack", "recipient": "#data-alerts", "subject": "Basket Size Analysis", "body": "Basket size data is ready for review.", "attach_step": "step_1"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT o.order_id, SUM(oi.quantity) AS item_count FROM orders o JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY o.order_id;" + }, + }, + { + "step_id": "step_2", + "tool": "send_notification", + "params": { + "channel": "slack", + "recipient": "#data-alerts", + "subject": "Basket Size Analysis", + "body": "Basket size data is ready for review.", + "attach_step": "step_1", + }, + }, ], "explanation": "Basket size = number of items (quantity), not dollar amount. Alerts -> Slack #data-alerts.", }, { "task": "Find products with no sales and generate a report", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT p.name, p.category, p.price FROM products p LEFT JOIN order_items oi ON oi.product_id = p.product_id WHERE oi.item_id IS NULL ORDER BY p.name;"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Unsold Products Report"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT p.name, p.category, p.price FROM products p LEFT JOIN order_items oi ON oi.product_id = p.product_id WHERE oi.item_id IS NULL ORDER BY p.name;" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Unsold Products Report"}, + }, ], "explanation": "Use LEFT JOIN + IS NULL for 'never sold'. Reports use markdown_table.", }, { "task": "Customer ranking pipeline: get spending data, transform, and cache", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT c.name, c.segment, SUM(oi.quantity * oi.unit_price) AS total_spending FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.segment ORDER BY total_spending DESC;"}}, - {"step_id": "step_2", "tool": "transform_data", "params": {"source_step": "step_1", "operation": "compute", "params": {"new_column": "spending_rank", "expression": "RANK() OVER (ORDER BY total_spending DESC)"}}}, - {"step_id": "step_3", "tool": "store_results", "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_customer_ranking"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT c.name, c.segment, SUM(oi.quantity * oi.unit_price) AS total_spending FROM customers c JOIN orders o ON o.customer_id = c.customer_id JOIN order_items oi ON oi.order_id = o.order_id WHERE o.status NOT IN ('cancelled', 'returned') GROUP BY c.customer_id, c.name, c.segment ORDER BY total_spending DESC;" + }, + }, + { + "step_id": "step_2", + "tool": "transform_data", + "params": { + "source_step": "step_1", + "operation": "compute", + "params": { + "new_column": "spending_rank", + "expression": "RANK() OVER (ORDER BY total_spending DESC)", + }, + }, + }, + { + "step_id": "step_3", + "tool": "store_results", + "params": {"source_step": "step_2", "target": "cache", "target_name": "dashboard_customer_ranking"}, + }, ], "explanation": "Customer spending uses order_items revenue. Pipelines include transform. Dashboard -> cache.", }, { "task": "Monthly valid orders trend and email the report", "tool_chain": [ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, COUNT(*) AS valid_orders FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;"}}, - {"step_id": "step_2", "tool": "generate_report", "params": {"source_step": "step_1", "format": "markdown_table", "title": "Monthly Valid Orders"}}, - {"step_id": "step_3", "tool": "send_notification", "params": {"channel": "email", "recipient": "manager@company.com", "subject": "Monthly Valid Orders Report", "body": "Valid orders trend report is attached.", "attach_step": "step_2"}}, + { + "step_id": "step_1", + "tool": "query_database", + "params": { + "query": "SELECT DATE_TRUNC('month', o.order_date)::date AS month, COUNT(*) AS valid_orders FROM orders o WHERE o.status NOT IN ('cancelled', 'returned') AND o.order_date >= '2025-01-01' AND o.order_date < '2026-01-01' GROUP BY month ORDER BY month;" + }, + }, + { + "step_id": "step_2", + "tool": "generate_report", + "params": {"source_step": "step_1", "format": "markdown_table", "title": "Monthly Valid Orders"}, + }, + { + "step_id": "step_3", + "tool": "send_notification", + "params": { + "channel": "email", + "recipient": "manager@company.com", + "subject": "Monthly Valid Orders Report", + "body": "Valid orders trend report is attached.", + "attach_step": "step_2", + }, + }, ], "explanation": "Valid orders exclude cancelled AND returned. Reports go via email.", }, @@ -150,7 +342,4 @@ def _build_trace(task: str, chain: list[dict], explanation: str) -> ExecutionTra def get_seed_traces() -> list[ExecutionTrace]: """Return all 12 seed traces as ExecutionTrace objects.""" - return [ - _build_trace(raw["task"], raw["tool_chain"], raw["explanation"]) - for raw in SEED_TRACES_RAW - ] + return [_build_trace(raw["task"], raw["tool_chain"], raw["explanation"]) for raw in SEED_TRACES_RAW] diff --git a/src/behavioral_memory/evaluation/statistics.py b/src/behavioral_memory/evaluation/statistics.py index 1b6c1d6..7f9d1f7 100644 --- a/src/behavioral_memory/evaluation/statistics.py +++ b/src/behavioral_memory/evaluation/statistics.py @@ -39,9 +39,7 @@ def bootstrap_ci( return point_estimate, means[lower_idx], means[upper_idx] -def mcnemar_test( - results_a: list[bool], results_b: list[bool] -) -> dict[str, Any]: +def mcnemar_test(results_a: list[bool], results_b: list[bool]) -> dict[str, Any]: """McNemar's exact test for paired binary outcomes. Compares two methods on the same set of tasks. diff --git a/src/behavioral_memory/evaluation/strategies.py b/src/behavioral_memory/evaluation/strategies.py index ece0ce1..df46c6b 100644 --- a/src/behavioral_memory/evaluation/strategies.py +++ b/src/behavioral_memory/evaluation/strategies.py @@ -22,16 +22,12 @@ def generate(self, query: str, tool_schemas: list[ToolSchema]) -> Plan: class StaticFewShotStrategy: """Baseline: a fixed set of 3 examples is included for all tasks.""" - def __init__( - self, engine: PlanEngine, static_traces: list[ExecutionTrace] - ) -> None: + def __init__(self, engine: PlanEngine, static_traces: list[ExecutionTrace]) -> None: self._engine = engine self._static = static_traces[:3] def generate(self, query: str, tool_schemas: list[ToolSchema]) -> Plan: - return self._engine.generate_static_few_shot( - query, tool_schemas, self._static - ) + return self._engine.generate_static_few_shot(query, tool_schemas, self._static) class DynamicRetrievalStrategy: diff --git a/src/behavioral_memory/gatekeeper/pipeline.py b/src/behavioral_memory/gatekeeper/pipeline.py index aebcd63..0a68392 100644 --- a/src/behavioral_memory/gatekeeper/pipeline.py +++ b/src/behavioral_memory/gatekeeper/pipeline.py @@ -37,9 +37,7 @@ def __init__( self._settings = settings or Settings() self._schema_validator = SchemaValidator(registry) self._sandbox = SandboxExecutor(settings=self._settings) - self._dedup_gate = DeduplicationGate( - Deduplicator(store=store, settings=self._settings) - ) + self._dedup_gate = DeduplicationGate(Deduplicator(store=store, settings=self._settings)) def evaluate(self, trace: ExecutionTrace) -> GatekeeperResult: """Run a trace through all three gates and return the result. diff --git a/src/behavioral_memory/gatekeeper/sandbox.py b/src/behavioral_memory/gatekeeper/sandbox.py index b45f281..8944417 100644 --- a/src/behavioral_memory/gatekeeper/sandbox.py +++ b/src/behavioral_memory/gatekeeper/sandbox.py @@ -76,8 +76,7 @@ def _dry_run(self, trace: ExecutionTrace) -> tuple[bool, str]: if right_step and right_step not in available_outputs: return ( False, - f"Step '{step.step_id}' join references '{right_step}' " - f"which has no output yet", + f"Step '{step.step_id}' join references '{right_step}' which has no output yet", ) available_outputs.add(step.step_id) diff --git a/src/behavioral_memory/gatekeeper/schema_validator.py b/src/behavioral_memory/gatekeeper/schema_validator.py index 9976ca0..e83d949 100644 --- a/src/behavioral_memory/gatekeeper/schema_validator.py +++ b/src/behavioral_memory/gatekeeper/schema_validator.py @@ -58,13 +58,9 @@ def _check_dependencies(self, trace: ExecutionTrace) -> list[str]: for step in trace.tool_chain: for dep in step.depends_on: if dep not in all_ids: - failures.append( - f"Step '{step.step_id}' depends on non-existent step '{dep}'" - ) + failures.append(f"Step '{step.step_id}' depends on non-existent step '{dep}'") elif id_order.get(dep, -1) >= id_order.get(step.step_id, 0): - failures.append( - f"Step '{step.step_id}' depends on '{dep}' which comes after it" - ) + failures.append(f"Step '{step.step_id}' depends on '{dep}' which comes after it") return failures def _check_required_params(self, trace: ExecutionTrace) -> list[str]: @@ -77,12 +73,9 @@ def _check_required_params(self, trace: ExecutionTrace) -> list[str]: return failures @staticmethod - def _validate_step_params( - step: ToolCall, schema: ToolSchema, failures: list[str] - ) -> None: + def _validate_step_params(step: ToolCall, schema: ToolSchema, failures: list[str]) -> None: for req_param in schema.required_params: if req_param not in step.parameters: failures.append( - f"Step '{step.step_id}': missing required param " - f"'{req_param}' for tool '{step.tool_name}'" + f"Step '{step.step_id}': missing required param '{req_param}' for tool '{step.tool_name}'" ) diff --git a/src/behavioral_memory/memory/__init__.py b/src/behavioral_memory/memory/__init__.py index c9e48c8..23ad77f 100644 --- a/src/behavioral_memory/memory/__init__.py +++ b/src/behavioral_memory/memory/__init__.py @@ -1,6 +1,14 @@ from behavioral_memory.memory.dedup import Deduplicator from behavioral_memory.memory.in_memory_store import InMemoryTraceStore -from behavioral_memory.memory.store import TraceStore from behavioral_memory.memory.token_budget import select_traces_within_budget + +def __getattr__(name: str): # type: ignore[no-untyped-def] + if name == "TraceStore": + from behavioral_memory.memory.store import TraceStore + + return TraceStore + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = ["Deduplicator", "InMemoryTraceStore", "TraceStore", "select_traces_within_budget"] diff --git a/src/behavioral_memory/memory/in_memory_store.py b/src/behavioral_memory/memory/in_memory_store.py index 3c9dc82..4a6129b 100644 --- a/src/behavioral_memory/memory/in_memory_store.py +++ b/src/behavioral_memory/memory/in_memory_store.py @@ -13,6 +13,7 @@ from __future__ import annotations import logging +from typing import Any import numpy as np from langchain_core.embeddings import Embeddings @@ -39,9 +40,7 @@ def __init__( self._traces: list[ExecutionTrace] = [] self._vectors: list[list[float]] = [] - def search( - self, query: str, k: int | None = None - ) -> list[tuple[ExecutionTrace, float]]: + def search(self, query: str, k: int | None = None) -> list[tuple[ExecutionTrace, float]]: k = k or self._settings.few_shot_k if not self._traces: return [] @@ -79,9 +78,7 @@ def count(self) -> int: return len(self._traces) @staticmethod - def _cosine_similarities( - query_vec: list[float], doc_vecs: list[list[float]] - ) -> np.ndarray: + def _cosine_similarities(query_vec: list[float], doc_vecs: list[list[float]]) -> Any: q = np.array(query_vec) d = np.array(doc_vecs) q_norm = q / (np.linalg.norm(q) + 1e-10) diff --git a/src/behavioral_memory/memory/store.py b/src/behavioral_memory/memory/store.py index 8f0a1a9..7cd98a5 100644 --- a/src/behavioral_memory/memory/store.py +++ b/src/behavioral_memory/memory/store.py @@ -3,16 +3,18 @@ Stores validated execution traces as vector-embedded documents and retrieves the most semantically similar traces for a given query. Corresponds to Section III.B of the paper. + +Requires: pip install behavioral-memory[postgres] """ from __future__ import annotations import json import logging +from typing import Any from langchain_core.documents import Document from langchain_core.embeddings import Embeddings -from langchain_postgres import PGVector from behavioral_memory.core.config import Settings from behavioral_memory.core.exceptions import MemoryStoreError @@ -20,6 +22,21 @@ logger = logging.getLogger(__name__) +try: + from langchain_postgres import PGVector + + _HAS_PGVECTOR = True +except ImportError: + _HAS_PGVECTOR = False + PGVector = None # type: ignore[assignment, misc] + + +def _check_postgres_deps() -> None: + if not _HAS_PGVECTOR: + raise ImportError( + "TraceStore requires PostgreSQL dependencies. Install them with: pip install behavioral-memory[postgres]" + ) + class TraceStore: """Vector store for validated execution traces. @@ -27,6 +44,8 @@ class TraceStore: Accepts any LangChain-compatible Embeddings model, making the framework model-agnostic. The caller decides which embedding provider to use (Gemini, OpenAI, local, etc.). + + Requires: pip install behavioral-memory[postgres] """ def __init__( @@ -36,14 +55,15 @@ def __init__( collection_name: str | None = None, settings: Settings | None = None, ) -> None: + _check_postgres_deps() self._settings = settings or Settings() self._connection_url = connection_url or self._settings.vector_store_url self._collection_name = collection_name or self._settings.vector_store_collection self._embeddings = embeddings - self._vectorstore: PGVector | None = None + self._vectorstore: Any = None @property - def vectorstore(self) -> PGVector: + def vectorstore(self) -> Any: if self._vectorstore is None: try: self._vectorstore = PGVector( @@ -56,14 +76,8 @@ def vectorstore(self) -> PGVector: raise MemoryStoreError(f"Failed to connect to vector store: {e}") from e return self._vectorstore - def search( - self, query: str, k: int | None = None - ) -> list[tuple[ExecutionTrace, float]]: - """Retrieve the top-k most similar traces for a query. - - Returns a list of (trace, similarity_score) tuples sorted by - descending similarity. - """ + def search(self, query: str, k: int | None = None) -> list[tuple[ExecutionTrace, float]]: + """Retrieve the top-k most similar traces for a query.""" k = k or self._settings.few_shot_k try: results = self.vectorstore.similarity_search_with_score(query, k=k) @@ -112,15 +126,9 @@ def count(self) -> int: except Exception: return 0 - # -- Serialization helpers -- - @staticmethod def _trace_to_doc(trace: ExecutionTrace) -> Document: - """Serialize an ExecutionTrace into a LangChain Document. - - The page_content is the task description (what gets embedded). - The metadata carries the full structured trace as JSON. - """ + """Serialize an ExecutionTrace into a LangChain Document.""" tool_chain_data = [step.model_dump() for step in trace.tool_chain] return Document( page_content=trace.task_description, diff --git a/src/behavioral_memory/observability/annotation.py b/src/behavioral_memory/observability/annotation.py index fa9cbc4..497baa8 100644 --- a/src/behavioral_memory/observability/annotation.py +++ b/src/behavioral_memory/observability/annotation.py @@ -40,9 +40,7 @@ def __init__( self._poller = poller self._gatekeeper = gatekeeper - def process_feedback( - self, traces: list[ExecutionTrace] | None = None - ) -> FeedbackStats: + def process_feedback(self, traces: list[ExecutionTrace] | None = None) -> FeedbackStats: """Process a batch of traces through the gatekeeper. If traces are not provided, polls Langfuse for positive ones. @@ -57,25 +55,16 @@ def process_feedback( result = self._gatekeeper.submit(trace) if result.accepted: stats.accepted += 1 - stats.details.append( - f"Accepted: {trace.task_description[:60]}" - ) + stats.details.append(f"Accepted: {trace.task_description[:60]}") elif result.is_duplicate: stats.rejected_duplicate += 1 - stats.details.append( - f"Duplicate: {trace.task_description[:60]}" - ) + stats.details.append(f"Duplicate: {trace.task_description[:60]}") elif not result.schema_valid: stats.rejected_validation += 1 - stats.details.append( - f"Invalid: {trace.task_description[:60]} — " - f"{', '.join(result.failures)}" - ) + stats.details.append(f"Invalid: {trace.task_description[:60]} — {', '.join(result.failures)}") elif not result.sandbox_passed: stats.rejected_sandbox += 1 - stats.details.append( - f"Sandbox fail: {trace.task_description[:60]}" - ) + stats.details.append(f"Sandbox fail: {trace.task_description[:60]}") except Exception as e: stats.errors += 1 logger.warning("Error processing trace: %s", e) diff --git a/src/behavioral_memory/observability/feedback.py b/src/behavioral_memory/observability/feedback.py index 951361f..8f200dd 100644 --- a/src/behavioral_memory/observability/feedback.py +++ b/src/behavioral_memory/observability/feedback.py @@ -100,9 +100,9 @@ def _trace_to_execution_trace(self, trace: Any) -> ExecutionTrace | None: tool_chain = [ ToolCall( - step_id=s.get("step_id", f"step_{i+1}"), - tool_name=s.get("tool_name", s.get("tool", "")), - parameters=s.get("parameters", s.get("params", {})), + step_id=s.get("step_id", f"step_{i + 1}"), + tool_name=str(s.get("tool_name", s.get("tool", ""))), + parameters=dict(s.get("parameters", s.get("params", {}))), # type: ignore[arg-type] depends_on=s.get("depends_on", []), ) for i, s in enumerate(steps_data) @@ -126,9 +126,7 @@ def poll_once(self) -> list[ExecutionTrace]: """Single poll cycle.""" return self.fetch_positive_traces() - def poll_loop( - self, callback: Any = None, max_iterations: int | None = None - ) -> None: + def poll_loop(self, callback: Any = None, max_iterations: int | None = None) -> None: """Continuous polling loop. Calls callback(trace) for each positive trace found. diff --git a/src/behavioral_memory/observability/tracer.py b/src/behavioral_memory/observability/tracer.py index 2e1d29c..0245974 100644 --- a/src/behavioral_memory/observability/tracer.py +++ b/src/behavioral_memory/observability/tracer.py @@ -84,9 +84,7 @@ def log_plan( input=plan.query, output=plan.raw_llm_output, metadata={ - "retrieved_examples": [ - t.task_description for t in plan.retrieved_traces - ], + "retrieved_examples": [t.task_description for t in plan.retrieved_traces], }, ) diff --git a/src/behavioral_memory/planner/engine.py b/src/behavioral_memory/planner/engine.py index 0c1b52d..93ae9d2 100644 --- a/src/behavioral_memory/planner/engine.py +++ b/src/behavioral_memory/planner/engine.py @@ -115,6 +115,4 @@ def generate_static_few_shot( static_traces: list[ExecutionTrace], ) -> Plan: """Generate a plan with fixed static examples (baseline).""" - return self.generate( - query=query, tool_schemas=tool_schemas, traces=static_traces - ) + return self.generate(query=query, tool_schemas=tool_schemas, traces=static_traces) diff --git a/src/behavioral_memory/planner/postprocess.py b/src/behavioral_memory/planner/postprocess.py index 3860819..096a022 100644 --- a/src/behavioral_memory/planner/postprocess.py +++ b/src/behavioral_memory/planner/postprocess.py @@ -9,6 +9,7 @@ import json import logging import re +from typing import Any from behavioral_memory.core.exceptions import PlanGenerationError from behavioral_memory.core.schemas import ToolCall @@ -16,7 +17,7 @@ logger = logging.getLogger(__name__) -def extract_json_array(raw: str) -> list[dict]: +def extract_json_array(raw: str) -> list[dict[str, Any]]: """Extract a JSON array from raw LLM output.""" text = raw.strip() @@ -42,7 +43,7 @@ def extract_json_array(raw: str) -> list[dict]: return parsed -def parse_tool_calls(raw_steps: list[dict]) -> list[ToolCall]: +def parse_tool_calls(raw_steps: list[dict[str, Any]]) -> list[ToolCall]: """Convert raw JSON steps into validated ToolCall models.""" calls: list[ToolCall] = [] for i, step in enumerate(raw_steps): diff --git a/src/behavioral_memory/tools/mcp_client.py b/src/behavioral_memory/tools/mcp_client.py index 19bcec2..ce53938 100644 --- a/src/behavioral_memory/tools/mcp_client.py +++ b/src/behavioral_memory/tools/mcp_client.py @@ -27,9 +27,7 @@ async def fetch_mcp_schemas(server_url: str) -> list[ToolSchema]: from mcp import ClientSession from mcp.client.sse import sse_client except ImportError as e: - raise MCPConnectionError( - "MCP SDK not installed. Install with: pip install mcp" - ) from e + raise MCPConnectionError("MCP SDK not installed. Install with: pip install mcp") from e schemas: list[ToolSchema] = [] try: @@ -59,9 +57,7 @@ async def fetch_mcp_schemas(server_url: str) -> list[ToolSchema]: return schemas -def load_schemas_into_registry( - schemas: list[ToolSchema], registry: ToolRegistry | None = None -) -> ToolRegistry: +def load_schemas_into_registry(schemas: list[ToolSchema], registry: ToolRegistry | None = None) -> ToolRegistry: """Load a list of ToolSchemas into a ToolRegistry.""" reg = registry or ToolRegistry() reg.register_many(schemas) diff --git a/src/behavioral_memory/tools/mock_tools.py b/src/behavioral_memory/tools/mock_tools.py index 10e137b..b53c35b 100644 --- a/src/behavioral_memory/tools/mock_tools.py +++ b/src/behavioral_memory/tools/mock_tools.py @@ -15,9 +15,11 @@ from __future__ import annotations +from typing import Any + from behavioral_memory.core.schemas import ToolSchema -TOOL_DEFINITIONS: list[dict] = [ +TOOL_DEFINITIONS: list[dict[str, Any]] = [ { "name": "query_database", "description": ( diff --git a/tests/e2e/test_full_pipeline.py b/tests/e2e/test_full_pipeline.py index c2cb10c..71e1af7 100644 --- a/tests/e2e/test_full_pipeline.py +++ b/tests/e2e/test_full_pipeline.py @@ -106,9 +106,7 @@ def test_difficulty_distribution(self): def test_all_gold_chains_use_known_tools(self, registry): for task in EVALUATION_TASKS: for step in task["gold_tool_chain"]: - assert registry.has_tool(step["tool"]), ( - f"Task {task['task_id']}: unknown tool '{step['tool']}'" - ) + assert registry.has_tool(step["tool"]), f"Task {task['task_id']}: unknown tool '{step['tool']}'" def test_gold_chain_step_ids_are_unique(self): for task in EVALUATION_TASKS: @@ -185,10 +183,22 @@ class TestPostprocessorEndToEnd: """LLM response parsing handles realistic outputs.""" def test_parses_perfect_response(self): - raw = json.dumps([ - {"step_id": "step_1", "tool_name": "query_database", "parameters": {"query": "SELECT 1"}, "depends_on": []}, - {"step_id": "step_2", "tool_name": "generate_report", "parameters": {"source_step": "step_1", "format": "csv", "title": "t"}, "depends_on": ["step_1"]}, - ]) + raw = json.dumps( + [ + { + "step_id": "step_1", + "tool_name": "query_database", + "parameters": {"query": "SELECT 1"}, + "depends_on": [], + }, + { + "step_id": "step_2", + "tool_name": "generate_report", + "parameters": {"source_step": "step_1", "format": "csv", "title": "t"}, + "depends_on": ["step_1"], + }, + ] + ) steps = postprocess_plan(raw) assert len(steps) == 2 assert steps[0].tool_name == "query_database" @@ -212,9 +222,11 @@ def test_parses_response_with_trailing_comma(self): assert len(steps) == 1 def test_parses_alternative_key_names(self): - raw = json.dumps([ - {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT 1"}}, - ]) + raw = json.dumps( + [ + {"step_id": "step_1", "tool": "query_database", "params": {"query": "SELECT 1"}}, + ] + ) steps = postprocess_plan(raw) assert steps[0].tool_name == "query_database" @@ -270,8 +282,14 @@ def test_valid_trace_accepted(self, registry): trace = ExecutionTrace( task_description="Get revenue and generate a report", tool_chain=[ - ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT SUM(amount) FROM orders"}), - ToolCall(step_id="s2", tool_name="generate_report", parameters={"source_step": "s1", "format": "markdown_table", "title": "Revenue"}), + ToolCall( + step_id="s1", tool_name="query_database", parameters={"query": "SELECT SUM(amount) FROM orders"} + ), + ToolCall( + step_id="s2", + tool_name="generate_report", + parameters={"source_step": "s1", "format": "markdown_table", "title": "Revenue"}, + ), ], source="execution", ) @@ -296,7 +314,11 @@ def test_broken_data_flow_rejected_at_gate2(self): task_description="test", tool_chain=[ ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT 1"}), - ToolCall(step_id="s2", tool_name="generate_report", parameters={"source_step": "s99", "format": "csv", "title": "t"}), + ToolCall( + step_id="s2", + tool_name="generate_report", + parameters={"source_step": "s99", "format": "csv", "title": "t"}, + ), ], ) passed, msg = sandbox.execute(trace) @@ -311,9 +333,16 @@ class TestPlanEngineEndToEnd: """PlanEngine orchestration with a mocked LLM and mocked store.""" def test_generate_produces_plan(self): - expected_output = json.dumps([ - {"step_id": "step_1", "tool_name": "query_database", "parameters": {"query": "SELECT COUNT(*) FROM customers"}, "depends_on": []}, - ]) + expected_output = json.dumps( + [ + { + "step_id": "step_1", + "tool_name": "query_database", + "parameters": {"query": "SELECT COUNT(*) FROM customers"}, + "depends_on": [], + }, + ] + ) mock_response = MagicMock() mock_response.content = expected_output @@ -346,9 +375,16 @@ def test_generate_produces_plan(self): def test_generate_with_traces_passes_them_to_prompt(self): seed_traces = get_seed_traces()[:2] - expected_output = json.dumps([ - {"step_id": "step_1", "tool_name": "query_database", "parameters": {"query": "SELECT 1"}, "depends_on": []}, - ]) + expected_output = json.dumps( + [ + { + "step_id": "step_1", + "tool_name": "query_database", + "parameters": {"query": "SELECT 1"}, + "depends_on": [], + }, + ] + ) mock_response = MagicMock() mock_response.content = expected_output @@ -371,9 +407,16 @@ def test_generate_with_traces_passes_them_to_prompt(self): assert "REFERENCE EXAMPLES" in prompt_text def test_generate_zero_shot(self): - expected_output = json.dumps([ - {"step_id": "step_1", "tool_name": "query_database", "parameters": {"query": "SELECT 1"}, "depends_on": []}, - ]) + expected_output = json.dumps( + [ + { + "step_id": "step_1", + "tool_name": "query_database", + "parameters": {"query": "SELECT 1"}, + "depends_on": [], + }, + ] + ) mock_response = MagicMock() mock_response.content = expected_output mock_llm = MagicMock() @@ -466,15 +509,17 @@ def test_task_0_round_trip(self): assert "customer" in prompt.lower() gold_chain = task["gold_tool_chain"] - fake_llm_output = json.dumps([ - { - "step_id": s["step_id"], - "tool_name": s["tool"], - "parameters": s.get("params", {}), - "depends_on": [], - } - for s in gold_chain - ]) + fake_llm_output = json.dumps( + [ + { + "step_id": s["step_id"], + "tool_name": s["tool"], + "parameters": s.get("params", {}), + "depends_on": [], + } + for s in gold_chain + ] + ) steps = postprocess_plan(fake_llm_output) assert len(steps) == len(gold_chain) @@ -492,10 +537,7 @@ def test_task_0_round_trip(self): passed, msg = sandbox.execute(trace) assert passed, msg - predicted_chain = [ - {"tool": s.tool_name, "params": s.parameters} - for s in steps - ] + predicted_chain = [{"tool": s.tool_name, "params": s.parameters} for s in steps] metrics = compute_metrics(predicted_chain, gold_chain) assert metrics["tsa"] is True assert metrics["esa"] is True @@ -510,10 +552,17 @@ def test_all_simple_tasks_round_trip(self): simple_tasks = [t for t in EVALUATION_TASKS if t["difficulty"] == "simple"] for task in simple_tasks: gold = task["gold_tool_chain"] - fake_output = json.dumps([ - {"step_id": s["step_id"], "tool_name": s["tool"], "parameters": s.get("params", {}), "depends_on": []} - for s in gold - ]) + fake_output = json.dumps( + [ + { + "step_id": s["step_id"], + "tool_name": s["tool"], + "parameters": s.get("params", {}), + "depends_on": [], + } + for s in gold + ] + ) steps = postprocess_plan(fake_output) trace = ExecutionTrace( task_description=task["task"], @@ -534,10 +583,17 @@ def test_all_challenging_tasks_round_trip(self): challenging = [t for t in EVALUATION_TASKS if t["difficulty"] == "challenging"] for task in challenging: gold = task["gold_tool_chain"] - fake_output = json.dumps([ - {"step_id": s["step_id"], "tool_name": s["tool"], "parameters": s.get("params", {}), "depends_on": []} - for s in gold - ]) + fake_output = json.dumps( + [ + { + "step_id": s["step_id"], + "tool_name": s["tool"], + "parameters": s.get("params", {}), + "depends_on": [], + } + for s in gold + ] + ) steps = postprocess_plan(fake_output) trace = ExecutionTrace( task_description=task["task"], @@ -586,8 +642,16 @@ def test_round_trip_preserves_data(self): trace = ExecutionTrace( task_description="Get revenue by category", tool_chain=[ - ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT category, SUM(revenue) FROM products GROUP BY category"}), - ToolCall(step_id="s2", tool_name="generate_report", parameters={"source_step": "s1", "format": "markdown_table", "title": "Revenue Report"}), + ToolCall( + step_id="s1", + tool_name="query_database", + parameters={"query": "SELECT category, SUM(revenue) FROM products GROUP BY category"}, + ), + ToolCall( + step_id="s2", + tool_name="generate_report", + parameters={"source_step": "s1", "format": "markdown_table", "title": "Revenue Report"}, + ), ], validated=True, source="seed", diff --git a/tests/integration/test_gatekeeper.py b/tests/integration/test_gatekeeper.py index 2a4ef49..55c32ff 100644 --- a/tests/integration/test_gatekeeper.py +++ b/tests/integration/test_gatekeeper.py @@ -16,10 +16,36 @@ def test_validates_complete_pipeline(self, benchmark_registry): task_description="Revenue pipeline", tool_chain=[ ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT 1"}), - ToolCall(step_id="s2", tool_name="transform_data", parameters={"source_step": "s1", "operation": "compute", "params": {"new_column": "x", "expression": "1+1"}}), - ToolCall(step_id="s3", tool_name="generate_report", parameters={"source_step": "s2", "format": "markdown_table", "title": "Report"}), - ToolCall(step_id="s4", tool_name="send_notification", parameters={"channel": "email", "recipient": "a@b.com", "subject": "Report", "body": "Done", "attach_step": "s3"}), - ToolCall(step_id="s5", tool_name="store_results", parameters={"source_step": "s2", "target": "csv_file", "target_name": "out.csv"}), + ToolCall( + step_id="s2", + tool_name="transform_data", + parameters={ + "source_step": "s1", + "operation": "compute", + "params": {"new_column": "x", "expression": "1+1"}, + }, + ), + ToolCall( + step_id="s3", + tool_name="generate_report", + parameters={"source_step": "s2", "format": "markdown_table", "title": "Report"}, + ), + ToolCall( + step_id="s4", + tool_name="send_notification", + parameters={ + "channel": "email", + "recipient": "a@b.com", + "subject": "Report", + "body": "Done", + "attach_step": "s3", + }, + ), + ToolCall( + step_id="s5", + tool_name="store_results", + parameters={"source_step": "s2", "target": "csv_file", "target_name": "out.csv"}, + ), ], source="execution", ) @@ -35,7 +61,11 @@ def test_valid_data_flow(self): task_description="test", tool_chain=[ ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT 1"}), - ToolCall(step_id="s2", tool_name="generate_report", parameters={"source_step": "s1", "format": "csv", "title": "t"}), + ToolCall( + step_id="s2", + tool_name="generate_report", + parameters={"source_step": "s1", "format": "csv", "title": "t"}, + ), ], ) passed, _msg = sandbox.execute(trace) @@ -47,7 +77,11 @@ def test_invalid_source_ref(self): task_description="test", tool_chain=[ ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT 1"}), - ToolCall(step_id="s2", tool_name="generate_report", parameters={"source_step": "nonexistent", "format": "csv", "title": "t"}), + ToolCall( + step_id="s2", + tool_name="generate_report", + parameters={"source_step": "nonexistent", "format": "csv", "title": "t"}, + ), ], ) passed, msg = sandbox.execute(trace) diff --git a/tests/unit/test_schema_validator.py b/tests/unit/test_schema_validator.py index b362c3c..274c48c 100644 --- a/tests/unit/test_schema_validator.py +++ b/tests/unit/test_schema_validator.py @@ -29,7 +29,11 @@ def test_duplicate_step_ids(self, benchmark_registry): task_description="test", tool_chain=[ ToolCall(step_id="s1", tool_name="query_database", parameters={"query": "SELECT 1"}), - ToolCall(step_id="s1", tool_name="generate_report", parameters={"source_step": "s1", "format": "csv", "title": "t"}), + ToolCall( + step_id="s1", + tool_name="generate_report", + parameters={"source_step": "s1", "format": "csv", "title": "t"}, + ), ], ) is_valid, failures = validator.validate(trace) From 1b7005b59d029f84d10923dc3541baaa80d0f2be Mon Sep 17 00:00:00 2001 From: harskuma Date: Tue, 19 May 2026 16:03:21 +0530 Subject: [PATCH 3/4] Fix mypy strict mode: use TYPE_CHECKING guard and scoped CLI override - store.py: Replace `PGVector = None # type: ignore` with proper TYPE_CHECKING import + deferred factory function. Avoids unused-ignore error when langchain-postgres is transitively installed. - pyproject.toml: Remove global disallow_untyped_decorators=false and warn_unused_ignores=false. Add scoped [[tool.mypy.overrides]] for behavioral_memory.cli only, since typer decorators lack type stubs. Co-authored-by: Cursor --- pyproject.toml | 4 ++++ src/behavioral_memory/memory/store.py | 30 +++++++++++++++++++++------ 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2d4b0c2..bc74579 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -142,6 +142,10 @@ module = [ ] ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "behavioral_memory.cli" +disallow_untyped_decorators = false + [tool.pytest.ini_options] testpaths = ["tests"] pythonpath = ["src"] diff --git a/src/behavioral_memory/memory/store.py b/src/behavioral_memory/memory/store.py index 7cd98a5..ff9869c 100644 --- a/src/behavioral_memory/memory/store.py +++ b/src/behavioral_memory/memory/store.py @@ -11,7 +11,7 @@ import json import logging -from typing import Any +from typing import TYPE_CHECKING, Any from langchain_core.documents import Document from langchain_core.embeddings import Embeddings @@ -22,13 +22,16 @@ logger = logging.getLogger(__name__) +_HAS_PGVECTOR = False try: - from langchain_postgres import PGVector + import langchain_postgres # noqa: F401 _HAS_PGVECTOR = True except ImportError: - _HAS_PGVECTOR = False - PGVector = None # type: ignore[assignment, misc] + pass + +if TYPE_CHECKING: + from langchain_postgres import PGVector def _check_postgres_deps() -> None: @@ -38,6 +41,22 @@ def _check_postgres_deps() -> None: ) +def _create_pgvector( + collection_name: str, + connection: str, + embeddings: Embeddings, +) -> PGVector: + """Deferred import + construction so the module loads without postgres deps.""" + from langchain_postgres import PGVector + + return PGVector( + collection_name=collection_name, + connection=connection, + embeddings=embeddings, + use_jsonb=True, + ) + + class TraceStore: """Vector store for validated execution traces. @@ -66,11 +85,10 @@ def __init__( def vectorstore(self) -> Any: if self._vectorstore is None: try: - self._vectorstore = PGVector( + self._vectorstore = _create_pgvector( collection_name=self._collection_name, connection=self._connection_url, embeddings=self._embeddings, - use_jsonb=True, ) except Exception as e: raise MemoryStoreError(f"Failed to connect to vector store: {e}") from e From ef64580415ab4bd72e0af582d592d0fb72a21e63 Mon Sep 17 00:00:00 2001 From: harskuma Date: Tue, 19 May 2026 16:24:40 +0530 Subject: [PATCH 4/4] Update Langfuse integration for SDK v4 API - tracer.py: Replace deprecated client.trace() with client.start_observation() + span.start_observation() for child generations. Matches Langfuse SDK v4.6+ API. - feedback.py: Replace client.fetch_traces/fetch_scores with client.api.trace.list/scores.list for v4 compatibility. Co-authored-by: Cursor --- .../observability/feedback.py | 10 ++++---- src/behavioral_memory/observability/tracer.py | 23 +++++++++++++------ 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/src/behavioral_memory/observability/feedback.py b/src/behavioral_memory/observability/feedback.py index 8f200dd..86fe8ba 100644 --- a/src/behavioral_memory/observability/feedback.py +++ b/src/behavioral_memory/observability/feedback.py @@ -44,14 +44,13 @@ def fetch_positive_traces(self) -> list[ExecutionTrace]: """Fetch traces with positive SME scores from Langfuse. Returns ExecutionTrace objects ready for gatekeeper evaluation. + Compatible with Langfuse SDK v4+ (client.api.trace/scores). """ if self.client is None: return [] try: - traces_response = self.client.fetch_traces( - tags=["behavioral-memory"], - ) + traces_response = self.client.api.trace.list(tags="behavioral-memory") traces = traces_response.data if hasattr(traces_response, "data") else [] except Exception as e: logger.warning("Failed to fetch traces from Langfuse: %s", e) @@ -70,7 +69,10 @@ def fetch_positive_traces(self) -> list[ExecutionTrace]: def _has_positive_score(self, trace: Any) -> bool: """Check if a trace has a positive score meeting the threshold.""" try: - scores_response = self.client.fetch_scores(trace_id=trace.id) + scores_response = self.client.api.scores.list( + trace_id=trace.id, + config_id=None, + ) scores = scores_response.data if hasattr(scores_response, "data") else [] for score in scores: if ( diff --git a/src/behavioral_memory/observability/tracer.py b/src/behavioral_memory/observability/tracer.py index 0245974..cdfba6d 100644 --- a/src/behavioral_memory/observability/tracer.py +++ b/src/behavioral_memory/observability/tracer.py @@ -2,6 +2,8 @@ Logs execution traces, retrieved examples, and generated plans to Langfuse for observability and later SME review (Section III.F). + +Compatible with Langfuse SDK v4+. """ from __future__ import annotations @@ -55,6 +57,7 @@ def log_plan( """Log a generated plan to Langfuse. Returns the trace ID if successful, None otherwise. + Compatible with Langfuse SDK v4+ (start_as_current_observation API). """ if not self.enabled or self.client is None: return None @@ -64,32 +67,38 @@ def log_plan( if tags: trace_tags.extend(tags) - trace = self.client.trace( + output_data = json.dumps([s.model_dump() for s in plan.steps], indent=2) + + span = self.client.start_observation( name="plan_generation", + as_type="span", input=plan.query, - output=json.dumps([s.model_dump() for s in plan.steps], indent=2), - user_id=user_id, - session_id=session_id, - tags=trace_tags, + output=output_data, metadata={ "retrieved_traces_count": len(plan.retrieved_traces), "schemas_used": [s.name for s in plan.schemas_used], "token_budget_used": plan.token_budget_used, "steps_count": len(plan.steps), + "tags": trace_tags, + "user_id": user_id, + "session_id": session_id, }, ) - trace.generation( + generation = span.start_observation( name="llm_plan_generation", + as_type="generation", input=plan.query, output=plan.raw_llm_output, metadata={ "retrieved_examples": [t.task_description for t in plan.retrieved_traces], }, ) + generation.end() + span.end() self.client.flush() - trace_id: str = trace.id + trace_id: str = span.trace_id logger.info("Logged plan to Langfuse: trace_id=%s", trace_id) return trace_id