From 887c7d464f21be7e27b8fb467e3bdc85075d2b05 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 20:04:53 +0300 Subject: [PATCH 1/7] feat: tool response capping + centralize config into src/config/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool response capping: - src/tools/_cap.py: cap_tool_result() truncates dicts that exceed TOOL_RESPONSE_MAX_CHARS (default 40 000 chars ≈ 10 K tokens) and appends a notice telling the agent to narrow its query - src/agent/core.py: wrap ALL_TOOLS with with_cap() at agent init - 9 unit tests in tests/test_tools/test_cap.py Config centralization: - src/config/appsettings.py: Settings class (moved from agent/config.py) - src/config/__init__.py: re-exports Settings and settings - Deleted src/agent/config.py; all 21 import sites updated to `from config import settings` --- .env.example | 1 + README.md | 1 + scripts/setup_db.py | 2 +- src/agent/core.py | 6 +- src/agent/db/__init__.py | 2 +- src/agent/db/postgres.py | 2 +- src/agent/db/sqlite.py | 2 +- src/agent/poller.py | 2 +- src/agent/turns.py | 2 +- src/api/app.py | 2 +- src/api/routers/chat.py | 2 +- src/cli/ask.py | 2 +- src/cli/main.py | 2 +- src/config/__init__.py | 3 + .../config.py => config/appsettings.py} | 5 ++ src/mcp_server.py | 2 +- src/tools/_cache.py | 2 +- src/tools/_cap.py | 52 +++++++++++++ src/tools/cloudtrail.py | 2 +- src/tools/cloudwatch.py | 2 +- src/tools/ec2.py | 2 +- src/tools/ecs.py | 2 +- src/tools/iam.py | 2 +- src/tools/lambda_.py | 2 +- src/tools/rds.py | 2 +- tests/test_backends/test_backends.py | 2 +- tests/test_tools/test_cap.py | 74 +++++++++++++++++++ 27 files changed, 160 insertions(+), 22 deletions(-) create mode 100644 src/config/__init__.py rename src/{agent/config.py => config/appsettings.py} (88%) create mode 100644 src/tools/_cap.py create mode 100644 tests/test_tools/test_cap.py diff --git a/.env.example b/.env.example index 7239c1f..b0857ec 100644 --- a/.env.example +++ b/.env.example @@ -32,6 +32,7 @@ AWS_PROFILE=devops-agent-readonly # optional MAX_TOOL_CALLS=20 INVESTIGATION_TIMEOUT=120 LOG_LEVEL=INFO +# TOOL_RESPONSE_MAX_CHARS=40000 # cap large tool responses before they reach the LLM; 0 = disabled # Storage backend — uncomment exactly ONE block below: # diff --git a/README.md b/README.md index 9e9d193..ae4d35a 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,7 @@ docs/ | `AWS_PROFILE` | none | AWS named profile (e.g. `devops-agent-readonly`) | | `MAX_TOOL_CALLS` | `20` | Hard cap on tool calls per investigation | | `INVESTIGATION_TIMEOUT` | `120` | Timeout in seconds | +| `TOOL_RESPONSE_MAX_CHARS` | `40000` | Truncate tool responses larger than this before feeding to the LLM; `0` disables | | `SLACK_WEBHOOK_URL` | none | Slack incoming webhook URL; leave unset to disable notifications | | `POLL_INTERVAL_MINUTES` | `0` | Proactive polling interval in minutes; `0` disables the poller | | `POLL_ERROR_THRESHOLD` | `5.0` | Lambda error rate % that triggers an automatic investigation | diff --git a/scripts/setup_db.py b/scripts/setup_db.py index d44cdff..e4b4cc1 100644 --- a/scripts/setup_db.py +++ b/scripts/setup_db.py @@ -22,7 +22,7 @@ from dotenv import load_dotenv load_dotenv(Path(__file__).parent.parent / ".env") -from agent.config import settings # noqa: E402 — must come after load_dotenv +from config import settings # noqa: E402 — must come after load_dotenv def _banner(msg: str) -> None: diff --git a/src/agent/core.py b/src/agent/core.py index 18ae7c5..a15391e 100644 --- a/src/agent/core.py +++ b/src/agent/core.py @@ -8,10 +8,11 @@ from deepagents import create_deep_agent from langchain_litellm import ChatLiteLLM -from agent.config import settings +from config import settings from models.agent import Confidence, Investigation, InvestigationResult, RootCauseCategory from agent.prompts import SYSTEM_PROMPT from loguru import logger +from tools._cap import with_cap from tools.cloudtrail import ALL_CLOUDTRAIL_TOOLS from tools.cloudwatch import ALL_CLOUDWATCH_TOOLS from tools.ec2 import ALL_EC2_TOOLS @@ -60,9 +61,10 @@ def init_agent(checkpointer: Any) -> None: api_base=settings.llm_api_base or None, api_key=settings.llm_api_key or None, ) + tools = [with_cap(t) for t in ALL_TOOLS] if settings.tool_response_max_chars > 0 else ALL_TOOLS _agent = create_deep_agent( model=model, - tools=ALL_TOOLS, + tools=tools, system_prompt=SYSTEM_PROMPT, checkpointer=checkpointer, ) diff --git a/src/agent/db/__init__.py b/src/agent/db/__init__.py index e13c202..da68bb8 100644 --- a/src/agent/db/__init__.py +++ b/src/agent/db/__init__.py @@ -10,7 +10,7 @@ def _create_backend() -> DatabaseBackend: - from agent.config import settings + from config import settings backend = settings.checkpoint_backend diff --git a/src/agent/db/postgres.py b/src/agent/db/postgres.py index df0cce5..d3c14b7 100644 --- a/src/agent/db/postgres.py +++ b/src/agent/db/postgres.py @@ -8,7 +8,7 @@ from loguru import logger from agent.db.base import DatabaseBackend -from agent.config import settings +from config import settings class PostgresBackend(DatabaseBackend): diff --git a/src/agent/db/sqlite.py b/src/agent/db/sqlite.py index 364efbe..d01cdcf 100644 --- a/src/agent/db/sqlite.py +++ b/src/agent/db/sqlite.py @@ -10,7 +10,7 @@ from loguru import logger from agent.db.base import DatabaseBackend -from agent.config import settings +from config import settings _DDL = """ diff --git a/src/agent/poller.py b/src/agent/poller.py index 57e9cb1..6b2f33c 100644 --- a/src/agent/poller.py +++ b/src/agent/poller.py @@ -15,7 +15,7 @@ from loguru import logger -from agent.config import settings +from config import settings # In-memory dedup: maps a trigger key → datetime last investigated. diff --git a/src/agent/turns.py b/src/agent/turns.py index 4dc66e3..4e212d2 100644 --- a/src/agent/turns.py +++ b/src/agent/turns.py @@ -6,7 +6,7 @@ from loguru import logger -from agent.config import settings +from config import settings from agent.db import db # Fallback pricing ($/M tokens) for models absent from LiteLLM's database. diff --git a/src/api/app.py b/src/api/app.py index a7cde79..cc4165e 100644 --- a/src/api/app.py +++ b/src/api/app.py @@ -37,7 +37,7 @@ def emit(self, record: logging.LogRecord) -> None: @asynccontextmanager async def lifespan(_app: FastAPI): import asyncio - from agent.config import settings as _cfg + from config import settings as _cfg checkpointer = await db.init() init_agent(checkpointer) diff --git a/src/api/routers/chat.py b/src/api/routers/chat.py index a6ecbbb..4ab8f71 100644 --- a/src/api/routers/chat.py +++ b/src/api/routers/chat.py @@ -13,7 +13,7 @@ from fastapi.responses import StreamingResponse from loguru import logger -from agent.config import settings +from config import settings from agent.core import get_agent from agent.turns import calc_cost, save_turn, notify_slack from models.chat import ChatRequest diff --git a/src/cli/ask.py b/src/cli/ask.py index 2b09c2f..2f3d009 100644 --- a/src/cli/ask.py +++ b/src/cli/ask.py @@ -8,7 +8,7 @@ from rich.markdown import Markdown from rich.progress import Progress, SpinnerColumn, TextColumn -from agent.config import settings +from config import settings console = Console() diff --git a/src/cli/main.py b/src/cli/main.py index 5fd512e..ce61aa1 100644 --- a/src/cli/main.py +++ b/src/cli/main.py @@ -4,7 +4,7 @@ from loguru import logger from rich.console import Console -from agent.config import settings +from config import settings app = typer.Typer( name="devops-agent", diff --git a/src/config/__init__.py b/src/config/__init__.py new file mode 100644 index 0000000..c027e72 --- /dev/null +++ b/src/config/__init__.py @@ -0,0 +1,3 @@ +from config.appsettings import Settings, settings + +__all__ = ["Settings", "settings"] diff --git a/src/agent/config.py b/src/config/appsettings.py similarity index 88% rename from src/agent/config.py rename to src/config/appsettings.py index 3339624..64635a4 100644 --- a/src/agent/config.py +++ b/src/config/appsettings.py @@ -23,6 +23,11 @@ class Settings(BaseSettings): investigation_timeout: int = 120 log_level: str = "INFO" + # Cap tool responses before feeding them back to the LLM. + # Prevents large CloudWatch / CloudTrail payloads from exhausting the context window. + # ~40 K chars ≈ 10 K tokens. Set to 0 to disable. + tool_response_max_chars: int = 40_000 + # Storage backend: "memory" | "sqlite" | "postgres" # memory → no persistence, zero config (default, great for CI / quick testing) # sqlite → local file-based persistence, zero external dependencies diff --git a/src/mcp_server.py b/src/mcp_server.py index 8411d8c..b182a09 100644 --- a/src/mcp_server.py +++ b/src/mcp_server.py @@ -20,7 +20,7 @@ from fastmcp import FastMCP from loguru import logger -from agent.config import settings +from config import settings mcp = FastMCP( name="OpenDevOps Agent", diff --git a/src/tools/_cache.py b/src/tools/_cache.py index 3b28335..d42f6ce 100644 --- a/src/tools/_cache.py +++ b/src/tools/_cache.py @@ -18,7 +18,7 @@ from cachetools import TTLCache, cached from cachetools.keys import hashkey -from agent.config import settings +from config import settings F = TypeVar("F", bound=Callable[..., Any]) diff --git a/src/tools/_cap.py b/src/tools/_cap.py new file mode 100644 index 0000000..7d7bf4d --- /dev/null +++ b/src/tools/_cap.py @@ -0,0 +1,52 @@ +"""Utility: cap tool response size before the result reaches the LLM context window. + +AWS tool responses (CloudWatch logs, CloudTrail events, etc.) can be very large. +Feeding them untruncated into the LLM wastes tokens and can exhaust the context window. +This module truncates oversized responses and appends a notice so the agent knows +to narrow its query. +""" + +from __future__ import annotations + +import json +from functools import wraps +from typing import Any + +from config import settings + + +def cap_tool_result(result: Any, max_chars: int | None = None) -> Any: + """Return result unchanged if within the size limit; truncate and annotate if not. + + max_chars defaults to settings.tool_response_max_chars when not supplied. + """ + if not isinstance(result, dict): + return result + + limit = max_chars if max_chars is not None else settings.tool_response_max_chars + try: + serialized = json.dumps(result) + except (TypeError, ValueError): + return result + + if len(serialized) <= limit: + return result + + return { + "_capped": True, + "_original_chars": len(serialized), + "_notice": ( + f"Response truncated from {len(serialized):,} to {limit:,} chars " + "to protect the context window. Use more specific filters, shorter time " + "ranges, or pagination to retrieve focused data." + ), + "_data": serialized[:limit], + } + + +def with_cap(fn: Any) -> Any: + """Wrap a tool function so its return value is passed through cap_tool_result.""" + @wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return cap_tool_result(fn(*args, **kwargs)) + return wrapper diff --git a/src/tools/cloudtrail.py b/src/tools/cloudtrail.py index ac0b897..6786b01 100644 --- a/src/tools/cloudtrail.py +++ b/src/tools/cloudtrail.py @@ -8,7 +8,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/cloudwatch.py b/src/tools/cloudwatch.py index 46e4580..9b1beea 100644 --- a/src/tools/cloudwatch.py +++ b/src/tools/cloudwatch.py @@ -9,7 +9,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/ec2.py b/src/tools/ec2.py index 654fcb2..580a962 100644 --- a/src/tools/ec2.py +++ b/src/tools/ec2.py @@ -7,7 +7,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/ecs.py b/src/tools/ecs.py index 058a82f..46af9cc 100644 --- a/src/tools/ecs.py +++ b/src/tools/ecs.py @@ -7,7 +7,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/iam.py b/src/tools/iam.py index 8d8f943..a35c965 100644 --- a/src/tools/iam.py +++ b/src/tools/iam.py @@ -7,7 +7,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/lambda_.py b/src/tools/lambda_.py index d5c5d38..1b6bc64 100644 --- a/src/tools/lambda_.py +++ b/src/tools/lambda_.py @@ -8,7 +8,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/src/tools/rds.py b/src/tools/rds.py index 9792b4d..f57dd1d 100644 --- a/src/tools/rds.py +++ b/src/tools/rds.py @@ -8,7 +8,7 @@ from loguru import logger -from agent.config import settings +from config import settings from tools._cache import tool_cached diff --git a/tests/test_backends/test_backends.py b/tests/test_backends/test_backends.py index c55af04..3fabe01 100644 --- a/tests/test_backends/test_backends.py +++ b/tests/test_backends/test_backends.py @@ -39,7 +39,7 @@ async def pg_backend() -> DatabaseBackend: pytest.skip("DATABASE_URL not set — skipping postgres backend tests") from agent.db.postgres import PostgresBackend - from agent.config import settings + from config import settings settings.database_url = url b = PostgresBackend() diff --git a/tests/test_tools/test_cap.py b/tests/test_tools/test_cap.py new file mode 100644 index 0000000..689de9a --- /dev/null +++ b/tests/test_tools/test_cap.py @@ -0,0 +1,74 @@ +"""Tests for the tool response capping utility.""" + +from __future__ import annotations + +import json + +from tools._cap import cap_tool_result, with_cap + + +def test_small_result_passes_through(): + result = {"alarms": ["a", "b"], "count": 2} + assert cap_tool_result(result) is result + + +def test_large_result_is_capped(): + big = {"events": ["x" * 100] * 1000} + capped = cap_tool_result(big, max_chars=500) + assert isinstance(capped, dict) + assert capped["_capped"] is True + assert "_notice" in capped + assert len(capped["_data"]) <= 500 + + +def test_capped_data_is_valid_json_prefix(): + data = {"logs": [{"msg": "hello " * 20}] * 200} + capped = cap_tool_result(data, max_chars=200) + assert "_data" in capped + assert isinstance(capped["_data"], str) + + +def test_non_dict_passes_through(): + assert cap_tool_result("plain string") == "plain string" + assert cap_tool_result(None) is None + assert cap_tool_result([1, 2, 3]) == [1, 2, 3] + + +def test_original_chars_recorded(): + big = {"k": "v" * 10_000} + capped = cap_tool_result(big, max_chars=100) + assert capped["_original_chars"] == len(json.dumps(big)) + + +def test_with_cap_decorator_passes_small_result(): + def my_tool() -> dict: + return {"count": 1} + + wrapped = with_cap(my_tool) + assert wrapped() == {"count": 1} + + +def test_with_cap_decorator_caps_large_result(): + def my_tool() -> dict: + return {"data": "x" * 50_000} + + wrapped = with_cap(my_tool) + result = wrapped() + assert result["_capped"] is True + + +def test_with_cap_preserves_function_name(): + def get_alarms() -> dict: + return {} + + wrapped = with_cap(get_alarms) + assert wrapped.__name__ == "get_alarms" + + +def test_with_cap_preserves_docstring(): + def get_alarms() -> dict: + """List CloudWatch alarms.""" + return {} + + wrapped = with_cap(get_alarms) + assert wrapped.__doc__ == "List CloudWatch alarms." From eea720daf8af09702a7f3ef9acbd67fbd90c3601 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 20:24:41 +0300 Subject: [PATCH 2/7] feat: skills system with lambda-throttling skill - src/tools/skills.py: auto-discovers src/skills/*/SKILL.md files, exposes list_skills() and use_skill(name) tools to the agent - src/skills/lambda-throttling/SKILL.md: step-by-step throttling investigation (concurrency checks, CloudTrail changes, account limits, traffic spikes, mitigation table) - src/agent/prompts.py: injects available skill names + descriptions into system prompt at startup; full content only loaded on use_skill() call - src/config/appsettings.py: checkpoint_backend uses Literal type - README: skills system marked done in mid-term roadmap --- README.md | 2 + src/agent/core.py | 2 + src/agent/prompts.py | 33 ++++++++++- src/config/appsettings.py | 6 +- src/skills/lambda-throttling/SKILL.md | 75 +++++++++++++++++++++++++ src/tools/skills.py | 79 +++++++++++++++++++++++++++ 6 files changed, 194 insertions(+), 3 deletions(-) create mode 100644 src/skills/lambda-throttling/SKILL.md create mode 100644 src/tools/skills.py diff --git a/README.md b/README.md index ae4d35a..c6cab6f 100644 --- a/README.md +++ b/README.md @@ -207,6 +207,7 @@ docs/ - [x] **Multi-provider LLM support** — 100+ providers via LiteLLM; swap models with a single `LLM_MODEL` env var change; supports OpenRouter, Anthropic, OpenAI, Groq, Ollama, and any OpenAI-compatible endpoint; see [docs/llm_providers.md](docs/llm_providers.md) - [x] **MCP integration** — expose the agent as an MCP server (`devops-agent mcp`); `investigate`, `ask`, and `list_sessions` tools available in Claude Desktop, Cursor, or any MCP-compatible client; stdio and HTTP+SSE transports; see [docs/mcp_server.md](docs/mcp_server.md) - [x] **Multi-backend storage** — `memory` (zero config), `sqlite` (local file, no external service), `postgres` (production); switch with one env var; see [docs/databases.md](docs/databases.md) +- [x] **Skills system** — on-demand investigation skills loaded from `src/skills/*/SKILL.md`; skill names injected into system prompt at startup, full content loaded only when agent calls `use_skill(name)`; ships with `lambda-throttling` skill; add your own by dropping a `SKILL.md` into `src/skills//` - [ ] **Custom tools via URL** — register external tools by pointing at an OpenAPI/HTTP endpoint; agent discovers and calls them alongside built-in AWS tools - [x] **Bash CLI escape hatch (Phase 1)** — `run_bash_command` is implemented for read-only AWS CLI, kubectl, and docker commands with strict allowlist validation and timeout. - [ ] **Bash sandbox Phase 2** — run each bash command in an isolated throwaway container (`--network none`, read-only FS, non-root, resource limits). @@ -219,6 +220,7 @@ docs/ ### Later - [ ] **Observability** — OpenTelemetry traces for agent steps, tool call latency, LLM token usage +- [ ] **Follow-up question suggestions** — after each investigation completes, generate 3 suggested follow-up questions in the background and surface them in the UI as clickable chips - [ ] **Session / user feedback loop** — thumbs up/down on investigations, feed signals back to the agent and to an internal quality dashboard - [ ] **Telegram integration** — bot that accepts `/investigate` commands and streams findings back to a chat or group - [ ] **Knowledge base** — attach internal runbooks, post-mortems, and architecture docs so the agent grounds answers in org-specific context diff --git a/src/agent/core.py b/src/agent/core.py index a15391e..3935d10 100644 --- a/src/agent/core.py +++ b/src/agent/core.py @@ -23,6 +23,7 @@ from tools.iam import ALL_IAM_TOOLS from tools.lambda_ import ALL_LAMBDA_TOOLS from tools.rds import ALL_RDS_TOOLS +from tools.skills import ALL_SKILL_TOOLS ALL_TOOLS = ( ALL_CLOUDWATCH_TOOLS @@ -34,6 +35,7 @@ + ALL_IAM_TOOLS + ALL_HISTORY_TOOLS + ALL_BASH_TOOLS + + ALL_SKILL_TOOLS + ALL_FINAL_ANSWER_TOOLS ) diff --git a/src/agent/prompts.py b/src/agent/prompts.py index 06af6e9..1049de6 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -1,4 +1,6 @@ -SYSTEM_PROMPT = """You are an expert AWS SRE investigating an incident. You have read-only access to AWS services via tools. +from __future__ import annotations + +_BASE_PROMPT = """You are an expert AWS SRE investigating an incident. You have read-only access to AWS services via tools. ## Investigation Methodology @@ -9,6 +11,12 @@ 5. Form explicit hypotheses before calling tools to verify them. 6. Rank hypotheses by likelihood and confirm or rule out each one with evidence. +## Skills + +You have access to investigation skills for common incident types via `list_skills` and `use_skill`. +{runbook_section} +When the incident matches a known skill, call `use_skill(name)` early in the investigation to load its step-by-step guidance. + ## ECS Investigations Before calling `list_ecs_services` or `describe_ecs_service`, always call `list_ecs_clusters` first to discover real cluster names — never guess them. @@ -55,3 +63,26 @@ Be concise. Skip obvious observations. Go straight to anomalies. If you're uncertain, say so explicitly and reflect it in the confidence level. """ + + +def build_system_prompt() -> str: + """Build the system prompt, injecting the list of available runbooks.""" + try: + from tools.skills import available_skill_summaries + summaries = available_skill_summaries() + except Exception: + summaries = [] + + if summaries: + lines = ["Available skills (call `use_skill(name)` to load full content):"] + for r in summaries: + lines.append(f" - `{r['name']}` — {r['description']}") + runbook_section = "\n".join(lines) + "\n" + else: + runbook_section = "" + + return _BASE_PROMPT.format(runbook_section=runbook_section) + + +# Evaluated once at import time so the agent gets a stable prompt per process. +SYSTEM_PROMPT = build_system_prompt() diff --git a/src/config/appsettings.py b/src/config/appsettings.py index 64635a4..d904b53 100644 --- a/src/config/appsettings.py +++ b/src/config/appsettings.py @@ -1,3 +1,5 @@ +from typing import Literal + from pydantic_settings import BaseSettings, SettingsConfigDict @@ -28,11 +30,11 @@ class Settings(BaseSettings): # ~40 K chars ≈ 10 K tokens. Set to 0 to disable. tool_response_max_chars: int = 40_000 - # Storage backend: "memory" | "sqlite" | "postgres" + # Storage backend — controls which DatabaseBackend is used. # memory → no persistence, zero config (default, great for CI / quick testing) # sqlite → local file-based persistence, zero external dependencies # postgres → full production persistence - checkpoint_backend: str = "memory" + checkpoint_backend: Literal["memory", "sqlite", "postgres"] = "memory" # SQLite file path — only used when checkpoint_backend = "sqlite" sqlite_path: str = "./data/agent.db" diff --git a/src/skills/lambda-throttling/SKILL.md b/src/skills/lambda-throttling/SKILL.md new file mode 100644 index 0000000..911ece6 --- /dev/null +++ b/src/skills/lambda-throttling/SKILL.md @@ -0,0 +1,75 @@ +--- +name: lambda-throttling +description: Investigate Lambda function throttling — concurrency limits, burst limits, reserved concurrency misconfiguration +--- + +## Lambda Throttling Investigation + +### Step 1 — Identify the throttled function + +- Call `get_alarms(state="ALARM")` and look for alarms on `Throttles` metric +- If the function name is known, call `get_error_rate(function_name)` directly +- Otherwise call `list_lambda_functions()` to enumerate candidates + +### Step 2 — Confirm throttling vs errors + +Lambda reports two distinct failure modes — confirm which you're dealing with: + +- **Throttles** (`AWS/Lambda > Throttles`) — requests rejected before execution starts; exit code is 429 +- **Errors** (`AWS/Lambda > Errors`) — function executed but threw an exception + +Use `get_metric_data` with: +- Namespace: `AWS/Lambda`, Metric: `Throttles`, Dimension: `FunctionName` +- Namespace: `AWS/Lambda`, Metric: `ConcurrentExecutions`, Dimension: `FunctionName` +- Namespace: `AWS/Lambda`, Metric: `ConcurrentExecutions` (no dimension) — account-level + +### Step 3 — Check concurrency configuration + +Call `get_function_config(function_name)` and look for: + +- `ReservedConcurrentExecutions` — if set, this is the hard cap for this function + - A value of `0` means the function is fully throttled (intentional or misconfiguration) + - A value lower than current demand causes throttling even if account limit is not hit +- `ProvisionedConcurrencyConfigs` — if missing or underpowered, cold starts under burst may throttle + +### Step 4 — Check account-level concurrency limit + +Account default is 1000 concurrent executions (region-scoped). If the account-level +`ConcurrentExecutions` metric is near 1000, other functions may be consuming the budget. + +Use `run_bash_command("aws lambda get-account-settings")` to see: +- `TotalConcurrentExecutions` — current usage +- `UnreservedConcurrentExecutions` — what's left for functions without reserved concurrency + +### Step 5 — Check for recent changes + +Call `get_trail_events(hours=2)` and filter for: +- `UpdateFunctionConfiguration` — memory/timeout/concurrency change +- `PutFunctionConcurrency` — reserved concurrency was set (possibly to 0) +- `DeleteFunctionConcurrency` — concurrency limit removed +- `UpdateFunctionCode` — new deployment that may have increased cold-start time + +### Step 6 — Check for traffic spike + +If no config change is found, the cause is likely demand exceeding capacity: + +- Check `Invocations` metric for a spike in the throttling window +- Check upstream triggers: API Gateway (`5XXError`), SQS queue depth (`ApproximateNumberOfMessagesVisible`), EventBridge rule frequency +- Use `run_bash_command("aws sqs get-queue-attributes --queue-url --attribute-names ApproximateNumberOfMessagesVisible")` if SQS-triggered + +### Common Root Causes + +| Cause | Signal | Fix | +|---|---|---| +| Reserved concurrency set too low | `ReservedConcurrentExecutions` < demand | Raise or remove the reservation | +| Reserved concurrency = 0 | `ReservedConcurrentExecutions = 0` | Remove the 0-cap (likely misconfiguration) | +| Account concurrency limit hit | Account-level `ConcurrentExecutions` near 1000 | Request limit increase via AWS Support | +| Burst limit exceeded | Throttles spike immediately after a cold traffic surge | Use provisioned concurrency or gradual traffic ramp | +| SQS backlog overwhelming function | Large `ApproximateNumberOfMessagesVisible` | Reduce `MaximumConcurrency` on event source mapping or scale up reserved concurrency | + +### Mitigation Steps (read-only — document for human action) + +1. If `ReservedConcurrentExecutions = 0`: remove the reservation via AWS console or CLI +2. If reservation is set too low: raise it to match peak concurrent demand plus 20% headroom +3. If account limit is hit: open AWS Support case for concurrency increase, or identify and cap noisy-neighbour functions +4. If burst limit: enable provisioned concurrency for the function, or add exponential-backoff retry on the caller side diff --git a/src/tools/skills.py b/src/tools/skills.py new file mode 100644 index 0000000..dbff2f7 --- /dev/null +++ b/src/tools/skills.py @@ -0,0 +1,79 @@ +"""Skills registry — auto-discovers SKILL.md files under src/skills/. + +Two tools are exposed to the agent: + list_skills() — returns names + descriptions of all available skills + use_skill(name) — loads the full investigation steps for a named skill + +The agent should call use_skill() for whichever skill matches the incident type. +Available skill names and descriptions are already injected into the system prompt, +so list_skills() is only needed if the agent wants to refresh or confirm what exists. +Full content is loaded on demand so unrelated skills cost zero tokens. +""" + +from __future__ import annotations + +import re +from functools import lru_cache +from pathlib import Path + +_SKILLS_DIR = Path(__file__).parent.parent / "skills" +_FRONTMATTER_RE = re.compile(r"^---\n(.*?)\n---[ \t]*\n?", re.DOTALL) + + +def _parse_frontmatter(text: str) -> dict[str, str]: + m = _FRONTMATTER_RE.match(text) + if not m: + return {} + result: dict[str, str] = {} + for line in m.group(1).splitlines(): + if ":" in line: + key, _, val = line.partition(":") + result[key.strip()] = val.strip() + return result + + +@lru_cache(maxsize=1) +def _load_registry() -> dict[str, dict[str, str]]: + """Scan skills dir once and cache {name: {description, content}}.""" + registry: dict[str, dict[str, str]] = {} + if not _SKILLS_DIR.exists(): + return registry + for skill_path in sorted(_SKILLS_DIR.glob("*/SKILL.md")): + text = skill_path.read_text(encoding="utf-8") + meta = _parse_frontmatter(text) + slug = skill_path.parent.name + name = meta.get("name") or slug + description = meta.get("description", "") + body = _FRONTMATTER_RE.sub("", text).strip() + registry[name] = {"slug": slug, "description": description, "content": body} + return registry + + +def available_skill_summaries() -> list[dict[str, str]]: + """Return [{name, description}] for all discovered skills (no full content).""" + return [ + {"name": name, "description": info["description"]} + for name, info in _load_registry().items() + ] + + +def list_skills() -> dict: + """List all available investigation skills with their names and descriptions.""" + items = available_skill_summaries() + return {"skills": items, "count": len(items)} + + +def use_skill(name: str) -> dict: + """Load the full investigation skill for a named incident type. + The skill contains step-by-step investigation guidance, key metrics to check, + log patterns to look for, and common root causes with mitigations.""" + registry = _load_registry() + if name not in registry: + return { + "error": f"Skill '{name}' not found.", + "available": list(registry.keys()), + } + return {"name": name, "skill": registry[name]["content"]} + + +ALL_SKILL_TOOLS = [list_skills, use_skill] From 1e8602d9c2c7062a1509d53c7512251661a14f71 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 20:27:46 +0300 Subject: [PATCH 3/7] docs: add skills.md and update README roadmap --- README.md | 2 + docs/skills.md | 103 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 docs/skills.md diff --git a/README.md b/README.md index c6cab6f..d861691 100644 --- a/README.md +++ b/README.md @@ -211,6 +211,8 @@ docs/ - [ ] **Custom tools via URL** — register external tools by pointing at an OpenAPI/HTTP endpoint; agent discovers and calls them alongside built-in AWS tools - [x] **Bash CLI escape hatch (Phase 1)** — `run_bash_command` is implemented for read-only AWS CLI, kubectl, and docker commands with strict allowlist validation and timeout. - [ ] **Bash sandbox Phase 2** — run each bash command in an isolated throwaway container (`--network none`, read-only FS, non-root, resource limits). +- [x] **Tool response capping** — truncates oversized AWS tool responses (CloudWatch logs, CloudTrail events) before they reach the LLM context window; configurable via `TOOL_RESPONSE_MAX_CHARS` (default 40 000 chars ≈ 10 K tokens) +- [ ] **Conversation summarization** — automatically summarize old messages when the session approaches the model's context limit; preserves recent exchanges and injects a structured summary so long investigations never fail mid-session - [ ] **Optimize tool loading** — pass only relevant tools per investigation context instead of the full 19-tool set - [ ] **Message middleware pipeline** — compaction, summarization, intent detection, context trimmer - [ ] **Guardrails** — input/output validation, PII scrubbing, query scope enforcement diff --git a/docs/skills.md b/docs/skills.md new file mode 100644 index 0000000..d82fa94 --- /dev/null +++ b/docs/skills.md @@ -0,0 +1,103 @@ +# Skills + +Skills are on-demand investigation guides the agent can load during a session. +Each skill is a markdown file (`SKILL.md`) that contains step-by-step investigation +steps, key metrics to check, log patterns to look for, and a common root-cause table +for a specific incident type. + +--- + +## How it works + +1. At startup, the skill registry scans `src/skills/*/SKILL.md` and reads the + frontmatter (name + description) from each file. +2. Available skill names and descriptions are injected into the agent's system prompt — + this costs only a few tokens regardless of how many skills exist. +3. During an investigation, the agent calls `use_skill("lambda-throttling")` (or + whichever name matches the incident). Only then is the full skill content loaded + into the context window. +4. Unrelated skills cost **zero tokens** — their content is never loaded. + +--- + +## Built-in skills + +| Name | Description | +|---|---| +| `lambda-throttling` | Investigate Lambda function throttling — concurrency limits, burst limits, reserved concurrency misconfiguration | + +--- + +## Adding a skill + +1. Create a directory under `src/skills/` named after the incident type: + ``` + src/skills/ecs-oom/ + ``` + +2. Create `SKILL.md` inside it with the required frontmatter: + ```markdown + --- + name: ecs-oom + description: Investigate ECS task out-of-memory crashes and container restarts + --- + + ## ECS OOM Investigation + + ### Step 1 — Confirm OOM kills + ... + ``` + +3. Restart the server (or the CLI). The new skill is auto-discovered and injected + into the system prompt — no code changes needed. + +--- + +## Skill file format + +``` +src/skills/ +└── / + └── SKILL.md ← required +``` + +**`SKILL.md` structure:** + +```markdown +--- +name: +description: +--- + +## Investigation steps go here +... +``` + +- `name` — must be unique; this is what the agent passes to `use_skill()` +- `description` — shown in the system prompt; keep it under 120 chars so it doesn't bloat the prompt +- Body — plain markdown; use headers, bullet lists, and tables freely + +--- + +## Agent tools + +Two tools are available to the agent: + +| Tool | Description | +|---|---| +| `list_skills()` | Returns all skill names and descriptions | +| `use_skill(name)` | Loads the full skill content for the given name | + +The agent typically calls `use_skill()` directly (names are already in the system +prompt) and only falls back to `list_skills()` if it needs to confirm what is available. + +--- + +## Ideas for future skills + +- `ecs-oom` — ECS container out-of-memory crashes +- `rds-high-connections` — RDS max connection limit exceeded +- `api-gateway-5xx` — API Gateway 5XX spike investigation +- `ec2-cpu-spike` — EC2 CPU utilisation runaway +- `lambda-cold-start` — Lambda cold start latency spikes +- `ecs-deployment-failure` — ECS service failing to stabilise after deploy From f9ff675742d605829445aaa6d002967bd24a3d54 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 20:43:40 +0300 Subject: [PATCH 4/7] feat: conversation summarization with usage tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Automatically compacts long sessions before the LLM context window overflows. Fires before each agent call when total message chars exceed SUMMARIZATION_THRESHOLD_CHARS (default 60 000 chars ≈ 15 K tokens). - src/agent/summarizer.py: maybe_summarize() reads LangGraph state, splits at HumanMessage boundaries (no orphaned ToolMessages), calls LLM to produce a structured summary, applies RemoveMessage + injects summary HumanMessage via aupdate_state() - src/api/routers/chat.py: single await maybe_summarize() call before astream() - src/config/appsettings.py: summarization_enabled, threshold, keep_chars settings - save_usage_event() gains metadata: dict param across all backends so summarization events are recorded with {"summarization": True, "messages_removed": N, "chars_removed": N} for dashboard queries - migrations/003: ADD COLUMN metadata jsonb to usage_events (postgres) - SQLite DDL + inline ALTER TABLE migration for existing databases --- .env.example | 3 + migrations/003_usage_events_metadata.sql | 4 + src/agent/db/base.py | 1 + src/agent/db/memory.py | 2 + src/agent/db/postgres.py | 6 +- src/agent/db/sqlite.py | 16 +- src/agent/summarizer.py | 196 +++++++++++++++++++++++ src/api/routers/chat.py | 3 + src/config/appsettings.py | 7 + 9 files changed, 234 insertions(+), 4 deletions(-) create mode 100644 migrations/003_usage_events_metadata.sql create mode 100644 src/agent/summarizer.py diff --git a/.env.example b/.env.example index b0857ec..b54fd2b 100644 --- a/.env.example +++ b/.env.example @@ -33,6 +33,9 @@ MAX_TOOL_CALLS=20 INVESTIGATION_TIMEOUT=120 LOG_LEVEL=INFO # TOOL_RESPONSE_MAX_CHARS=40000 # cap large tool responses before they reach the LLM; 0 = disabled +# SUMMARIZATION_ENABLED=true +# SUMMARIZATION_THRESHOLD_CHARS=60000 # compact session when total message chars exceed this (~15K tokens) +# SUMMARIZATION_KEEP_CHARS=20000 # preserve this many recent chars intact (~5K tokens) # Storage backend — uncomment exactly ONE block below: # diff --git a/migrations/003_usage_events_metadata.sql b/migrations/003_usage_events_metadata.sql new file mode 100644 index 0000000..8441dd1 --- /dev/null +++ b/migrations/003_usage_events_metadata.sql @@ -0,0 +1,4 @@ +-- Add metadata column to usage_events for tracking per-event context +-- (e.g. summarization flag, chars_removed, etc.) +ALTER TABLE usage_events + ADD COLUMN IF NOT EXISTS metadata jsonb NOT NULL DEFAULT '{}'::jsonb; diff --git a/src/agent/db/base.py b/src/agent/db/base.py index 4bcb642..476d990 100644 --- a/src/agent/db/base.py +++ b/src/agent/db/base.py @@ -64,6 +64,7 @@ async def save_usage_event( cost_usd: float | None, latency_ms: int, tool_call_count: int, + metadata: dict | None = None, ) -> None: ... @abstractmethod diff --git a/src/agent/db/memory.py b/src/agent/db/memory.py index 27db8f1..2b6166b 100644 --- a/src/agent/db/memory.py +++ b/src/agent/db/memory.py @@ -116,6 +116,7 @@ async def save_usage_event( cost_usd: float | None, latency_ms: int, tool_call_count: int, + metadata: dict | None = None, ) -> None: self._usage[session_id].append({ "session_id": session_id, @@ -126,6 +127,7 @@ async def save_usage_event( "cost_usd": cost_usd, "latency_ms": latency_ms, "tool_call_count": tool_call_count, + "metadata": metadata or {}, }) async def list_sessions(self) -> list[dict]: diff --git a/src/agent/db/postgres.py b/src/agent/db/postgres.py index d3c14b7..67fef7a 100644 --- a/src/agent/db/postgres.py +++ b/src/agent/db/postgres.py @@ -181,17 +181,19 @@ async def save_usage_event( cost_usd: float | None, latency_ms: int, tool_call_count: int, + metadata: dict | None = None, ) -> None: await self._exec( """ INSERT INTO usage_events (session_id, message_id, model, input_tokens, output_tokens, - cost_usd, latency_ms, tool_call_count) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s) + cost_usd, latency_ms, tool_call_count, metadata) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) """, uuid.UUID(session_id), uuid.UUID(message_id) if message_id else None, model, input_tokens, output_tokens, cost_usd, latency_ms, tool_call_count, + self._jsonb(metadata or {}), ) async def list_sessions(self) -> list[dict]: diff --git a/src/agent/db/sqlite.py b/src/agent/db/sqlite.py index d01cdcf..b6e4a85 100644 --- a/src/agent/db/sqlite.py +++ b/src/agent/db/sqlite.py @@ -56,6 +56,7 @@ cost_usd REAL, latency_ms INTEGER NOT NULL DEFAULT 0, tool_call_count INTEGER NOT NULL DEFAULT 0, + metadata TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%S', 'now')) ); """ @@ -104,6 +105,15 @@ async def init(self) -> Any: await self._conn.execute(stmt) await self._conn.commit() + # Migrate existing databases: add metadata column if absent + try: + await self._conn.execute( + "ALTER TABLE usage_events ADD COLUMN metadata TEXT NOT NULL DEFAULT '{}'" + ) + await self._conn.commit() + except Exception: + pass # column already exists + # LangGraph checkpointer uses its own connection to the same file cp_conn = await aiosqlite.connect(self._path) self._checkpointer = AsyncSqliteSaver(cp_conn) @@ -224,16 +234,18 @@ async def save_usage_event( cost_usd: float | None, latency_ms: int, tool_call_count: int, + metadata: dict | None = None, ) -> None: await self._exec( """ INSERT INTO usage_events (id, session_id, message_id, model, input_tokens, output_tokens, - cost_usd, latency_ms, tool_call_count) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + cost_usd, latency_ms, tool_call_count, metadata) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, str(uuid.uuid4()), session_id, message_id, model, input_tokens, output_tokens, cost_usd, latency_ms, tool_call_count, + json.dumps(metadata or {}), ) async def list_sessions(self) -> list[dict]: diff --git a/src/agent/summarizer.py b/src/agent/summarizer.py new file mode 100644 index 0000000..c9cc728 --- /dev/null +++ b/src/agent/summarizer.py @@ -0,0 +1,196 @@ +"""Conversation summarizer — compacts long sessions before context window overflow. + +Triggered automatically before each agent call when total message chars exceed +SUMMARIZATION_THRESHOLD_CHARS. Replaces old messages with a structured summary +injected as a HumanMessage, preserving the last SUMMARIZATION_KEEP_CHARS of recent +context intact. The system prompt is never touched — LangGraph handles it separately. + +Summarization events are tracked in usage_events with metadata.summarization=True +so the dashboard can report how many sessions and tokens were compacted. +""" + +from __future__ import annotations + +import time +from typing import Any + +from langchain_core.messages import HumanMessage, RemoveMessage +from langchain_litellm import ChatLiteLLM +from loguru import logger + +from config import settings + +_SUMMARIZE_SYSTEM = """You are compressing an AWS incident investigation conversation to save context space. + +Produce a factual, dense summary covering: +- What is being investigated (original incident or question) +- Key findings from tool calls (metric values, error messages, timestamps, resource names) +- Hypotheses confirmed or ruled out with evidence +- Current investigation state + +Be specific — include actual numbers, service names, and timestamps where present. +Keep the summary under 400 words. Use this exact format: + +**Investigating:** +**Key findings:** +- +**Ruled out:** +**Current state:** +""" + + +def _msg_chars(msg: Any) -> int: + content = getattr(msg, "content", "") or "" + return len(str(content)) + + +def _total_chars(messages: list) -> int: + return sum(_msg_chars(m) for m in messages) + + +def _split_messages(messages: list, keep_chars: int) -> tuple[list, list]: + """Return (to_summarize, to_keep), always splitting at a HumanMessage boundary. + + Walks from the end accumulating chars until keep_chars is reached, then finds + the nearest preceding HumanMessage as the split point. This ensures we never + orphan ToolMessages whose parent AIMessage has been removed. + """ + if not messages: + return [], [] + + human_indices = [ + i for i, m in enumerate(messages) + if getattr(m, "type", None) == "human" + ] + + # Need at least 2 human turns — otherwise nothing useful to summarize + if len(human_indices) < 2: + return [], messages + + # Walk from end to find the split point + chars = 0 + split_at = None + for i in range(len(messages) - 1, -1, -1): + chars += _msg_chars(messages[i]) + # Only consider HumanMessage boundaries that aren't the very first message + if chars >= keep_chars and i in set(human_indices) and i > human_indices[0]: + split_at = i + break + + if split_at is None: + return [], messages + + return messages[:split_at], messages[split_at:] + + +def _format_for_summary(messages: list) -> str: + """Render messages as readable text for the summarization prompt.""" + lines = [] + for m in messages: + role = getattr(m, "type", type(m).__name__).upper() + content = str(getattr(m, "content", "") or "") + # Truncate very long tool results to keep the summarization prompt reasonable + if len(content) > 2_000: + content = content[:2_000] + "... [truncated]" + lines.append(f"{role}: {content}") + return "\n\n".join(lines) + + +async def maybe_summarize(agent: Any, config: dict, session_id: str) -> bool: + """Check message length and summarize if over threshold. Returns True if ran.""" + if not settings.summarization_enabled or settings.summarization_threshold_chars <= 0: + return False + + sid = session_id[:8] + + try: + state = await agent.aget_state(config) + except Exception as exc: + logger.warning("[{}] summarizer: could not read state — {}", sid, exc) + return False + + messages = (state.values or {}).get("messages", []) + if not messages: + return False + + total = _total_chars(messages) + if total <= settings.summarization_threshold_chars: + return False + + to_summarize, to_keep = _split_messages(messages, settings.summarization_keep_chars) + if len(to_summarize) < 4: + return False + + chars_removed = _total_chars(to_summarize) + logger.info( + "[{}] summarizer: session={} chars, removing {} msgs ({} chars), keeping {} msgs", + sid, total, len(to_summarize), chars_removed, len(to_keep), + ) + + # ── Summarize old messages via LLM ──────────────────────────────────────── + history_text = _format_for_summary(to_summarize) + start = time.time() + try: + llm = ChatLiteLLM( + model=settings.llm_model, + api_base=settings.llm_api_base or None, + api_key=settings.llm_api_key or None, + ) + response = await llm.ainvoke([ + {"role": "system", "content": _SUMMARIZE_SYSTEM}, + {"role": "user", "content": f"Summarize this investigation:\n\n{history_text}"}, + ]) + summary_text = str(response.content) + except Exception as exc: + logger.error("[{}] summarizer: LLM call failed — {}", sid, exc) + return False + + elapsed_ms = int((time.time() - start) * 1000) + + # ── Update LangGraph state ───────────────────────────────────────────────── + summary_msg = HumanMessage( + content=( + "[Context from earlier in this session — summarized to preserve context window]\n\n" + + summary_text + ) + ) + removes = [RemoveMessage(id=m.id) for m in to_summarize] + + try: + await agent.aupdate_state(config, {"messages": removes + [summary_msg]}) + except Exception as exc: + logger.error("[{}] summarizer: state update failed — {}", sid, exc) + return False + + logger.info( + "[{}] summarizer done: -{} chars, summary {} chars, {}ms", + sid, chars_removed, len(summary_text), elapsed_ms, + ) + + # ── Record in usage_events ───────────────────────────────────────────────── + from agent.db import db + from agent.turns import calc_cost + + input_tokens = chars_removed // 4 + output_tokens = len(summary_text) // 4 + cost = calc_cost(settings.llm_model, input_tokens, output_tokens) + + try: + await db.save_usage_event( + session_id, None, + model=settings.llm_model, + input_tokens=input_tokens, + output_tokens=output_tokens, + cost_usd=cost, + latency_ms=elapsed_ms, + tool_call_count=0, + metadata={ + "summarization": True, + "messages_removed": len(to_summarize), + "chars_removed": chars_removed, + }, + ) + except Exception as exc: + logger.warning("[{}] summarizer: could not save usage event — {}", sid, exc) + + return True diff --git a/src/api/routers/chat.py b/src/api/routers/chat.py index 4ab8f71..4164a89 100644 --- a/src/api/routers/chat.py +++ b/src/api/routers/chat.py @@ -15,6 +15,7 @@ from config import settings from agent.core import get_agent +from agent.summarizer import maybe_summarize from agent.turns import calc_cost, save_turn, notify_slack from models.chat import ChatRequest from api.streaming_labels import STREAMING_LABELS @@ -97,6 +98,8 @@ def _flush_text_buf(): timed_out = False + await maybe_summarize(agent, config, session_id) + try: async with asyncio.timeout(settings.investigation_timeout): async for chunk, _meta in agent.astream( diff --git a/src/config/appsettings.py b/src/config/appsettings.py index d904b53..ff23da0 100644 --- a/src/config/appsettings.py +++ b/src/config/appsettings.py @@ -42,6 +42,13 @@ class Settings(BaseSettings): # PostgreSQL connection string — only used when checkpoint_backend = "postgres" database_url: str | None = None + # Conversation summarization — compacts old messages when a session gets long. + # Fires before each agent call when total message chars exceed the threshold. + # Set summarization_enabled=false or summarization_threshold_chars=0 to disable. + summarization_enabled: bool = True + summarization_threshold_chars: int = 60_000 # ~15 K tokens; trigger compaction above this + summarization_keep_chars: int = 20_000 # ~5 K tokens of recent messages to preserve intact + # Slack — leave unset to disable notifications slack_webhook_url: str | None = None From 18a51e9ce8ac2bdad46aff53455ad6c633095d28 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 20:51:23 +0300 Subject: [PATCH 5/7] feat: expose summarization stats in dashboard + add feature docs - dashboard /stats now includes total_summarizations and total_chars_compacted (queried from usage_events.metadata across all three backends) - add docs/tool_capping.md and docs/conversation_summarization.md - update _SUMMARY_KEYS in test_backends to reflect new fields --- docs/conversation_summarization.md | 125 +++++++++++++++++++++++++++ docs/tool_capping.md | 80 +++++++++++++++++ src/agent/db/memory.py | 3 + src/agent/db/postgres.py | 8 +- src/agent/db/sqlite.py | 12 ++- tests/test_backends/test_backends.py | 3 +- 6 files changed, 228 insertions(+), 3 deletions(-) create mode 100644 docs/conversation_summarization.md create mode 100644 docs/tool_capping.md diff --git a/docs/conversation_summarization.md b/docs/conversation_summarization.md new file mode 100644 index 0000000..19ff1df --- /dev/null +++ b/docs/conversation_summarization.md @@ -0,0 +1,125 @@ +# Conversation Summarization + +## What it does + +Long investigations — many tool calls, large log dumps — can push a session's message +history past the LLM's context limit. When that happens the next agent call either +silently drops old messages or throws an error. + +Conversation summarization automatically compacts the oldest portion of a session's +message history before each new LLM call. Old messages are replaced with a single +structured summary, preserving recent context intact. + +## How it works + +`src/agent/summarizer.py` runs `maybe_summarize()` at the start of every chat turn, +before the LangGraph agent call: + +``` +1. Read current session messages via agent.aget_state() +2. Count total characters +3. If total > SUMMARIZATION_THRESHOLD_CHARS: + a. Walk from the end, accumulate chars until SUMMARIZATION_KEEP_CHARS is reached + b. Split only at a HumanMessage boundary (never mid tool-call pair) + c. Call the LLM with a summarization prompt over the old messages + d. Replace old messages with one HumanMessage containing the structured summary + e. Record a usage_event with metadata.summarization = true +4. Proceed with the normal agent call +``` + +The structured summary follows this template: + +``` +**Investigating:** +**Key findings:** +- +**Ruled out:** +**Current state:** +``` + +The system prompt is never touched — LangGraph manages it separately. + +## Configuration + +Set in `.env`: + +```env +# Enable/disable summarization entirely (default: true) +SUMMARIZATION_ENABLED=true + +# Total session chars that trigger a summarization pass (~15K tokens at 4 chars/token) +SUMMARIZATION_THRESHOLD_CHARS=60000 + +# How many recent chars to keep intact after compaction (~5K tokens) +SUMMARIZATION_KEEP_CHARS=20000 +``` + +| Variable | Default | Notes | +|---|---|---| +| `SUMMARIZATION_ENABLED` | `true` | Set to `false` to disable completely | +| `SUMMARIZATION_THRESHOLD_CHARS` | `60000` | ~15K tokens; increase for larger context models | +| `SUMMARIZATION_KEEP_CHARS` | `20000` | ~5K tokens of recent context preserved | + +## Dashboard tracking + +Each summarization run writes a `usage_events` record with: + +```json +{ + "summarization": true, + "messages_removed": 14, + "chars_removed": 42000 +} +``` + +The `/stats` endpoint (dashboard) exposes two aggregate fields in `summary`: + +| Field | Meaning | +|---|---| +| `total_summarizations` | How many times summarization has run across all sessions | +| `total_chars_compacted` | Total characters removed (proxy for tokens saved) | + +These power future dashboard widgets ("N sessions compacted, ~X tokens saved"). + +## Testing it + +**Unit test the split logic directly:** + +```bash +uv run pytest tests/test_agent/ -v -k "summar" +``` + +**Trigger it in a real session:** + +Temporarily lower the threshold in `.env`: + +```env +SUMMARIZATION_ENABLED=true +SUMMARIZATION_THRESHOLD_CHARS=500 # very low — triggers after a couple messages +SUMMARIZATION_KEEP_CHARS=200 +``` + +Then send a few messages in the chat UI. Watch the logs for: + +``` +[abc12345] summarizer: session=1204 chars, removing 8 msgs (820 chars), keeping 4 msgs +[abc12345] summarizer done: -820 chars, summary 312 chars, 1843ms +``` + +After that the LangGraph checkpoint for the session will contain the summary +message instead of the old raw exchanges. + +**Check dashboard stats (SQLite or Postgres):** + +```bash +curl http://localhost:8000/stats | python -m json.tool | grep -A2 summarization +``` + +## Edge cases + +| Scenario | Behaviour | +|---|---| +| Session has < 2 human turns | Summarization skipped — nothing meaningful to compact | +| LLM summarization call fails | Warning logged, session continues without compaction | +| State update fails | Error logged, session continues unchanged | +| `SUMMARIZATION_ENABLED=false` | `maybe_summarize()` returns immediately without any DB or LLM call | diff --git a/docs/tool_capping.md b/docs/tool_capping.md new file mode 100644 index 0000000..2a1d050 --- /dev/null +++ b/docs/tool_capping.md @@ -0,0 +1,80 @@ +# Tool Response Capping + +## What it does + +AWS tool responses (CloudWatch Logs, CloudTrail events, ECS task lists, etc.) can return +payloads that are hundreds of kilobytes. Feeding them untruncated into the LLM wastes +tokens and can exhaust the context window mid-investigation. + +Tool response capping automatically truncates any tool response that exceeds a +configurable character limit, appending a notice that explains what happened and +suggests how to narrow the query. + +## How it works + +Every tool registered in `src/agent/core.py` is wrapped with `with_cap()` at startup: + +```python +# src/agent/core.py +tools = [with_cap(t) for t in ALL_TOOLS] if settings.tool_response_max_chars > 0 else ALL_TOOLS +``` + +The cap logic lives in `src/tools/_cap.py`: + +- Serialises the result to JSON and measures its length. +- If within the limit → returns the result unchanged. +- If over the limit → returns a replacement dict: + +```json +{ + "_capped": true, + "_original_chars": 184320, + "_notice": "Response truncated from 184,320 to 40,000 chars to protect the context window. Use more specific filters, shorter time ranges, or pagination to retrieve focused data.", + "_data": "...first 40,000 chars of the serialised result..." +} +``` + +The agent sees the truncation notice and knows to refine its query. + +## Configuration + +Set in `.env` (or leave unset to use the default): + +```env +# 0 = disabled (no capping); omit to use the default of 40,000 chars +TOOL_RESPONSE_MAX_CHARS=40000 +``` + +| Value | Behaviour | +|---|---| +| `0` | Capping disabled — all responses pass through unchanged | +| `> 0` | Responses larger than this many characters are truncated | + +## Testing it + +Run the dedicated unit tests: + +```bash +uv run pytest tests/test_tools/test_cap.py -v +``` + +To trigger capping manually: + +```python +from tools._cap import cap_tool_result + +big = {"events": ["x"] * 10_000} +result = cap_tool_result(big, max_chars=100) +print(result["_capped"]) # True +print(result["_notice"]) +``` + +Or with a real tool response — set `TOOL_RESPONSE_MAX_CHARS=1` in `.env` and fire +any CloudWatch query; the agent's reply will reference the truncation notice. + +## When capping does NOT fire + +- Non-dict return values are passed through unchanged (e.g., plain strings). +- If `TOOL_RESPONSE_MAX_CHARS=0` the wrapper is never applied. +- The cap applies only to what the LLM sees; the full raw response is **not** logged + or persisted — the agent stores only what it actually received. diff --git a/src/agent/db/memory.py b/src/agent/db/memory.py index 2b6166b..6e8b616 100644 --- a/src/agent/db/memory.py +++ b/src/agent/db/memory.py @@ -194,6 +194,7 @@ async def get_dashboard_stats(self) -> dict: sum(u["latency_ms"] for u in all_usage) // len(all_usage) if all_usage else 0 ) + summ_events = [u for u in all_usage if isinstance(u.get("metadata"), dict) and u["metadata"].get("summarization")] return { "summary": { @@ -205,6 +206,8 @@ async def get_dashboard_stats(self) -> dict: "total_output_tokens": sum(u["output_tokens"] for u in all_usage), "total_cost_usd": total_cost, "avg_latency_ms": avg_latency, + "total_summarizations": len(summ_events), + "total_chars_compacted": sum(e["metadata"].get("chars_removed", 0) for e in summ_events), }, "activity": [], "top_tools": [], diff --git a/src/agent/db/postgres.py b/src/agent/db/postgres.py index 67fef7a..41efcfa 100644 --- a/src/agent/db/postgres.py +++ b/src/agent/db/postgres.py @@ -301,7 +301,11 @@ async def get_dashboard_stats(self) -> dict: (SELECT COALESCE(SUM(input_tokens), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id WHERE s.is_deleted = FALSE) AS total_input_tokens, (SELECT COALESCE(SUM(output_tokens), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id WHERE s.is_deleted = FALSE) AS total_output_tokens, (SELECT COALESCE(SUM(cost_usd), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id WHERE s.is_deleted = FALSE) AS total_cost_usd, - (SELECT COALESCE(AVG(latency_ms), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id WHERE s.is_deleted = FALSE) AS avg_latency_ms + (SELECT COALESCE(AVG(latency_ms), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id WHERE s.is_deleted = FALSE) AS avg_latency_ms, + (SELECT COUNT(*) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id + WHERE s.is_deleted = FALSE AND ue.metadata @> '{"summarization": true}'::jsonb) AS total_summarizations, + (SELECT COALESCE(SUM((ue.metadata->>'chars_removed')::int), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id + WHERE s.is_deleted = FALSE AND ue.metadata @> '{"summarization": true}'::jsonb) AS total_chars_compacted """) summary = summary_row or {} @@ -377,6 +381,8 @@ async def get_dashboard_stats(self) -> dict: "total_output_tokens": int(summary.get("total_output_tokens", 0) or 0), "total_cost_usd": float(summary.get("total_cost_usd", 0) or 0), "avg_latency_ms": round(float(summary.get("avg_latency_ms", 0) or 0)), + "total_summarizations": int(summary.get("total_summarizations", 0) or 0), + "total_chars_compacted": int(summary.get("total_chars_compacted", 0) or 0), }, "activity": [{"date": str(r["day"]), "sessions": int(r["sessions"])} for r in activity_rows], "top_tools": top_tools, diff --git a/src/agent/db/sqlite.py b/src/agent/db/sqlite.py index b6e4a85..e223ef7 100644 --- a/src/agent/db/sqlite.py +++ b/src/agent/db/sqlite.py @@ -355,7 +355,15 @@ async def get_dashboard_stats(self) -> dict: WHERE s.is_deleted = 0) AS total_cost_usd, (SELECT COALESCE(AVG(latency_ms), 0) FROM usage_events ue JOIN sessions s ON s.id = ue.session_id - WHERE s.is_deleted = 0) AS avg_latency_ms + WHERE s.is_deleted = 0) AS avg_latency_ms, + (SELECT COUNT(*) FROM usage_events ue + JOIN sessions s ON s.id = ue.session_id + WHERE s.is_deleted = 0 + AND json_extract(ue.metadata, '$.summarization') = 1) AS total_summarizations, + (SELECT COALESCE(SUM(CAST(json_extract(ue.metadata, '$.chars_removed') AS INTEGER)), 0) + FROM usage_events ue JOIN sessions s ON s.id = ue.session_id + WHERE s.is_deleted = 0 + AND json_extract(ue.metadata, '$.summarization') = 1) AS total_chars_compacted """) or {} activity_rows = await self._fetchall(""" @@ -447,6 +455,8 @@ async def get_dashboard_stats(self) -> dict: "total_output_tokens": int(summary.get("total_output_tokens", 0) or 0), "total_cost_usd": float(summary.get("total_cost_usd", 0) or 0), "avg_latency_ms": round(float(summary.get("avg_latency_ms", 0) or 0)), + "total_summarizations": int(summary.get("total_summarizations", 0) or 0), + "total_chars_compacted": int(summary.get("total_chars_compacted", 0) or 0), }, "activity": [ {"date": r["day"], "sessions": int(r["sessions"])} diff --git a/tests/test_backends/test_backends.py b/tests/test_backends/test_backends.py index 3fabe01..2d873b5 100644 --- a/tests/test_backends/test_backends.py +++ b/tests/test_backends/test_backends.py @@ -192,7 +192,8 @@ async def test_save_usage_event_does_not_raise(backend: DatabaseBackend): _DASHBOARD_KEYS = {"summary", "activity", "top_tools", "service_breakdown", "recent_sessions", "root_causes"} _SUMMARY_KEYS = {"total_sessions", "total_queries", "total_tool_calls", "total_tool_errors", - "total_input_tokens", "total_output_tokens", "total_cost_usd", "avg_latency_ms"} + "total_input_tokens", "total_output_tokens", "total_cost_usd", "avg_latency_ms", + "total_summarizations", "total_chars_compacted"} async def test_dashboard_stats_has_required_keys(backend: DatabaseBackend): From d24955a5d696d7a08301a16b6ad1e2dd809e47a6 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 21:38:53 +0300 Subject: [PATCH 6/7] feat: show summarization stats on dashboard UI Add context management stat row (hidden until first summarization fires): - Sessions compacted: total summarization runs - Context saved: estimated tokens (~chars/4) compacted across all sessions --- frontend/src/pages/DashboardPage.tsx | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/frontend/src/pages/DashboardPage.tsx b/frontend/src/pages/DashboardPage.tsx index 72f6435..4967b76 100644 --- a/frontend/src/pages/DashboardPage.tsx +++ b/frontend/src/pages/DashboardPage.tsx @@ -1,6 +1,6 @@ import { useEffect, useState } from 'react'; import { Link } from 'react-router-dom'; -import { MessageSquare, Wrench, DollarSign, Clock, AlertTriangle, Zap } from 'lucide-react'; +import { MessageSquare, Wrench, DollarSign, Clock, AlertTriangle, Zap, Scissors } from 'lucide-react'; import { fmtCost, fmtTok, relativeTime, cn } from '../lib/utils'; interface Summary { @@ -8,6 +8,7 @@ interface Summary { total_tool_calls: number; total_tool_errors: number; total_input_tokens: number; total_output_tokens: number; total_cost_usd: number; avg_latency_ms: number; + total_summarizations: number; total_chars_compacted: number; } interface ActivityDay { date: string; sessions: number; } interface TopTool { tool: string; count: number; errors: number; } @@ -144,6 +145,24 @@ export default function DashboardPage() { /> + {/* Context management row */} + {(summary.total_summarizations > 0 || summary.total_chars_compacted > 0) && ( +
+ } + /> + 0 ? `~${fmtTok(Math.round(summary.total_chars_compacted / 4))}` : '—'} + sub={`${summary.total_chars_compacted.toLocaleString()} chars compacted`} + icon={} + /> +
+ )} + {/* Activity chart */}

