TAMPO: Task- and Model-Aware Automatic Prompt Optimization for Auto-Routing in LLM-based Systems

Automatic Prompt Optimization (APO) enables Large Language Models (LLMs) to adapt to specific tasks while minimizing manual engineering costs. However, since existing APO approaches either rely solely on multi-round iterative procedures or use model-specific generators tailored to optimizing prompts for a single model and objective, they are not readily applicable to auto-routing scenarios, which require operating over diverse LLMs and juggling multiple, often competing, trade-offs. To address this issue, we propose TAMPO, a novel task- and model-aware APO framework for auto-routing in LLM-based systems. Specifically, to capture performance variation across a broad range of tasks and models, we construct a comprehensive, heterogeneity-aware dataset to train an uncertainty-aware reward model. Serving as an offline proxy, this reward model can greatly mitigate reward hacking, allowing TAMPO to learn an optimal multi-objective conditional policy for robust prompt generation. Based on the user requirements encoded in our defined preference vector, this policy enables flexible control over prompt generation and supports a cost-effective deployment strategy. Extensive experiments across 86 tasks demonstrate that TAMPO effectively maintains performance stability across diverse tasks and models, providing a robust, controllable solution for auto-routing in various LLM-based systems.

Large language models (LLMs) act as intelligent assistants, but their performance heavily depends on how you ask the question—the "prompt". In real-world applications, systems often dynamically switch between different AI models to balance quality, speed, and computational cost. However, a prompt that works perfectly for one model might fail on another, and manually designing effective prompts for every situation is slow and expensive. To address these problems, we propose TAMPO, a framework that automatically generates prompts tailored to both the user's task and the specific AI model. Unlike previous approaches, TAMPO learns from a massive dataset to understand how different models respond to various prompts and instantly generates a prompt, while allowing users to flexibly steer the AI's focus, whether they want to prioritize high accuracy, fast response, or low token usage. Evaluated across 86 diverse tasks, the experimental results demonstrate that TAMPO can keep AI systems stable and efficient even when switching between models, eliminate endless manual tweaking, and pave the way for more reliable and controllable AI systems.