CAPO: A Primal-Dual Framework for Constraint-Aware Prompt Optimization

Large Language Models (LLMs) are increasingly deployed in agentic contexts, where the model relies on a system prompt to use tools and complete tasks. Implicit in these agentic settings are operational requirements: using the right tools, keeping prompts at reasonable length, achieving parsimonious solution paths, and complying with safety and formatting policies. For many practitioners, assembling domain-specific supervised data to post-train LLMs to satisfy such requirements is infeasible. In this paper, we introduce CAPO (Constraint-Aware Prompt Optimization), an \emph{explicit threshold-constrained} prompt-optimization algorithm based on primal--dual updates. CAPO employs a primal--dual framework to optimize system prompts under explicit constraints via Lagrangian relaxation. Our results across three agentic benchmarks show that CAPO more reliably reaches empirically feasible operating points while improving agentic performance. We also demonstrate that the algorithm generalizes beyond agentic use cases, achieving strong performance on assistant-style evaluations that require satisfying output-format and safety/privacy constraints. Finally, we show that CAPO's rewrite policy can be amortized into DCAPO, a feedback-aware trainable rewriter that matches CAPO accuracy while improving the constrained trade-off.