diff --git a/CLAUDE.md b/CLAUDE.md index f3a8ad3..fa8203a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -90,18 +90,25 @@ Behind the `openapi` Cargo feature. `OpenApiToolAdapter` parses an OpenAPI 3.0 s - Actions: `get`, `set`, `list`, `remove` - Does **not** touch the core agent loop — wired entirely through `SubAgentTool` -### Sub-Agent Multi-Provider Support +### Construction API -`SubAgentTool` supports any provider via `with_model_config()`. Without it, sub-agents default to Anthropic. For non-Anthropic providers (OpenAI, xAI, Groq, etc.), pass the appropriate `ModelConfig`: +The primary constructor is `Agent::from_config(ModelConfig)`: it selects the built-in provider from `config.api`, sets the model id from `config.id`, and resolves the API key from the provider-conventional env var (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `XAI_API_KEY`, …; see `provider::resolve_api_key`). This removes the provider↔config pairing footgun and the doubly-specified model id. ```rust -let config = ModelConfig::xai("grok-3-mini-fast", "Grok 3 Mini Fast"); -let sub = SubAgentTool::new("analyst", Arc::new(OpenAiCompatProvider)) - .with_model(&config.id) - .with_api_key(&key) - .with_model_config(config); +// provider auto-selected, key from XAI_API_KEY +let agent = Agent::from_config(ModelConfig::xai("grok-4-1-fast", "Grok 4.1 Fast")); ``` +Other constructors: +- `Agent::from_provider(provider, config)` — explicit provider (custom impls, test doubles). Pair with `ModelConfig::mock()` in tests. +- `Agent::from_config_with(®istry, config) -> Result<_, AgentBuildError>` — resolve against a custom `ProviderRegistry`. +- `Agent::set_model(config)` — switch model mid-session (re-resolves the env key; re-selects the provider only when it was registry-resolved, never clobbering an explicit one; explicit keys preserved). +- `Agent::new(provider)` + `with_model`/`with_model_config` — the original builder, still supported. + +`SubAgentTool` mirrors these: `from_config`, `from_config_with`, `from_provider`. + +`SubAgentTool` mirrors these: `SubAgentTool::from_config(name, config)` and `SubAgentTool::from_provider(name, provider, config)`. + `AgentLoopConfig` also supports `turn_delay: Option` — an inter-turn delay to throttle API calls for rate-limit-sensitive providers. Exposed on `SubAgentTool` via `with_turn_delay()`. ### Testing diff --git a/README.md b/README.md index c446160..057c072 100644 --- a/README.md +++ b/README.md @@ -91,15 +91,15 @@ tokio = { version = "1", features = ["full"] } ```rust use yoagent::agent::Agent; -use yoagent::provider::AnthropicProvider; +use yoagent::provider::ModelConfig; use yoagent::types::*; #[tokio::main] async fn main() { - let mut agent = Agent::new(AnthropicProvider) - .with_system_prompt("You are a helpful assistant.") - .with_model("claude-sonnet-5") - .with_api_key(std::env::var("ANTHROPIC_API_KEY").unwrap()); + // Provider is selected from the config's protocol; the key is read from + // ANTHROPIC_API_KEY. Call `.with_api_key(key)` to pass one explicitly. + let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) + .with_system_prompt("You are a helpful assistant."); let mut rx = agent.prompt("What is Rust's ownership model?").await; @@ -131,7 +131,7 @@ use yoagent::SkillSet; let skills = SkillSet::load(&["./skills"])?; -let agent = Agent::new(AnthropicProvider) +let agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) .with_system_prompt("You are a coding assistant.") .with_skills(skills) // Injects skill index into system prompt .with_tools(tools); diff --git a/docs/getting-started/quick-start.md b/docs/getting-started/quick-start.md index 658bce0..81170d7 100644 --- a/docs/getting-started/quick-start.md +++ b/docs/getting-started/quick-start.md @@ -4,15 +4,15 @@ ```rust use yoagent::{Agent, AgentEvent, StreamDelta}; -use yoagent::provider::AnthropicProvider; +use yoagent::provider::ModelConfig; use yoagent::tools::default_tools; #[tokio::main] async fn main() { - let mut agent = Agent::new(AnthropicProvider) + // The provider is selected from the config's protocol, and the API key is + // read from ANTHROPIC_API_KEY. Add `.with_api_key(key)` to override it. + let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) .with_system_prompt("You are a helpful coding assistant.") - .with_model("claude-sonnet-5") - .with_api_key(std::env::var("ANTHROPIC_API_KEY").unwrap()) .with_tools(default_tools()); let mut rx = agent.prompt("List the files in the current directory").await; @@ -45,19 +45,18 @@ async fn main() { ## Example with OpenAI-Compatible Provider -For OpenAI, xAI, Groq, DeepSeek, Mistral, MiniMax, Z.ai, Qwen, Ollama, or any compatible API, use `OpenAiCompatProvider` with a `ModelConfig`: +For OpenAI, xAI, Groq, DeepSeek, Mistral, MiniMax, Z.ai, Qwen, Ollama, or any compatible API, pass the matching `ModelConfig` preset — `from_config` picks the OpenAI-compatible provider and resolves the provider's env key automatically: ```rust use yoagent::{Agent, AgentEvent}; -use yoagent::provider::OpenAiCompatProvider; +use yoagent::provider::ModelConfig; use yoagent::tools::default_tools; #[tokio::main] async fn main() { - let mut agent = Agent::new(OpenAiCompatProvider) + // Provider inferred from the preset; key from OPENAI_API_KEY. + let mut agent = Agent::from_config(ModelConfig::openai("gpt-5.5", "GPT-5.5")) .with_system_prompt("You are a helpful assistant.") - .with_model("gpt-5.5") - .with_api_key(std::env::var("OPENAI_API_KEY").unwrap()) .with_tools(default_tools()); let mut rx = agent.prompt("What is 2 + 2?").await; @@ -82,15 +81,13 @@ async fn main() { ```rust use yoagent::{Agent, AgentEvent, StreamDelta}; -use yoagent::provider::AnthropicProvider; +use yoagent::provider::ModelConfig; use yoagent::tools::default_tools; #[tokio::main] async fn main() { - let mut agent = Agent::new(AnthropicProvider) + let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) .with_system_prompt("You are a helpful assistant.") - .with_model("claude-sonnet-5") - .with_api_key(std::env::var("ANTHROPIC_API_KEY").unwrap()) .with_tools(default_tools()); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); diff --git a/examples/basic.rs b/examples/basic.rs index 374d52d..f4a68c2 100644 --- a/examples/basic.rs +++ b/examples/basic.rs @@ -3,17 +3,16 @@ //! Run with: ANTHROPIC_API_KEY=sk-... cargo run --example basic use yoagent::agent::Agent; -use yoagent::provider::AnthropicProvider; +use yoagent::provider::ModelConfig; use yoagent::*; #[tokio::main] async fn main() { - let api_key = std::env::var("ANTHROPIC_API_KEY").expect("Set ANTHROPIC_API_KEY"); - - let mut agent = Agent::new(AnthropicProvider) - .with_system_prompt("You are a helpful assistant. Be concise.") - .with_model("claude-sonnet-5") - .with_api_key(api_key); + // The Anthropic provider is selected from the config's protocol, and the + // API key is read from ANTHROPIC_API_KEY. Add `.with_api_key(key)` to + // pass one explicitly. + let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) + .with_system_prompt("You are a helpful assistant. Be concise."); println!("Sending prompt..."); diff --git a/examples/callbacks.rs b/examples/callbacks.rs index 8626dc8..9fec842 100644 --- a/examples/callbacks.rs +++ b/examples/callbacks.rs @@ -70,10 +70,8 @@ async fn main() { } } - let mut agent = Agent::new(provider) + let mut agent = Agent::from_provider(provider, yoagent::provider::ModelConfig::mock()) .with_system_prompt("You are helpful.") - .with_model("mock") - .with_api_key("test") .with_tools(vec![Box::new(GreetTool)]) // Limit to 5 turns (plenty for this example) .on_before_turn(|messages, turn| { diff --git a/examples/cli.rs b/examples/cli.rs index 394808e..c629a35 100644 --- a/examples/cli.rs +++ b/examples/cli.rs @@ -26,7 +26,7 @@ use std::io::{self, BufRead, Write}; use yoagent::agent::Agent; -use yoagent::provider::{AnthropicProvider, GoogleProvider, ModelConfig, OpenAiCompatProvider}; +use yoagent::provider::ModelConfig; use yoagent::skills::SkillSet; use yoagent::tools::default_tools; use yoagent::*; @@ -126,20 +126,10 @@ async fn main() { SkillSet::load(&skill_dirs).expect("Failed to load skills") }; - let mut agent = if let Some(ref url) = api_url { - if provider_name.as_deref() == Some("ollama") { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::ollama(url, &model)) - } else { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::local(url, &model)) - } - } else if let Some(ref prov) = provider_name { - make_provider_agent(prov, &model) - } else { - Agent::new(AnthropicProvider) - }; - agent = agent + let mut agent = build_agent(&api_url, &provider_name, &model) .with_system_prompt(SYSTEM_PROMPT) - .with_model(&model) + // from_config already resolves the provider-conventional env key; this + // override preserves the CLI's API_KEY fallback (empty = leave to env). .with_api_key(&api_key) .with_skills(skills.clone()) .with_tools(default_tools()); @@ -189,22 +179,8 @@ async fn main() { } s if s.starts_with("/model ") => { let new_model = s.trim_start_matches("/model ").trim(); - agent = if let Some(ref url) = api_url { - if provider_name.as_deref() == Some("ollama") { - Agent::new(OpenAiCompatProvider) - .with_model_config(ModelConfig::ollama(url, new_model)) - } else { - Agent::new(OpenAiCompatProvider) - .with_model_config(ModelConfig::local(url, new_model)) - } - } else if let Some(ref prov) = provider_name { - make_provider_agent(prov, new_model) - } else { - Agent::new(AnthropicProvider) - }; - agent = agent + agent = build_agent(&api_url, &provider_name, new_model) .with_system_prompt(SYSTEM_PROMPT) - .with_model(new_model) .with_api_key(&api_key) .with_skills(skills.clone()) .with_tools(default_tools()); @@ -334,36 +310,46 @@ async fn main() { println!("\n{DIM} bye 👋{RESET}\n"); } +/// Select the config for the requested provider/URL and build an agent from +/// it. A local/OpenAI-compatible URL wins; then a named provider; else +/// Anthropic. Every branch flows through `from_config`, so the provider, +/// model id, and context window all come from a single `ModelConfig`. +fn build_agent(api_url: &Option, provider_name: &Option, model: &str) -> Agent { + if let Some(url) = api_url { + let config = if provider_name.as_deref() == Some("ollama") { + ModelConfig::ollama(url, model) + } else { + ModelConfig::local(url, model) + }; + Agent::from_config(config) + } else if let Some(prov) = provider_name { + make_provider_agent(prov, model) + } else { + Agent::from_config(ModelConfig::anthropic(model, model)) + } +} + fn make_provider_agent(provider: &str, model: &str) -> Agent { - match provider { - "zai" => Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::zai(model, model)), - "qwen" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::qwen(model, model)) - } - "openai" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::openai(model, model)) - } - "xai" => Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::xai(model, model)), - "groq" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::groq(model, model)) - } - "deepseek" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::deepseek(model, model)) - } - "mistral" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::mistral(model, model)) - } - "minimax" => { - Agent::new(OpenAiCompatProvider).with_model_config(ModelConfig::minimax(model, model)) - } - "ollama" => Agent::new(OpenAiCompatProvider) - .with_model_config(ModelConfig::ollama("http://localhost:11434/v1", model)), - "google" => Agent::new(GoogleProvider).with_model_config(ModelConfig::google(model, model)), + // One `from_config` per provider: the config carries the protocol (so the + // right provider is selected), the model id, context window, and pricing. + // No provider↔config pairing to get wrong, and no model id passed twice. + let config = match provider { + "zai" => ModelConfig::zai(model, model), + "qwen" => ModelConfig::qwen(model, model), + "openai" => ModelConfig::openai(model, model), + "xai" => ModelConfig::xai(model, model), + "groq" => ModelConfig::groq(model, model), + "deepseek" => ModelConfig::deepseek(model, model), + "mistral" => ModelConfig::mistral(model, model), + "minimax" => ModelConfig::minimax(model, model), + "ollama" => ModelConfig::ollama("http://localhost:11434/v1", model), + "google" => ModelConfig::google(model, model), other => { eprintln!("Unknown provider: {other}. Supported: zai, qwen, openai, xai, groq, deepseek, mistral, minimax, ollama, google."); std::process::exit(1); } - } + }; + Agent::from_config(config) } fn truncate(s: &str, max: usize) -> &str { diff --git a/examples/persistence.rs b/examples/persistence.rs index a204429..e45e246 100644 --- a/examples/persistence.rs +++ b/examples/persistence.rs @@ -14,10 +14,8 @@ use yoagent::types::*; async fn main() { // --- Phase 1: Initial conversation --- let provider = MockProvider::text("The capital of France is Paris."); - let mut agent = Agent::new(provider) - .with_system_prompt("You are a helpful assistant.") - .with_model("mock") - .with_api_key("test"); + let mut agent = Agent::from_provider(provider, yoagent::provider::ModelConfig::mock()) + .with_system_prompt("You are a helpful assistant."); println!("=== Phase 1: Initial conversation ==="); let mut rx = agent.prompt("What is the capital of France?").await; @@ -42,10 +40,8 @@ async fn main() { // --- Phase 2: Restore and continue --- let provider2 = MockProvider::text("Paris is also known as the City of Light."); - let mut agent2 = Agent::new(provider2) - .with_system_prompt("You are a helpful assistant.") - .with_model("mock") - .with_api_key("test"); + let mut agent2 = Agent::from_provider(provider2, yoagent::provider::ModelConfig::mock()) + .with_system_prompt("You are a helpful assistant."); agent2.restore_messages(&json).expect("Failed to restore"); println!( diff --git a/examples/sub_agent.rs b/examples/sub_agent.rs index 1184eca..27f0212 100644 --- a/examples/sub_agent.rs +++ b/examples/sub_agent.rs @@ -10,27 +10,27 @@ use std::sync::Arc; use yoagent::agent::Agent; -use yoagent::provider::{AnthropicProvider, StreamProvider}; +use yoagent::provider::ModelConfig; use yoagent::sub_agent::SubAgentTool; use yoagent::tools; use yoagent::*; +// Each construction resolves the Anthropic provider from the config's protocol +// and reads ANTHROPIC_API_KEY from the environment. +fn sonnet() -> ModelConfig { + ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5") +} + #[tokio::main] async fn main() { - let api_key = std::env::var("ANTHROPIC_API_KEY").expect("Set ANTHROPIC_API_KEY"); - let model = "claude-sonnet-5"; - let provider: Arc = Arc::new(AnthropicProvider); - // Sub-agent 1: a researcher with file reading tools - let researcher = SubAgentTool::new("researcher", Arc::clone(&provider)) + let researcher = SubAgentTool::from_config("researcher", sonnet()) .with_description( "Searches and reads files to gather information. Delegate research tasks here.", ) .with_system_prompt( "You are a research assistant. Read files and summarize findings concisely.", ) - .with_model(model) - .with_api_key(&api_key) .with_tools(vec![ Arc::new(tools::ReadFileTool::new()), Arc::new(tools::SearchTool::new()), @@ -39,11 +39,9 @@ async fn main() { .with_max_turns(10); // Sub-agent 2: a coder with file editing tools - let coder = SubAgentTool::new("coder", Arc::clone(&provider)) + let coder = SubAgentTool::from_config("coder", sonnet()) .with_description("Writes and edits code files. Delegate coding tasks here.") .with_system_prompt("You are a coding assistant. Write clean, correct code. Be concise.") - .with_model(model) - .with_api_key(&api_key) .with_tools(vec![ Arc::new(tools::ReadFileTool::new()), Arc::new(tools::WriteFileTool::new()), @@ -52,7 +50,7 @@ async fn main() { .with_max_turns(15); // Parent agent: coordinates between sub-agents - let mut agent = Agent::new(AnthropicProvider) + let mut agent = Agent::from_config(sonnet()) .with_system_prompt( "You are a coordinator agent. You have two sub-agents:\n\ - 'researcher': for reading files and gathering information\n\ @@ -60,8 +58,6 @@ async fn main() { Delegate tasks to the appropriate sub-agent. You can run both in parallel \ when the tasks are independent.", ) - .with_model(model) - .with_api_key(api_key) .with_sub_agent(researcher) .with_sub_agent(coder); diff --git a/src/agent.rs b/src/agent.rs index d4c8734..d184f88 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -36,6 +36,10 @@ pub struct Agent { messages: Vec, tools: Vec>, provider: Arc, + /// Whether `provider` was supplied explicitly by the caller (`new` / + /// `from_provider`) rather than resolved from a registry (`from_config`). + /// `set_model` uses this to avoid silently discarding a caller's provider. + provider_is_explicit: bool, // Queues (shared with the loop via Arc) steering_queue: Arc>>, @@ -71,8 +75,137 @@ pub struct Agent { pending_completion: Option>, Vec)>>, } +/// Error building an [`Agent`] from a [`ModelConfig`] and a registry. +/// +/// Only returned by [`Agent::from_config_with`] / +/// [`SubAgentTool::from_config_with`](crate::SubAgentTool::from_config_with); +/// the built-in `from_config` uses a complete registry and cannot fail. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[non_exhaustive] +pub enum AgentBuildError { + /// The registry has no provider registered for the config's protocol. + #[error( + "no provider registered for protocol {0}; register one on the \ + ProviderRegistry or use Agent::from_provider with an explicit provider" + )] + NoProviderForProtocol(crate::provider::ApiProtocol), +} + impl Agent { + /// Construct from an explicit provider, then configure with + /// [`with_model`](Self::with_model) / [`with_api_key`](Self::with_api_key). + /// + /// Prefer [`from_config`](Self::from_config) (provider + env key resolved + /// from one `ModelConfig`) or [`from_provider`](Self::from_provider) for a + /// custom provider; both avoid the provider↔config mismatch this + /// constructor allows. Kept for backward compatibility. pub fn new(provider: impl StreamProvider + 'static) -> Self { + Self::with_provider_arc(Arc::new(provider)) + } + + /// Build an [`Agent`] from a [`ModelConfig`], selecting the built-in + /// provider for the config's protocol and resolving the API key from the + /// provider-conventional environment variable. + /// + /// This is the primary constructor: the model id, provider, context + /// window, and pricing all come from the one `ModelConfig`, so there's no + /// provider↔config mismatch to get wrong and no model id to pass twice. + /// An explicit key set later via [`with_api_key`](Self::with_api_key) + /// always overrides the environment. + /// + /// ```no_run + /// use yoagent::{Agent, provider::ModelConfig}; + /// // provider auto-selected from config.api; key from ANTHROPIC_API_KEY + /// let agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")); + /// ``` + /// + /// # Panics + /// + /// Never panics — the default registry covers every [`ApiProtocol`] + /// variant (enforced by a unit test). Use + /// [`from_config_with`](Self::from_config_with) with a custom registry + /// when a protocol may be unregistered and you want a `Result` instead. + /// + /// [`ApiProtocol`]: crate::provider::ApiProtocol + pub fn from_config(config: ModelConfig) -> Self { + Self::from_config_with(&crate::provider::ProviderRegistry::default(), config) + .expect("default registry covers all built-in protocols") + } + + /// Like [`from_config`](Self::from_config) but resolves the provider from a + /// caller-supplied registry, returning an error if the config's protocol + /// isn't registered. + pub fn from_config_with( + registry: &crate::provider::ProviderRegistry, + config: ModelConfig, + ) -> Result { + let provider = registry + .resolve(&config.api) + .ok_or(AgentBuildError::NoProviderForProtocol(config.api))?; + let mut agent = Self::with_provider_arc(provider).configured_for(config); + // The provider came from the registry, not the caller — set_model may + // safely re-resolve it on a model switch. + agent.provider_is_explicit = false; + Ok(agent) + } + + /// Build an [`Agent`] from an explicit provider plus its [`ModelConfig`]. + /// + /// The escape hatch for custom [`StreamProvider`] implementations that + /// aren't in any registry (including test doubles — pair with + /// [`ModelConfig::mock`](crate::provider::ModelConfig::mock)). The config + /// is still required so the model id, context window, and pricing stay + /// defined together; the API key is resolved from the environment unless + /// set explicitly. + pub fn from_provider(provider: impl StreamProvider + 'static, config: ModelConfig) -> Self { + Self::with_provider_arc(Arc::new(provider)).configured_for(config) + } + + /// Switch the model mid-session, re-resolving the environment API key from + /// the new config's provider (an explicit key set via + /// [`with_api_key`](Self::with_api_key) is preserved, since the key is + /// resolved lazily and an explicit one always wins). + /// + /// Provider handling depends on how the agent was built: + /// - Built with [`from_config`](Self::from_config): the built-in provider + /// for the new protocol is selected from the default registry. + /// - Built with an **explicit** provider ([`from_provider`](Self::from_provider) + /// or [`new`](Self::new)): that provider is **kept** — it is never + /// silently replaced — and a warning is logged if it may not serve the + /// new protocol. Reconstruct with `from_provider` to change providers. + pub fn set_model(&mut self, config: ModelConfig) { + if self.provider_is_explicit { + tracing::warn!( + "set_model: keeping the explicitly-supplied provider; it may not \ + serve the new protocol {}. Reconstruct with from_provider to \ + change providers.", + config.api + ); + } else if let Some(provider) = + crate::provider::ProviderRegistry::default().resolve(&config.api) + { + self.provider = provider; + } else { + tracing::warn!( + "set_model: no built-in provider for protocol {}; keeping the \ + previous provider, which will not match the new model.", + config.api + ); + } + self.model = config.id.clone(); + self.model_config = Some(config); + } + + /// Apply a `ModelConfig` to a freshly-constructed agent: set the model id + /// and stash the config (provider is already wired). Key stays lazily + /// resolved from the config's provider env var unless set explicitly. + fn configured_for(mut self, config: ModelConfig) -> Self { + self.model = config.id.clone(); + self.model_config = Some(config); + self + } + + fn with_provider_arc(provider: Arc) -> Self { Self { system_prompt: String::new(), model: String::new(), @@ -83,7 +216,10 @@ impl Agent { model_config: None, messages: Vec::new(), tools: Vec::new(), - provider: Arc::new(provider), + provider, + // Explicit by default (new / from_provider); from_config_with + // flips this to false after resolving from the registry. + provider_is_explicit: true, steering_queue: Arc::new(Mutex::new(Vec::new())), follow_up_queue: Arc::new(Mutex::new(Vec::new())), steering_mode: QueueMode::OneAtATime, diff --git a/src/lib.rs b/src/lib.rs index b0b5c2e..03b0985 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,13 +7,13 @@ //! # Quick start //! //! ```no_run -//! use yoagent::{Agent, provider::AnthropicProvider, tools}; +//! use yoagent::{Agent, provider::ModelConfig, tools}; //! //! # #[tokio::main] //! # async fn main() { -//! let mut agent = Agent::new(AnthropicProvider) -//! .with_model("claude-sonnet-5") -//! .with_api_key(std::env::var("ANTHROPIC_API_KEY").unwrap()) +//! // Provider is selected from the config's protocol; the API key is read +//! // from ANTHROPIC_API_KEY. Call `.with_api_key(...)` to override. +//! let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5")) //! .with_system_prompt("You are a helpful coding assistant.") //! .with_tools(tools::default_tools()); //! @@ -65,7 +65,7 @@ pub mod types; #[cfg(feature = "openapi")] pub mod openapi; -pub use agent::Agent; +pub use agent::{Agent, AgentBuildError}; pub use agent_loop::{agent_loop, agent_loop_continue}; pub use context::{CompactionStrategy, DefaultCompaction}; pub use retry::RetryConfig; diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index beb893b..2ff8c7e 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -43,6 +43,10 @@ pub struct AnthropicProvider; #[async_trait] impl StreamProvider for AnthropicProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::AnthropicMessages) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/azure_openai.rs b/src/provider/azure_openai.rs index a0a8a87..1869687 100644 --- a/src/provider/azure_openai.rs +++ b/src/provider/azure_openai.rs @@ -19,6 +19,10 @@ pub struct AzureOpenAiProvider; #[async_trait] impl StreamProvider for AzureOpenAiProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::AzureOpenAiResponses) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/bedrock.rs b/src/provider/bedrock.rs index 6e288ff..9c86f95 100644 --- a/src/provider/bedrock.rs +++ b/src/provider/bedrock.rs @@ -22,6 +22,10 @@ pub struct BedrockProvider; #[async_trait] impl StreamProvider for BedrockProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::BedrockConverseStream) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/google.rs b/src/provider/google.rs index ab8cd90..bcc2c45 100644 --- a/src/provider/google.rs +++ b/src/provider/google.rs @@ -15,6 +15,10 @@ pub struct GoogleProvider; #[async_trait] impl StreamProvider for GoogleProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::GoogleGenerativeAi) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/google_vertex.rs b/src/provider/google_vertex.rs index f8fad3f..38410f1 100644 --- a/src/provider/google_vertex.rs +++ b/src/provider/google_vertex.rs @@ -27,6 +27,10 @@ impl GoogleVertexProvider { #[async_trait] impl StreamProvider for GoogleVertexProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::GoogleVertex) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/model.rs b/src/provider/model.rs index a867fc9..a14be0f 100644 --- a/src/provider/model.rs +++ b/src/provider/model.rs @@ -346,6 +346,30 @@ pub struct ModelConfig { } impl ModelConfig { + /// A minimal config for tests. `provider` is `"mock"`, cost rates are all + /// zero, and `base_url` points at a non-routable host. + /// + /// Use it **only** with + /// [`Agent::from_provider`](crate::Agent::from_provider) / + /// [`SubAgentTool::from_provider`](crate::SubAgentTool::from_provider) and a + /// [`MockProvider`](crate::provider::MockProvider): those take the provider + /// explicitly, so the config's protocol is never consulted. + /// + /// **Do not** pass it to + /// [`Agent::from_config`](crate::Agent::from_config) — that dispatches on + /// the protocol (here `AnthropicMessages`) and would build the **real** + /// Anthropic provider pointed at the non-routable `base_url`, so the first + /// prompt fails with a network error instead of returning a mock response. + pub fn mock() -> Self { + Self::custom( + ApiProtocol::AnthropicMessages, + "mock", + "http://mock.invalid", + "mock", + "Mock", + ) + } + /// Create a config for any protocol without a dedicated preset /// (Bedrock, Vertex, Azure, or future protocols). /// diff --git a/src/provider/openai_compat.rs b/src/provider/openai_compat.rs index 3bf48d8..fb812ce 100644 --- a/src/provider/openai_compat.rs +++ b/src/provider/openai_compat.rs @@ -20,6 +20,10 @@ pub struct OpenAiCompatProvider; #[async_trait] impl StreamProvider for OpenAiCompatProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::OpenAiCompletions) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/openai_responses.rs b/src/provider/openai_responses.rs index c0202d4..b92b0e0 100644 --- a/src/provider/openai_responses.rs +++ b/src/provider/openai_responses.rs @@ -17,6 +17,10 @@ pub struct OpenAiResponsesProvider; #[async_trait] impl StreamProvider for OpenAiResponsesProvider { + fn protocol(&self) -> Option { + Some(crate::provider::ApiProtocol::OpenAiResponses) + } + async fn stream( &self, config: StreamConfig, diff --git a/src/provider/registry.rs b/src/provider/registry.rs index 7364b39..1cb48cb 100644 --- a/src/provider/registry.rs +++ b/src/provider/registry.rs @@ -4,11 +4,12 @@ use super::model::{ApiProtocol, ModelConfig}; use super::traits::*; use crate::types::*; use std::collections::HashMap; +use std::sync::Arc; use tokio::sync::mpsc; /// Registry of all available stream providers, keyed by API protocol. pub struct ProviderRegistry { - providers: HashMap>, + providers: HashMap>, } impl ProviderRegistry { @@ -21,10 +22,21 @@ impl ProviderRegistry { /// Register a provider for a given protocol. pub fn register(&mut self, protocol: ApiProtocol, provider: impl StreamProvider + 'static) { - self.providers.insert(protocol, Box::new(provider)); + self.providers.insert(protocol, Arc::new(provider)); } - /// Get a provider for a given protocol. + /// Resolve the provider for a protocol as a shared handle. + /// + /// Unlike [`get`](Self::get), this hands out an owned `Arc` that can be + /// stored (e.g. by [`Agent::from_config`](crate::Agent::from_config)), + /// so callers don't have to borrow the registry for the provider's + /// lifetime. + pub fn resolve(&self, protocol: &ApiProtocol) -> Option> { + self.providers.get(protocol).cloned() + } + + /// Borrow the provider for a protocol for a transient call. Prefer + /// [`resolve`](Self::resolve) when you need to keep the handle. pub fn get(&self, protocol: &ApiProtocol) -> Option<&dyn StreamProvider> { self.providers.get(protocol).map(|p| p.as_ref()) } @@ -244,3 +256,74 @@ mod resolve_key_tests { std::env::remove_var("YOAGENT_API_KEY"); } } + +#[cfg(test)] +mod registry_tests { + use super::*; + + /// Every `ApiProtocol` variant. The wildcard-free `match` is a compile + /// forcing-function: adding a variant fails to build here until it's added + /// to the list (and, in turn, registered in `default()` below), so the + /// "default registry is complete" invariant behind `from_config`'s unwrap + /// cannot silently rot. + fn all_protocols() -> Vec { + fn _assert_exhaustive(p: ApiProtocol) { + match p { + ApiProtocol::AnthropicMessages + | ApiProtocol::OpenAiCompletions + | ApiProtocol::OpenAiResponses + | ApiProtocol::AzureOpenAiResponses + | ApiProtocol::GoogleGenerativeAi + | ApiProtocol::GoogleVertex + | ApiProtocol::BedrockConverseStream => {} + } + } + vec![ + ApiProtocol::AnthropicMessages, + ApiProtocol::OpenAiCompletions, + ApiProtocol::OpenAiResponses, + ApiProtocol::AzureOpenAiResponses, + ApiProtocol::GoogleGenerativeAi, + ApiProtocol::GoogleVertex, + ApiProtocol::BedrockConverseStream, + ] + } + + /// The default registry must resolve a provider for every protocol, so + /// `Agent::from_config` (which unwraps against it) can never panic. + #[test] + fn default_registry_covers_all_protocols() { + let registry = ProviderRegistry::default(); + for api in all_protocols() { + assert!( + registry.resolve(&api).is_some(), + "default registry missing a provider for {api}" + ); + } + } + + /// Each protocol must resolve to a provider that actually *speaks* that + /// protocol — guards against a mis-wired `register()` line (e.g. mapping + /// `GoogleVertex` to the Anthropic provider) that a mere is-some check + /// would miss. + #[test] + fn default_registry_maps_each_protocol_to_its_own_provider() { + let registry = ProviderRegistry::default(); + for api in all_protocols() { + let provider = registry.resolve(&api).expect("registered"); + assert_eq!( + provider.protocol(), + Some(api), + "protocol {api} resolved to a provider that reports {:?}", + provider.protocol() + ); + } + } + + #[test] + fn empty_registry_resolves_nothing() { + assert!(ProviderRegistry::new() + .resolve(&ApiProtocol::AnthropicMessages) + .is_none()); + } +} diff --git a/src/provider/traits.rs b/src/provider/traits.rs index bf48c54..97abf6a 100644 --- a/src/provider/traits.rs +++ b/src/provider/traits.rs @@ -71,6 +71,16 @@ pub trait StreamProvider: Send + Sync { tx: mpsc::UnboundedSender, cancel: tokio_util::sync::CancellationToken, ) -> Result; + + /// The API protocol this provider speaks, if it maps to a single one. + /// + /// Built-in providers return `Some(_)`; the default is `None` (for test + /// doubles and multi-protocol adapters). Used to verify registry wiring + /// (a resolved provider should report the protocol it was registered for) + /// and to enable protocol-mismatch diagnostics. + fn protocol(&self) -> Option { + None + } } #[derive(Debug, thiserror::Error)] diff --git a/src/sub_agent.rs b/src/sub_agent.rs index 8536078..4ac62dc 100644 --- a/src/sub_agent.rs +++ b/src/sub_agent.rs @@ -90,6 +90,62 @@ impl SubAgentTool { } } + /// Create a sub-agent from a name and [`ModelConfig`], selecting the + /// built-in provider for the config's protocol. + /// + /// Mirrors [`Agent::from_config`](crate::Agent::from_config): the model + /// id, provider, and pricing come from one config, and the API key is + /// resolved from the provider-conventional env var unless set explicitly + /// with [`with_api_key`](Self::with_api_key). + /// + /// # Panics + /// + /// Never panics — the default registry covers every [`ApiProtocol`] + /// variant. Use [`from_config_with`](Self::from_config_with) with a custom + /// registry when a protocol may be unregistered and you want a `Result`. + /// + /// [`ApiProtocol`]: crate::provider::ApiProtocol + pub fn from_config(name: impl Into, config: ModelConfig) -> Self { + Self::from_config_with(&crate::provider::ProviderRegistry::default(), name, config) + .expect("default registry covers all built-in protocols") + } + + /// Like [`from_config`](Self::from_config) but resolves the provider from a + /// caller-supplied registry, returning an error if the config's protocol + /// isn't registered. Mirrors + /// [`Agent::from_config_with`](crate::Agent::from_config_with). + pub fn from_config_with( + registry: &crate::provider::ProviderRegistry, + name: impl Into, + config: ModelConfig, + ) -> Result { + let provider = registry + .resolve(&config.api) + .ok_or(crate::AgentBuildError::NoProviderForProtocol(config.api))?; + Ok(Self::new(name, provider).configured_for(config)) + } + + /// Create a sub-agent from a name, explicit provider, and [`ModelConfig`]. + /// + /// The escape hatch for custom providers and test doubles (pair with + /// [`ModelConfig::mock`](crate::provider::ModelConfig::mock)). Mirrors + /// [`Agent::from_provider`](crate::Agent::from_provider). + pub fn from_provider( + name: impl Into, + provider: Arc, + config: ModelConfig, + ) -> Self { + Self::new(name, provider).configured_for(config) + } + + /// Set the model id and stash the config on a freshly-constructed + /// sub-agent (provider already wired). + fn configured_for(mut self, config: ModelConfig) -> Self { + self.model = config.id.clone(); + self.model_config = Some(config); + self + } + pub fn with_description(mut self, desc: impl Into) -> Self { self.tool_description = desc.into(); self diff --git a/tests/agent_test.rs b/tests/agent_test.rs index 9b574fa..4dcad7f 100644 --- a/tests/agent_test.rs +++ b/tests/agent_test.rs @@ -593,3 +593,135 @@ async fn test_env_var_fallback_resolves_api_key() { run_one_prompt(&mut agent).await; assert_eq!(*captured.lock().unwrap(), "cerebras-env-key"); } + +// --------------------------------------------------------------------------- +// 0.10 construction API: from_config / from_provider / from_config_with / set_model +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_from_provider_runs_end_to_end() { + // from_provider + ModelConfig::mock() is the test-double construction path. + let mut agent = Agent::from_provider( + MockProvider::text("Hi from mock"), + yoagent::provider::ModelConfig::mock(), + ); + assert_eq!(agent.model, "mock"); + + let mut rx = agent.prompt("hello").await; + let mut events = Vec::new(); + while let Some(e) = rx.recv().await { + events.push(e); + } + agent.finish().await; + assert_eq!(agent.messages().len(), 2); +} + +#[test] +fn test_from_config_wires_model_and_config() { + // from_config selects a built-in provider from config.api, sets the id, + // and stashes pricing so session_cost_usd can price the session. + let mut mc = yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"); + mc.cost.input_per_million = 3.0; + let agent = Agent::from_config(mc).with_messages(vec![assistant_with_usage(some_usage())]); + assert_eq!(agent.model, "claude-sonnet-5"); + assert!( + agent.session_cost_usd().is_some_and(|c| c > 0.0), + "from_config must carry the config's pricing" + ); +} + +#[tokio::test] +async fn test_from_config_resolves_env_key() { + std::env::set_var("OPENROUTER_API_KEY", "openrouter-from-config"); + let captured = Arc::new(std::sync::Mutex::new(String::new())); + // from_provider carries the config (so provider="openrouter" drives env + // resolution) while letting us capture the resolved key. + let mut agent = Agent::from_provider( + KeyCapturingProvider { + captured: captured.clone(), + }, + yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "openrouter", + "http://unused.invalid", + "m", + "M", + ), + ); + run_one_prompt(&mut agent).await; + assert_eq!(*captured.lock().unwrap(), "openrouter-from-config"); +} + +#[test] +fn test_from_config_with_errors_on_empty_registry() { + let registry = yoagent::provider::ProviderRegistry::new(); + // Agent isn't Debug, so match instead of expect_err. + let err = match Agent::from_config_with( + ®istry, + yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"), + ) { + Ok(_) => panic!("empty registry must fail"), + Err(e) => e, + }; + assert_eq!( + err, + yoagent::AgentBuildError::NoProviderForProtocol( + yoagent::provider::ApiProtocol::AnthropicMessages + ) + ); +} + +#[tokio::test] +async fn test_set_model_switches_model_id() { + let mut agent = Agent::from_config(yoagent::provider::ModelConfig::anthropic( + "claude-sonnet-5", + "Sonnet 5", + )); + assert_eq!(agent.model, "claude-sonnet-5"); + agent.set_model(yoagent::provider::ModelConfig::anthropic( + "claude-opus-4-8", + "Opus 4.8", + )); + assert_eq!(agent.model, "claude-opus-4-8"); +} + +#[test] +fn test_model_config_mock_is_unpriced() { + let mc = yoagent::provider::ModelConfig::mock(); + assert_eq!(mc.provider, "mock"); + assert!(!mc.cost.is_configured()); +} + +#[tokio::test] +async fn test_set_model_preserves_explicit_provider_and_key() { + // Regression guard for the set_model custom-provider clobber: an agent + // built with an explicit provider must keep BOTH that provider and an + // explicit key across a model switch — never silently swap in a built-in + // network provider. + let captured = Arc::new(std::sync::Mutex::new(String::new())); + let mut agent = Agent::from_provider( + KeyCapturingProvider { + captured: captured.clone(), + }, + yoagent::provider::ModelConfig::mock(), + ) + .with_api_key("explicit-key") + .with_retry_config(yoagent::RetryConfig::none()); + + // Switch to a different protocol whose built-in provider would hit the + // network if it clobbered ours. + agent.set_model(yoagent::provider::ModelConfig::anthropic( + "claude-sonnet-5", + "Sonnet 5", + )); + assert_eq!(agent.model, "claude-sonnet-5"); + + run_one_prompt(&mut agent).await; + // If the provider had been clobbered, our capture probe would never run + // (and the real Anthropic provider would have been used instead). + assert_eq!( + *captured.lock().unwrap(), + "explicit-key", + "set_model must keep the explicit provider and key" + ); +} diff --git a/tests/sub_agent_test.rs b/tests/sub_agent_test.rs index 3c69d0b..51039f0 100644 --- a/tests/sub_agent_test.rs +++ b/tests/sub_agent_test.rs @@ -834,3 +834,62 @@ async fn test_sub_agent_env_key_fallback() { run_sub_agent(&tool).await; assert_eq!(captured.lock().unwrap().0, "minimax-env-key"); } + +#[tokio::test] +async fn test_sub_agent_from_provider_construction() { + // from_provider + ModelConfig::mock() mirrors Agent's construction path. + let tool = SubAgentTool::from_provider( + "researcher", + Arc::new(MockProvider::text("Research result")), + yoagent::provider::ModelConfig::mock(), + ) + .with_description("Researches topics"); + + let result = tool + .execute( + serde_json::json!({"task": "go"}), + ToolContext { + tool_call_id: "tc-fp".into(), + tool_name: "researcher".into(), + cancel: CancellationToken::new(), + on_update: None, + on_progress: None, + }, + ) + .await + .expect("sub-agent should succeed"); + let text = match &result.content[0] { + Content::Text { text } => text, + other => panic!("expected text, got {other:?}"), + }; + assert!(text.contains("Research result")); +} + +#[test] +fn test_sub_agent_from_config_wires_model() { + // from_config selects a built-in provider from config.api and sets the id. + let tool = SubAgentTool::from_config( + "analyst", + yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"), + ); + assert_eq!(tool.name(), "analyst"); +} + +#[test] +fn test_sub_agent_from_config_with_errors_on_empty_registry() { + let registry = yoagent::provider::ProviderRegistry::new(); + let err = match SubAgentTool::from_config_with( + ®istry, + "analyst", + yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"), + ) { + Ok(_) => panic!("empty registry must fail"), + Err(e) => e, + }; + assert_eq!( + err, + yoagent::AgentBuildError::NoProviderForProtocol( + yoagent::provider::ApiProtocol::AnthropicMessages + ) + ); +}