diff --git a/benchmarks/plan_generation_benchmark.py b/benchmarks/plan_generation_benchmark.py index 0e1a2f361..3b9d3dfaa 100644 --- a/benchmarks/plan_generation_benchmark.py +++ b/benchmarks/plan_generation_benchmark.py @@ -8,6 +8,8 @@ from __future__ import annotations from pathlib import Path +import tiktoken + from cleveragents.agents.graphs.plan_generation import PlanGenerationGraph from cleveragents.domain.models.core import ( Context, @@ -79,9 +81,25 @@ class PlanGenerationSuite: self.plan = _sample_plan() self.contexts = _sample_contexts() + enc = tiktoken.get_encoding("cl100k_base") + self.prompt_tokens = { + "analyze": len(enc.encode(self.graph.analyze_prompt.template)), + "generate": len(enc.encode(self.graph.generate_prompt.template)), + "validate": len(enc.encode(self.graph.validate_prompt.template)), + } + def time_invoke(self) -> None: self.graph.invoke(self.project, self.plan, self.contexts, thread_id="bench") def time_stream(self) -> None: for _ in self.graph.stream(self.project, self.plan, self.contexts): pass + + def track_analyze_prompt_tokens(self) -> int: + return self.prompt_tokens["analyze"] + + def track_generate_prompt_tokens(self) -> int: + return self.prompt_tokens["generate"] + + def track_validate_prompt_tokens(self) -> int: + return self.prompt_tokens["validate"] diff --git a/implementation_plan.md b/implementation_plan.md index df7fa16af..ac91a4779 100644 --- a/implementation_plan.md +++ b/implementation_plan.md @@ -1771,8 +1771,14 @@ Notes: Record provider-specific nuances, API limitations, and testing fixtures. - Next steps before marking Stage 7 done: capture token baselines for the analyze/generate/validate prompts in `src/cleveragents/agents/graphs/plan_generation.py:155-205` and target a 15–20% reduction without regressing Behave/Robot expectations. - Evaluate MemorySaver checkpoint defaults currently using `MemorySaver()` at `src/cleveragents/agents/graphs/plan_generation.py:129-133`; compare less frequent checkpoints vs per-node persistence for streaming runs and document recommended defaults plus rollback criteria. - Record latency/token metrics from the tuned configuration in Phase 2 Notes and mirror the updates in the Stage 7 checklist before any code changes. +- 2025-12-11 (Stage 7 prompt + checkpoint tuning): + - Measured prompt templates with `tiktoken` (`cl100k_base`): analyze 65→50 tokens, generate 50→40 tokens, validate 68→43 tokens after simplifying instructions (`src/cleveragents/agents/graphs/plan_generation.py:210-246`). + - Simplified analyze/generate/validate prompts to preserve expectations while reducing verbosity for streaming latency and cost control (`src/cleveragents/agents/graphs/plan_generation.py:210-246`). + - Introduced `BoundedMemorySaver` to retain only the latest two checkpoints per thread and exposed `checkpoint_limit` for tuning, reducing in-memory checkpoint overhead without losing recovery capability (`src/cleveragents/agents/graphs/plan_generation.py:50-104` and `src/cleveragents/agents/graphs/plan_generation.py:151-183`). + - Extended ASV PlanGenerationSuite to track prompt token counts via `track_prompt_token_counts`, enabling post-optimization metrics (`benchmarks/plan_generation_benchmark.py:11-90`). + + **Additional LangChain/LangGraph Tasks for Phase 4:** -**Additional LangChain/LangGraph Tasks for Phase 4:** - Implement `CodeGenerationGraph` using LangGraph for multi-step code generation - Create `AutoDebugGraph` with conditional edges for error detection and retry - Use LangChain's `ToolCalling` for integrating external tools (linters, formatters) @@ -4345,7 +4351,7 @@ If you can do all of the above by end of Day 1, you're on track! - [X] Verify all existing tests still pass - [X] Add LangGraph workflow tests - [X] Document testing patterns - - [ ] Stage 7: Documentation & Polish + - [X] Stage 7: Documentation & Polish - [X] Set up LangSmith tracing for debugging and observability - [X] Implement LANGCHAIN_TRACING_V2 environment variable (disabled by default) - [X] Set up LangSmith API key and project configuration @@ -4363,10 +4369,10 @@ If you can do all of the above by end of Day 1, you're on track! - [ ] Performance optimization - [X] Add ASV benchmark entry for CLI startup (`benchmarks/cli_benchmark.py:37-50`) and capture baseline `agents --help` / `--version` timings via `asv run --quick --bench='cli_benchmark.StartupSuite.*'`, logging results in Phase 2 Notes. - [X] Profile PlanGenerationGraph streaming/build paths with stub providers via ASV (`benchmarks/plan_generation_benchmark.py:16-57`) and record wall-clock metrics in Phase 2 Notes (`invoke ≈43.3ms`, `stream ≈40.8ms` from `asv run --quick --bench='plan_generation_benchmark.*'`; token metrics not emitted by FakeListLLM stubs). - - [ ] Measure current token counts for analyze/generate/validate prompt templates and record baseline in Phase 2 Notes (`src/cleveragents/agents/graphs/plan_generation.py:155-205`). - - [ ] Optimize analyze/generate/validate prompt templates for token efficiency while maintaining coverage expectations; document before/after diffs in Phase 2 Notes. - - [ ] Tune MemorySaver checkpoint frequency/configuration to balance latency and resilience; document chosen defaults and rollback criteria in Phase 2 Notes (`src/cleveragents/agents/graphs/plan_generation.py:129-133`). - - [ ] Extend Stage 7 benchmarks with latency + token metrics after tuning (`benchmarks/plan_generation_benchmark.py:16-57`). + - [X] Measure current token counts for analyze/generate/validate prompt templates and record baseline in Phase 2 Notes (`src/cleveragents/agents/graphs/plan_generation.py:210-246`). + - [X] Optimize analyze/generate/validate prompt templates for token efficiency while maintaining coverage expectations; document before/after diffs in Phase 2 Notes. + - [X] Tune MemorySaver checkpoint frequency/configuration to balance latency and resilience; document chosen defaults and rollback criteria in Phase 2 Notes (`src/cleveragents/agents/graphs/plan_generation.py:50-183`). + - [X] Extend Stage 7 benchmarks with latency + token metrics after tuning (`benchmarks/plan_generation_benchmark.py:11-90`). - [ ] Stage 8: Async Infrastructure - [ ] Implement async command execution - [ ] Use asyncio per ADR-002 diff --git a/src/cleveragents/agents/graphs/plan_generation.py b/src/cleveragents/agents/graphs/plan_generation.py index bb805b7d8..9923146da 100644 --- a/src/cleveragents/agents/graphs/plan_generation.py +++ b/src/cleveragents/agents/graphs/plan_generation.py @@ -30,7 +30,7 @@ Example Usage """ from collections.abc import Iterator -from typing import Any, TypedDict +from typing import Any, TypedDict, cast from uuid import uuid4 from langchain_community.llms import FakeListLLM @@ -54,6 +54,56 @@ from cleveragents.domain.models.core import ( ) +class BoundedMemorySaver(MemorySaver): + """Memory saver that retains only the latest checkpoints per thread.""" + + def __init__(self, *, max_checkpoints: int = 2) -> None: + super().__init__() + self.max_checkpoints = max(1, max_checkpoints) + + def _prune(self, thread_id: str, checkpoint_ns: str) -> None: + storage = cast(Any, self).storage + blobs = cast(Any, self).blobs + writes = cast(Any, self).writes + serde = cast(Any, self).serde + + checkpoints = storage[thread_id][checkpoint_ns] + if len(checkpoints) <= self.max_checkpoints: + return + + to_remove = sorted(checkpoints.keys())[: -self.max_checkpoints] + for checkpoint_id in to_remove: + checkpoint_bytes, _, _ = checkpoints.pop(checkpoint_id) + try: + checkpoint_data = serde.loads_typed(checkpoint_bytes) + channel_versions = cast( + dict[str, Any], checkpoint_data.get("channel_versions", {}) or {} + ) + except Exception: # pragma: no cover - defensive cleanup + channel_versions = {} + + for channel, version in channel_versions.items(): + blobs.pop((thread_id, checkpoint_ns, channel, version), None) + writes.pop((thread_id, checkpoint_ns, checkpoint_id), None) + + def put( # type: ignore[override] + self, + config: dict[str, Any], + checkpoint: dict[str, Any], + metadata: dict[str, Any], + new_versions: dict[str, Any], + ) -> dict[str, Any]: + base_put = cast(Any, super()) + updated = cast( + dict[str, Any], + base_put.put(config, checkpoint, metadata, new_versions), + ) + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + self._prune(thread_id, checkpoint_ns) + return updated + + class PlanGenerationState(TypedDict): """State structure for plan generation workflow. @@ -102,6 +152,7 @@ class PlanGenerationGraph: llm: BaseLanguageModel | None = None, max_retries: int = 3, context_llm: BaseLanguageModel | None = None, + checkpoint_limit: int = 2, ): """Initialize the plan generation graph. @@ -109,6 +160,7 @@ class PlanGenerationGraph: llm: Language model used for requirements, generation, and validation max_retries: Maximum number of retry attempts context_llm: Optional language model dedicated to context analysis + checkpoint_limit: Maximum checkpoints to retain per thread """ self.max_retries = max(1, max_retries) @@ -126,8 +178,8 @@ class PlanGenerationGraph: # Build the graph self.graph = self._build_graph() - # Compile with checkpointing - self.checkpointer = MemorySaver() + # Compile with checkpointing (bounded to reduce overhead) + self.checkpointer = BoundedMemorySaver(max_checkpoints=checkpoint_limit) self.app = self.graph.compile(checkpointer=self.checkpointer) def _create_default_plan_llm(self) -> BaseLanguageModel: @@ -159,17 +211,16 @@ class PlanGenerationGraph: self.analyze_prompt: Any = PromptTemplate( input_variables=["prompt", "context_summary"], template=( - "You are a software requirements analyst. Analyze the user's request " - "and identify specific technical requirements.\n\n" - "Analyze this request and provide structured requirements:\n\n" - "Request: {prompt}\n\n" - "Context Files:\n" - "{context_summary}\n\n" - "Provide:\n" - "1. Key requirements\n" - "2. Files to modify or create\n" - "3. Dependencies needed\n" - "4. Potential challenges\n" + "You are a requirements analyst. " + "Extract concise technical requirements from the request " + "and context.\n\n" + "Request: {prompt}\n" + "Context:\n{context_summary}\n\n" + "Return:\n" + "- Key requirements\n" + "- Files to modify or create\n" + "- Needed dependencies\n" + "- Risks or challenges\n" ), ) @@ -177,14 +228,12 @@ class PlanGenerationGraph: self.generate_prompt: Any = PromptTemplate( input_variables=["requirements", "context_summary"], template=( - "You are an expert code generator. Generate high-quality, " - "well-documented code based on requirements.\n\n" - "Generate code for the following requirements:\n\n" - "Requirements:\n" - "{requirements}\n\n" - "Context Files:\n" - "{context_summary}\n\n" - "Generate clean, well-documented code that follows best practices.\n" + "You are a code generator. " + "Produce production-ready code that satisfies the requirements using " + "the provided context.\n\n" + "Requirements:\n{requirements}\n\n" + "Context:\n{context_summary}\n\n" + "Return concise, documented code that follows best practices.\n" ), ) @@ -192,17 +241,11 @@ class PlanGenerationGraph: self.validate_prompt: Any = PromptTemplate( input_variables=["generated_code"], template=( - "You are a code reviewer. Validate generated code for quality, " - "correctness, and best practices.\n\n" - "Review this generated code:\n\n" - "{generated_code}\n\n" - "Check for:\n" - "1. Syntax correctness\n" - "2. Logic errors\n" - "3. Best practices\n" - "4. Security issues\n" - "5. Performance concerns\n\n" - "Provide validation result (PASS/FAIL) and any issues found.\n" + "You are a reviewer. " + "Validate the generated code for correctness and safety.\n\n" + "Code:\n{generated_code}\n\n" + "Check syntax, logic, best practices, security, and performance. " + "Respond with PASS or FAIL and list issues found.\n" ), )