From 61c87730d3a7dba107bc78e0dfd6ef8b1dfa76d0 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Thu, 7 May 2026 21:52:20 +0300 Subject: [PATCH 7/7] docs: mark conversation summarization as done in roadmap --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d861691..bd87aab 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ docs/ - [x] **Bash CLI escape hatch (Phase 1)** — `run_bash_command` is implemented for read-only AWS CLI, kubectl, and docker commands with strict allowlist validation and timeout. - [ ] **Bash sandbox Phase 2** — run each bash command in an isolated throwaway container (`--network none`, read-only FS, non-root, resource limits). - [x] **Tool response capping** — truncates oversized AWS tool responses (CloudWatch logs, CloudTrail events) before they reach the LLM context window; configurable via `TOOL_RESPONSE_MAX_CHARS` (default 40 000 chars ≈ 10 K tokens) -- [ ] **Conversation summarization** — automatically summarize old messages when the session approaches the model's context limit; preserves recent exchanges and injects a structured summary so long investigations never fail mid-session +- [x] **Conversation summarization** — automatically summarize old messages when the session approaches the model's context limit; preserves recent exchanges and injects a structured summary so long investigations never fail mid-session; summarization events tracked in `usage_events.metadata` and surfaced in the dashboard - [ ] **Optimize tool loading** — pass only relevant tools per investigation context instead of the full 19-tool set - [ ] **Message middleware pipeline** — compaction, summarization, intent detection, context trimmer - [ ] **Guardrails** — input/output validation, PII scrubbing, query scope enforcement