diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b28a13c7..75b163cdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- *(agent)* [**breaking**] Make `AgentRunner` the only execution path for configured agents: remove the raw `Completion` and `StreamingCompletion` traits and their `Agent` implementations, make agent execution state private, add runner-backed per-request overrides, and route `Extractor` through the full hook lifecycle. Raw hook-free requests remain available explicitly through `CompletionModel`. + - For managed agent execution, replace `agent.completion(prompt, history).await?.send().await?` with `agent.runner(prompt).history(history).max_turns(3).run().await?`, choosing a turn budget large enough for tool follow-ups. + - For managed streaming execution, replace `agent.stream_completion(prompt, history).await?.stream().await?` with `agent.runner(prompt).history(history).max_turns(3).stream().await`. + - The runner consumes tool calls rather than returning the first raw model response. Callers that handled that response manually, and other intentionally hook-free transport, should start from `model.completion_request(prompt).messages(history)` and then call `.send().await?` or `.stream().await?`. + - `AgentRun::new(prompt).with_history(history)` remains a sans-I/O state machine for custom drivers; it contains no configured agent model, tools, memory, or hooks and is not an alternate configured-agent execution path. + - An `Agent`'s model is fixed and private. Former per-call `.model(...)` / `.model_opt(...)` users should retain the provider `CompletionModel` and use its raw request API, or construct a separate `Agent` for the selected model. - *(tool)* [**breaking**] Replace the parallel tool-execution APIs with one structured path. Typed tools now implement only `Tool::call(&mut ToolContext, Args) -> Result`; author-facing errors remain typed until private runtime erasure normalizes them into `ToolExecutionError`, `ToolContext` carries inbound values and host-only result metadata, `ToolResult` is the single runtime observation, and `ToolSet::execute` / `ToolServerHandle::execute` are the dispatch surfaces. Event-specific hook action types make invalid event/action combinations unrepresentable. - Tool implementations: retain one typed `type Error` for ordinary `?` propagation and direct-call tests; remove `classify_error`, `call_with_extensions`, and `call_structured`. The optional `map_error` method classifies domain failures at the erased boundary, while its default preserves the source as `Other`. Return refusals through `map_error` with `ToolExecutionError::refused`, and attach host-only result metadata with `ToolContext::insert_result`. - Context: replace `ToolCallExtensions` and `ToolResultExtensions` with `ToolContext`; replace request/runner `.tool_extensions(...)` with `.tool_context(...)`. Each dispatch snapshots inbound context exactly once, isolates tool-local mutations, and publishes only result metadata back to the caller and hooks. diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index 64ba85a1b..a2dd32e47 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -9,6 +9,69 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- *(agent)* [**breaking**] Make `AgentRunner` the only execution path for configured agents: remove the raw `Completion` and `StreamingCompletion` traits and their `Agent` implementations, make agent execution state private, add runner-backed per-request overrides, and route `Extractor` through the full hook lifecycle. Raw hook-free requests remain available explicitly through `CompletionModel`. + + Migration examples: + + ```rust + // Before: built a one-shot request from configured Agent state, but bypassed + // the AgentRunner lifecycle. + agent.completion(prompt, history).await?.send().await?; + + // After, for managed Agent execution: hooks, tools, retrieval, memory, and + // turn accounting all run. Budget enough calls for tool follow-ups. + agent + .runner(prompt) + .history(history) + .max_turns(3) + .run() + .await?; + ``` + + Streaming follows the same boundary: + + ```rust + // Before + agent.stream_completion(prompt, history).await?.stream().await?; + + // After + let stream = agent + .runner(prompt) + .history(history) + .max_turns(3) + .stream() + .await; + ``` + + The runner consumes tool calls instead of returning the first raw model + response. If the old caller handled that response itself, or for any other + intentionally hook-free provider transport, start from the model rather than + an `Agent`: + + ```rust + model + .completion_request(prompt) + .messages(history) + .send() + .await?; + + let stream = model + .completion_request(prompt) + .messages(history) + .stream() + .await?; + ``` + + `AgentRun::new(prompt).with_history(history)` remains the public sans-I/O + state machine for custom drivers. It contains no model, tools, memory, or hook + stack: callers must handle every `AgentRunStep`, perform provider/tool IO, and + feed results back explicitly. It is not a way to execute a configured + `Agent`; use `AgentRunner` for that. + + An `Agent` also keeps its configured model private and fixed. Applications + that previously called `.model(...)` or `.model_opt(...)` on the returned raw + request builder should retain the provider `CompletionModel` and use its raw + request API, or construct a separate `Agent` for that model selection. - *(providers)* [**breaking**] Move Together, OpenRouter, and Mistral embeddings onto the shared `GenericEmbeddingModel`, with provider-specific endpoint and typed request-shaping hooks. Together now forwards configured embedding dimensions, Mistral maps Codestral Embed dimensions to `output_dimension` while rejecting dimensions for fixed-size models, compatible providers may omit usage without weakening OpenAI's public response type, and Base64 response encoding is rejected before sending because the shared parser accepts numeric vectors. Remove the superseded provider-specific embedding response/data types, Together's API envelope module, and OpenRouter's duplicate `EncodingFormat`. - *(tool)* [**breaking**] Replace the parallel tool-execution APIs with one structured path. Typed tools now implement only `Tool::call(&mut ToolContext, Args) -> Result`; author-facing errors remain typed until private runtime erasure normalizes them into `ToolExecutionError`, `ToolContext` carries inbound values and host-only result metadata, `ToolResult` is the single runtime observation, and `ToolSet::execute` / `ToolServerHandle::execute` are the dispatch surfaces. Event-specific hook action types make invalid event/action combinations unrepresentable. - Tool implementations: retain one typed `type Error` for ordinary `?` propagation and direct-call tests; remove `classify_error`, `call_with_extensions`, and `call_structured`. The optional `map_error` method classifies domain failures at the erased boundary, while its default preserves the source as `Other`. Return refusals through `map_error` with `ToolExecutionError::refused`, and attach host-only result metadata with `ToolContext::insert_result`. diff --git a/crates/rig-core/src/agent/completion.rs b/crates/rig-core/src/agent/completion.rs index b49ce94b6..09657a29d 100644 --- a/crates/rig-core/src/agent/completion.rs +++ b/crates/rig-core/src/agent/completion.rs @@ -5,13 +5,13 @@ use super::runner::AgentRunner; use crate::{ agent::prompt_request::streaming::StreamingPromptRequest, completion::{ - Chat, Completion, CompletionError, CompletionModel, CompletionRequestBuilder, Document, - GetTokenUsage, Message, Prompt, PromptError, TypedPrompt, + Chat, CompletionError, CompletionModel, CompletionRequestBuilder, Document, GetTokenUsage, + Message, Prompt, PromptError, ToolDefinition, TypedPrompt, }, json_utils, message::ToolChoice, - streaming::{StreamingChat, StreamingCompletion, StreamingPrompt}, - tool::server::{ToolRegistrySnapshot, ToolServerHandle}, + streaming::{StreamingChat, StreamingPrompt}, + tool::server::{ToolRegistrySnapshot, ToolServerError, ToolServerHandle}, vector_store::{VectorStoreError, request::VectorSearchRequest}, wasm_compat::WasmCompatSend, }; @@ -224,49 +224,6 @@ pub(crate) fn allowed_tool_names_for_choice( Ok(allowed) } -/// Helper function to build a completion request from agent components. -/// This is used by `Agent::completion()` to preserve the public completion API. -#[allow(clippy::too_many_arguments)] -pub(crate) async fn build_completion_request( - model: &Arc, - prompt: Message, - chat_history: &[Message], - preamble: Option<&str>, - static_context: &[Document], - temperature: Option, - max_tokens: Option, - additional_params: Option<&serde_json::Value>, - record_telemetry_content: bool, - tool_choice: Option<&ToolChoice>, - tool_server_handle: &ToolServerHandle, - dynamic_context: &DynamicContextStore, - output_schema: Option<&schemars::Schema>, -) -> Result, CompletionError> { - Ok(build_prepared_completion_request( - model, - prompt, - chat_history, - preamble, - static_context, - temperature, - max_tokens, - additional_params, - record_telemetry_content, - tool_choice, - tool_server_handle, - dynamic_context, - output_schema, - // The single-shot `Agent::completion()` API has no run loop to consume an - // output-tool call, so it always uses native structured output (#1928). - &OutputMode::Native, - None, - // No agent run loop, so no `CompletionCall` hook to override the request. - None, - ) - .await? - .builder) -} - /// Helper function to build a completion request from agent components while /// preserving the executable Rig tool names sent to the provider. #[allow(clippy::too_many_arguments)] @@ -286,6 +243,8 @@ pub(crate) async fn build_prepared_completion_request( output_schema: Option<&schemars::Schema>, output_mode: &OutputMode, committed_output_tool: Option<&str>, + output_tool_description: Option<&str>, + augment_output_preamble: bool, request_patch: Option<&RequestPatch>, ) -> Result, CompletionError> { // Apply a per-turn request patch (the merged patch from every `CompletionCall` @@ -520,14 +479,17 @@ pub(crate) async fn build_prepared_completion_request( let effective_preamble: Option = { let base = preamble.map(str::to_owned); let instruction = match &resolved_mode { - OutputMode::Tool => output_tool_name.as_deref().map(|name| { - format!( - "When you have gathered enough information to answer, call the `{name}` \ + OutputMode::Tool if augment_output_preamble => { + output_tool_name.as_deref().map(|name| { + format!( + "When you have gathered enough information to answer, call the `{name}` \ tool exactly once with your final answer. Its arguments are the structured \ result and must satisfy the required schema. Do not return the final answer \ as plain text." - ) - }), + ) + }) + } + OutputMode::Tool => None, OutputMode::Prompted => output_schema.map(|schema| { let schema_json = serde_json::to_string(schema.as_value()).unwrap_or_default(); format!( @@ -568,9 +530,11 @@ pub(crate) async fn build_prepared_completion_request( if let (Some(name), Some(schema)) = (&output_tool_name, output_schema) { tooldefs.push(crate::completion::ToolDefinition { name: name.clone(), - description: "Call this tool exactly once with your final answer when you are done. \ - Its arguments are the structured result and must satisfy the output \ - schema." + description: output_tool_description + .unwrap_or( + "Call this tool exactly once with your final answer when you are done. \ + Its arguments are the structured result and must satisfy the output schema.", + ) .to_string(), parameters: schema.clone().to_value(), }); @@ -671,57 +635,66 @@ where M: CompletionModel, { /// Name of the agent used for logging and debugging - pub name: Option, + pub(crate) name: Option, /// Agent description. Primarily useful when using sub-agents as part of an agent workflow and converting agents to other formats. - pub description: Option, + pub(crate) description: Option, /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r) - pub model: Arc, + pub(crate) model: Arc, /// System prompt - pub preamble: Option, + pub(crate) preamble: Option, /// Context documents always available to the agent - pub static_context: Vec, + pub(crate) static_context: Vec, /// Temperature of the model - pub temperature: Option, + pub(crate) temperature: Option, /// Maximum number of tokens for the completion - pub max_tokens: Option, + pub(crate) max_tokens: Option, /// Additional parameters to be passed to the model - pub additional_params: Option, + pub(crate) additional_params: Option, /// Whether to record sensitive request, response, and tool content on GenAI spans. /// /// Defaults to `false`. Enabling this can expose prompts, retrieved context, /// tool results, model responses, and other sensitive or high-cardinality data /// through OpenTelemetry span attributes, which can increase observability /// backend storage and query costs. - pub record_telemetry_content: bool, - pub tool_server_handle: ToolServerHandle, + pub(crate) record_telemetry_content: bool, + pub(crate) tool_server_handle: ToolServerHandle, /// List of vector store, with the sample number - pub dynamic_context: DynamicContextStore, + pub(crate) dynamic_context: DynamicContextStore, /// Whether or not the underlying LLM should be forced to use a tool before providing a response. - pub tool_choice: Option, + pub(crate) tool_choice: Option, /// Default total model-call budget, including the initial call and every /// retry or continuation. `None` uses the implicit budget of one. - pub default_max_turns: Option, + pub(crate) default_max_turns: Option, /// Default hook stack applied to every prompt request and runner created /// from this agent. Empty by default. - pub hooks: HookStack, + pub(crate) hooks: HookStack, /// Optional JSON Schema for structured output. When set, providers that support /// native structured outputs will constrain the model's response to match this schema. - pub output_schema: Option, + pub(crate) output_schema: Option, /// How `output_schema` is enforced — tool call, native structured output, or /// prompt injection (see [`OutputMode`] and issue #1928). - pub output_mode: OutputMode, + pub(crate) output_mode: OutputMode, /// Optional conversation memory backend that loads/saves history per conversation id. - pub memory: Option>, + pub(crate) memory: Option>, /// Optional default conversation id used when none is set per-request. - pub default_conversation_id: Option, + pub(crate) default_conversation_id: Option, } impl Agent where M: CompletionModel, { - /// Returns the name of the agent. - pub(crate) fn name(&self) -> &str { + /// Returns the configured agent name. + pub fn name(&self) -> Option<&str> { + self.name.as_deref() + } + + /// Returns the configured agent description. + pub fn description(&self) -> Option<&str> { + self.description.as_deref() + } + + pub(crate) fn name_or_default(&self) -> &str { self.name.as_deref().unwrap_or(UNKNOWN_AGENT_NAME) } @@ -731,38 +704,16 @@ where pub fn runner(&self, prompt: impl Into) -> AgentRunner { AgentRunner::from_agent(self, prompt) } -} -impl Completion for Agent -where - M: CompletionModel, -{ - async fn completion( + /// Resolve the provider-facing tool definitions available for a prompt. + /// + /// This read-only view does not expose tool dispatch. Agent execution and + /// tool lifecycle hooks remain owned by [`Self::runner`]. + pub async fn tool_definitions( &self, - prompt: impl Into + WasmCompatSend, - chat_history: I, - ) -> Result, CompletionError> - where - I: IntoIterator, - T: Into, - { - let history: Vec = chat_history.into_iter().map(Into::into).collect(); - build_completion_request( - &self.model, - prompt.into(), - &history, - self.preamble.as_deref(), - &self.static_context, - self.temperature, - self.max_tokens, - self.additional_params.as_ref(), - self.record_telemetry_content, - self.tool_choice.as_ref(), - &self.tool_server_handle, - &self.dynamic_context, - self.output_schema.as_ref(), - ) - .await + prompt: Option, + ) -> Result, ToolServerError> { + self.tool_server_handle.get_tool_defs(prompt).await } } @@ -791,7 +742,7 @@ impl Prompt for &Agent where M: CompletionModel + 'static, { - #[tracing::instrument(skip(self, prompt), fields(agent_name = self.name()))] + #[tracing::instrument(skip(self, prompt), fields(agent_name = self.name_or_default()))] fn prompt( &self, prompt: impl Into + WasmCompatSend, @@ -805,7 +756,7 @@ impl Chat for Agent where M: CompletionModel + 'static, { - #[tracing::instrument(skip(self, prompt, chat_history), fields(agent_name = self.name()))] + #[tracing::instrument(skip(self, prompt, chat_history), fields(agent_name = self.name_or_default()))] async fn chat( &self, prompt: impl Into + WasmCompatSend, @@ -824,25 +775,6 @@ where } } -impl StreamingCompletion for Agent -where - M: CompletionModel, -{ - async fn stream_completion( - &self, - prompt: impl Into + WasmCompatSend, - chat_history: I, - ) -> Result, CompletionError> - where - I: IntoIterator + WasmCompatSend, - T: Into, - { - // Reuse the existing completion implementation to build the request - // This ensures streaming and non-streaming use the same request building logic - self.completion(prompt, chat_history).await - } -} - impl StreamingPrompt for Agent where M: CompletionModel + 'static, diff --git a/crates/rig-core/src/agent/mod.rs b/crates/rig-core/src/agent/mod.rs index f1e3d001e..9b9d6bce4 100644 --- a/crates/rig-core/src/agent/mod.rs +++ b/crates/rig-core/src/agent/mod.rs @@ -8,9 +8,10 @@ //! a simple bot with a specific system prompt to a complex RAG system with a set of dynamic //! context documents and tools. //! -//! The [Agent] struct implements the [crate::completion::Completion] and [crate::completion::Prompt] traits, -//! allowing it to be used for generating completions responses and prompts. The [Agent] struct also -//! implements the [crate::completion::Chat] trait, which allows it to be used for generating chat completions. +//! The [Agent] struct implements the runner-backed [crate::completion::Prompt], +//! [crate::completion::TypedPrompt], and [crate::completion::Chat] traits. All +//! agent execution goes through [AgentRunner], so hooks and lifecycle policies +//! cannot be bypassed through a raw agent request builder. //! //! The [AgentBuilder] implements the builder pattern for creating instances of [Agent]. //! It allows configuring the model, preamble, context documents, tools, temperature, and additional parameters @@ -20,7 +21,7 @@ //! ```no_run //! use rig_core::{ //! client::{CompletionClient, ProviderClient}, -//! completion::{Chat, Completion, Prompt}, +//! completion::{Chat, Prompt}, //! providers::openai, //! }; //! @@ -35,24 +36,15 @@ //! .temperature(0.8) //! .build(); //! -//! // Use the agent for completions and prompts +//! // Use the agent for chats and prompts //! // Generate a chat completion response from a prompt and chat history //! let chat_response = agent.chat("Prompt", &mut Vec::::new()).await?; //! //! // Generate a prompt completion response from a simple prompt //! let prompt_response = agent.prompt("Prompt").await?; //! -//! // Generate a completion request builder from a prompt and chat history. The builder -//! // will contain the agent's configuration (i.e.: preamble, context documents, tools, -//! // model parameters, etc.), but these can be overwritten. -//! let completion_req_builder = agent -//! .completion("Prompt", Vec::::new()) -//! .await?; -//! -//! let response = completion_req_builder -//! .temperature(0.9) // Overwrite the agent's temperature -//! .send() -//! .await?; +//! // Per-run overrides stay inside the hook-aware runner. +//! let response = agent.runner("Prompt").temperature(0.9).run().await?; //! # Ok(()) //! # } //! ``` diff --git a/crates/rig-core/src/agent/prompt_request/mod.rs b/crates/rig-core/src/agent/prompt_request/mod.rs index 505d3b51e..543e33955 100644 --- a/crates/rig-core/src/agent/prompt_request/mod.rs +++ b/crates/rig-core/src/agent/prompt_request/mod.rs @@ -42,6 +42,91 @@ macro_rules! forward_prompt_setters { self } + /// Override the agent preamble for this request. + pub fn preamble(mut self, preamble: impl Into) -> Self { + self.$recv = self.$recv.preamble(preamble); + self + } + + /// Remove the agent's configured preamble for this request. + pub fn without_preamble(mut self) -> Self { + self.$recv = self.$recv.without_preamble(); + self + } + + /// Append one static context document for this request. + pub fn document(mut self, document: crate::completion::Document) -> Self { + self.$recv = self.$recv.document(document); + self + } + + /// Append static context documents for this request. + pub fn documents( + mut self, + documents: impl IntoIterator, + ) -> Self { + self.$recv = self.$recv.documents(documents); + self + } + + /// Override the model temperature for this request. + pub fn temperature(mut self, temperature: f64) -> Self { + self.$recv = self.$recv.temperature(temperature); + self + } + + /// Remove the agent's configured temperature for this request. + pub fn without_temperature(mut self) -> Self { + self.$recv = self.$recv.without_temperature(); + self + } + + /// Override the maximum completion token count for this request. + pub fn max_tokens(mut self, max_tokens: u64) -> Self { + self.$recv = self.$recv.max_tokens(max_tokens); + self + } + + /// Remove the agent's configured maximum token count for this request. + pub fn without_max_tokens(mut self) -> Self { + self.$recv = self.$recv.without_max_tokens(); + self + } + + /// Shallow-merge object fields into the provider-specific parameters + /// for this request. Later fields win. + pub fn merge_additional_params( + mut self, + params: serde_json::Map, + ) -> Self { + self.$recv = self.$recv.merge_additional_params(params); + self + } + + /// Replace all provider-specific parameters for this request. + pub fn replace_additional_params(mut self, params: serde_json::Value) -> Self { + self.$recv = self.$recv.replace_additional_params(params); + self + } + + /// Remove the agent's configured provider-specific parameters for this request. + pub fn without_additional_params(mut self) -> Self { + self.$recv = self.$recv.without_additional_params(); + self + } + + /// Override the tool-choice policy for this request. + pub fn tool_choice(mut self, tool_choice: crate::message::ToolChoice) -> Self { + self.$recv = self.$recv.tool_choice(tool_choice); + self + } + + /// Remove the agent's configured tool-choice policy for this request. + pub fn without_tool_choice(mut self) -> Self { + self.$recv = self.$recv.without_tool_choice(); + self + } + /// Opt in or out of recording sensitive request, response, and tool /// content on GenAI telemetry spans for this request. /// @@ -297,6 +382,10 @@ pub struct PromptResponse { /// Where [`output`](Self::output) is the concatenated text, this preserves /// the individual content parts (text, reasoning, images, …). pub content: OneOrMany, + /// Number of synthetic output-tool calls in the turn that finalized this + /// response. Kept crate-private because it is runner bookkeeping rather + /// than provider-facing response content. + output_tool_calls: usize, } /// Serde shadow for [`PromptResponse`]. `content` is an `Option` here so runs @@ -315,6 +404,8 @@ struct PromptResponseRepr { messages: Option>, #[serde(default)] content: Option>, + #[serde(skip)] + output_tool_calls: usize, } impl From for PromptResponse { @@ -328,6 +419,7 @@ impl From for PromptResponse { completion_calls: repr.completion_calls, messages: repr.messages, content, + output_tool_calls: repr.output_tool_calls, } } } @@ -340,6 +432,7 @@ impl From for PromptResponseRepr { completion_calls: response.completion_calls, messages: response.messages, content: Some(response.content), + output_tool_calls: response.output_tool_calls, } } } @@ -359,6 +452,7 @@ impl PromptResponse { usage, completion_calls: Vec::new(), messages: None, + output_tool_calls: 0, } } @@ -384,6 +478,15 @@ impl PromptResponse { self } + pub(crate) fn with_output_tool_calls(mut self, count: usize) -> Self { + self.output_tool_calls = count; + self + } + + pub(crate) fn output_tool_calls(&self) -> usize { + self.output_tool_calls + } + /// The concatenated assistant text for the final turn. pub fn output(&self) -> &str { &self.output @@ -1101,6 +1204,18 @@ mod tests { assert_eq!(response.requests(), 2); } + #[test] + fn prompt_response_output_tool_marker_is_never_serialized() { + let response = PromptResponse::new("ok", usage(1, 2)).with_output_tool_calls(3); + + let value = serde_json::to_value(&response).expect("serialize prompt response"); + assert!(value.get("output_tool_calls").is_none()); + + let decoded: PromptResponse = + serde_json::from_value(value).expect("deserialize prompt response"); + assert_eq!(decoded.output_tool_calls(), 0); + } + #[test] fn prompt_response_deserializes_pre_monoid_null_usage_format() { // Fixture captured from rig before CompletionCall.usage dropped its diff --git a/crates/rig-core/src/agent/prompt_request/streaming.rs b/crates/rig-core/src/agent/prompt_request/streaming.rs index 621be9b8f..d0c7872cb 100644 --- a/crates/rig-core/src/agent/prompt_request/streaming.rs +++ b/crates/rig-core/src/agent/prompt_request/streaming.rs @@ -284,12 +284,10 @@ where M: CompletionModel + 'static, ::StreamingResponse: WasmCompatSend + GetTokenUsage, { - /// Create a new StreamingPromptRequest from an agent WITHOUT the agent's - /// default hooks. Use [`from_agent`](Self::from_agent) to include them. + /// Create a new `StreamingPromptRequest` from an agent, including its + /// default hooks. pub fn new(agent: Arc>, prompt: impl Into) -> StreamingPromptRequest { - let mut runner = AgentRunner::from_agent(agent.as_ref(), prompt); - runner.hooks = HookStack::new(); - StreamingPromptRequest { runner } + Self::from_agent(agent.as_ref(), prompt) } /// Create a new StreamingPromptRequest from an agent, cloning the agent's @@ -451,6 +449,15 @@ pub(crate) fn streaming_error_into_prompt(err: StreamingError) -> PromptError { } } +pub(crate) fn store_error_usage(runner: &AgentRunner, run: &AgentRun) +where + M: CompletionModel, +{ + if let Some(usage) = &runner.error_usage { + *usage.lock().unwrap_or_else(|error| error.into_inner()) = run.usage(); + } +} + /// The single agent drive loop, shared by the blocking and streaming surfaces. /// /// Owns the medium-independent loop — `next_step` dispatch, the `CompletionCall` @@ -485,6 +492,7 @@ where let step = match run.next_step() { Ok(step) => step, Err(err) => { + store_error_usage(&runner, &run); yield Err(Box::new(err).into()); break 'outer; } @@ -501,6 +509,7 @@ where let request_patch = match resolve_completion_call(&runner.hooks, &hook_ctx, &prompt, &history, turn).await { CompletionCallOutcome::Terminate(reason) => { + store_error_usage(&runner, &run); yield Err(StreamingError::Prompt(Box::new(run.cancel_error(reason)))); break 'outer; } @@ -537,12 +546,15 @@ where runner.output_schema.as_ref(), &runner.output_mode, committed_output_tool.as_deref(), + runner.output_tool_description.as_deref(), + runner.augment_output_preamble, request_patch.as_ref(), ) .await { Ok(prepared) => prepared, Err(err) => { + store_error_usage(&runner, &run); yield Err(err.into()); break 'outer; } @@ -564,25 +576,27 @@ where &agent_span, prompt, ); - let mut errored = false; + let mut turn_error = None; while let Some(item) = turn_stream.next().await { match item { Ok(item) => yield Ok(DriveItem::Item(item)), Err(err) => { - errored = true; - yield Err(err); + turn_error = Some(err); break; } } } drop(turn_stream); - if errored { + if let Some(err) = turn_error { + store_error_usage(&runner, &run); + yield Err(err); break 'outer; } pending_tool_snapshot = Some(turn_tool_snapshot); } AgentRunStep::CallTools { calls } => { let Some(tool_snapshot) = pending_tool_snapshot.take() else { + store_error_usage(&runner, &run); yield Err(StreamingError::Completion(CompletionError::ResponseError( "agent requested tool execution without a prepared registry snapshot" .to_string(), @@ -596,19 +610,20 @@ where calls, tool_snapshot, ); - let mut errored = false; + let mut tool_error = None; while let Some(item) = tool_stream.next().await { match item { Ok(item) => yield Ok(DriveItem::Item(item)), Err(err) => { - errored = true; - yield Err(err); + tool_error = Some(err); break; } } } drop(tool_stream); - if errored { + if let Some(err) = tool_error { + store_error_usage(&runner, &run); + yield Err(err); break 'outer; } } @@ -1583,6 +1598,45 @@ mod migrated_tests { use tracing_subscriber::layer::{Context, SubscriberExt}; use tracing_subscriber::{Layer, Registry, registry::LookupSpan}; + struct StopAgentStreamingBeforeCompletion; + + impl AgentHook for StopAgentStreamingBeforeCompletion { + async fn on_completion_call( + &self, + _ctx: &HookContext, + _event: crate::agent::CompletionCallEvent<'_>, + ) -> crate::agent::CompletionCallAction { + crate::agent::CompletionCallAction::stop("agent streaming stopped") + } + } + + #[tokio::test] + async fn public_streaming_request_constructor_preserves_agent_hooks() { + let model = MockCompletionModel::from_stream_turns([[ + MockStreamEvent::text("should not run"), + MockStreamEvent::final_response(Usage::new()), + ]]); + let agent = Arc::new( + AgentBuilder::new(model.clone()) + .add_hook(StopAgentStreamingBeforeCompletion) + .build(), + ); + + let mut stream = StreamingPromptRequest::new(agent, "go").await; + let error = stream + .try_next() + .await + .expect_err("the configured agent hook should terminate the stream"); + + assert!(matches!( + error, + StreamingError::Prompt(error) + if matches!(*error, PromptError::PromptCancelled { ref reason, .. } + if reason == "agent streaming stopped") + )); + assert_eq!(model.request_count(), 0); + } + #[test] fn finalize_streamed_choice_surfaces_output_over_tool_call_and_prose() { use crate::message::{ToolCall, ToolFunction}; diff --git a/crates/rig-core/src/agent/run/mod.rs b/crates/rig-core/src/agent/run/mod.rs index f40babfa6..360da6e46 100644 --- a/crates/rig-core/src/agent/run/mod.rs +++ b/crates/rig-core/src/agent/run/mod.rs @@ -22,9 +22,16 @@ //! carries no cross-version stability guarantee yet: resume with the same rig //! version that suspended the run. //! -//! [`crate::completion::Prompt::prompt`] on [`crate::agent::Agent`] drives -//! this machine internally; the same machine can be driven by hand for custom -//! control flow: +//! `AgentRun` deliberately contains no model, tool registry, memory backend, or +//! hook stack. Hand-driving it is a low-level provider integration: the caller +//! owns all IO and any lifecycle policy. To execute a configured [`Agent`](crate::agent::Agent) +//! with its hooks, tools, retrieval, and memory, use +//! [`Agent::runner`](crate::agent::Agent::runner); constructing an `AgentRun` +//! directly is not an alternate way to execute an `Agent`. +//! +//! [`crate::completion::Prompt::prompt`] and +//! [`Agent::runner`](crate::agent::Agent::runner) drive this machine internally; +//! the same machine can be driven by hand for custom provider control flow: //! //! ```rust,no_run //! use rig_core::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; @@ -611,6 +618,16 @@ impl AgentRun { _ => None, }) { + let output_tool_calls = items + .iter() + .filter(|item| { + matches!( + item, + AssistantContent::ToolCall(tc) + if tc.function.name == output_tool_name + ) + }) + .count(); let args = tool_call.function.arguments.clone(); let tool_call_id = tool_call.id.clone(); let output = json_utils::serialize_json_value(&args); @@ -658,7 +675,8 @@ impl AgentRun { let mut response = PromptResponse::new(output, self.usage) .with_messages(self.new_messages.clone()) - .with_completion_calls(self.completion_calls.clone()); + .with_completion_calls(self.completion_calls.clone()) + .with_output_tool_calls(output_tool_calls); if let Some(content) = final_content { response = response.with_content(content); } @@ -979,6 +997,50 @@ impl AgentRun { } } + /// Discard the pending invalid tool call without marking the turn as + /// recovered. + /// + /// This is the extractor driver's fallback after every invalid-tool hook + /// has declined to act. Keeping it distinct from [`InvalidToolCallAction::Skip`] + /// preserves the extractor's legacy response semantics: unrelated calls + /// disappear, a sibling output call can still finalize the turn, and raw + /// response observers still receive the provider response. + pub(crate) fn ignore_invalid_tool_call(&mut self) -> Result { + let mut resolving = match std::mem::replace(&mut self.state, RunState::Failed) { + RunState::ResolvingToolCalls(resolving) => resolving, + other => { + self.state = other; + return Err(self.protocol_violation( + "ignore_invalid_tool_call called without a pending invalid tool call", + )); + } + }; + + match resolving.items.get(resolving.next_index) { + Some(AssistantContent::ToolCall(tool_call)) + if !resolving + .allowed_tool_names + .contains(&tool_call.function.name) => {} + _ => { + self.state = RunState::ResolvingToolCalls(resolving); + return Err(self.protocol_violation( + "ignore_invalid_tool_call called without a pending invalid tool call", + )); + } + } + + resolving.items.remove(resolving.next_index); + resolving.has_tool_calls = resolving + .items + .iter() + .any(|item| matches!(item, AssistantContent::ToolCall(_))); + if resolving.items.is_empty() { + resolving.items.push(AssistantContent::text("")); + } + self.state = RunState::ResolvingToolCalls(resolving); + self.advance_resolution() + } + /// Feed the tool results for the pending [`AgentRunStep::CallTools`]. /// /// Results may be in any order; they are appended as a single user diff --git a/crates/rig-core/src/agent/runner.rs b/crates/rig-core/src/agent/runner.rs index 79ed67e85..9b8cbf73e 100644 --- a/crates/rig-core/src/agent/runner.rs +++ b/crates/rig-core/src/agent/runner.rs @@ -26,7 +26,7 @@ //! ``` use std::sync::{ - Arc, + Arc, Mutex, atomic::{AtomicU64, Ordering}, }; @@ -54,7 +54,7 @@ use super::{ }, }; use crate::{ - completion::{CompletionError, CompletionModel, Document, Message, PromptError}, + completion::{CompletionError, CompletionModel, Document, Message, PromptError, Usage}, json_utils, memory::ConversationMemory, message::{ToolCall, ToolChoice, UserContent}, @@ -66,6 +66,13 @@ use crate::{ use super::UNKNOWN_AGENT_NAME; +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) enum UnhandledInvalidToolCallPolicy { + #[default] + Fail, + IgnoreForExtractor, +} + /// Build the per-turn `chat` span shared by both turn sources. /// /// The span *name* must be a string literal — `tracing` bakes it into static @@ -190,10 +197,15 @@ where pub(crate) tool_choice: Option, pub(crate) output_schema: Option, pub(crate) output_mode: OutputMode, + pub(crate) output_tool_name: Option, + pub(crate) output_tool_description: Option, + pub(crate) augment_output_preamble: bool, + pub(crate) unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy, pub(crate) concurrency: usize, pub(crate) memory: Option>, pub(crate) conversation_id: Option, pub(crate) hooks: HookStack, + pub(crate) error_usage: Option>>, } impl AgentRunner @@ -222,10 +234,15 @@ where tool_choice: agent.tool_choice.clone(), output_schema: agent.output_schema.clone(), output_mode: agent.output_mode.clone(), + output_tool_name: None, + output_tool_description: None, + augment_output_preamble: true, + unhandled_invalid_tool_call_policy: UnhandledInvalidToolCallPolicy::Fail, concurrency: 1, memory: agent.memory.clone(), conversation_id: agent.default_conversation_id.clone(), hooks: agent.hooks.clone(), + error_usage: None, } } @@ -273,6 +290,124 @@ where self } + /// Override the agent preamble for this run. + pub fn preamble(mut self, preamble: impl Into) -> Self { + self.preamble = Some(preamble.into()); + self + } + + /// Remove the agent's configured preamble for this run. + pub fn without_preamble(mut self) -> Self { + self.preamble = None; + self + } + + /// Append one static context document for this run. + pub fn document(mut self, document: Document) -> Self { + self.static_context.push(document); + self + } + + /// Append static context documents for this run. + pub fn documents(mut self, documents: impl IntoIterator) -> Self { + self.static_context.extend(documents); + self + } + + /// Override the model temperature for this run. + pub fn temperature(mut self, temperature: f64) -> Self { + self.temperature = Some(temperature); + self + } + + /// Remove the agent's configured temperature for this run. + pub fn without_temperature(mut self) -> Self { + self.temperature = None; + self + } + + /// Override the maximum completion token count for this run. + pub fn max_tokens(mut self, max_tokens: u64) -> Self { + self.max_tokens = Some(max_tokens); + self + } + + /// Remove the agent's configured maximum token count for this run. + pub fn without_max_tokens(mut self) -> Self { + self.max_tokens = None; + self + } + + /// Shallow-merge object fields into the provider-specific parameters for + /// this run. Later fields win. A non-object baseline is replaced by the + /// supplied object. A later completion-call hook patch has final + /// precedence: object values shallow-merge, while a non-object on either + /// side causes wholesale replacement by the hook value. + pub fn merge_additional_params( + mut self, + params: serde_json::Map, + ) -> Self { + let params = serde_json::Value::Object(params); + self.additional_params = Some(match self.additional_params.take() { + Some(baseline) if baseline.is_object() => crate::json_utils::merge(baseline, params), + _ => params, + }); + self + } + + /// Replace all provider-specific parameters for this run. A later + /// completion-call hook patch has final precedence: object values + /// shallow-merge, while a non-object on either side causes wholesale + /// replacement by the hook value. + pub fn replace_additional_params(mut self, params: serde_json::Value) -> Self { + self.additional_params = Some(params); + self + } + + /// Remove the agent's configured provider-specific parameters for this run. + /// A later completion-call hook may still supply its own parameters. + pub fn without_additional_params(mut self) -> Self { + self.additional_params = None; + self + } + + /// Override the tool-choice policy for this run. + pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self { + self.tool_choice = Some(tool_choice); + self + } + + /// Remove the agent's configured tool-choice policy for this run. + pub fn without_tool_choice(mut self) -> Self { + self.tool_choice = None; + self + } + + /// Configure the synthetic tool used by an internal Tool-output flow. + pub(crate) fn output_tool( + mut self, + name: impl Into, + description: impl Into, + augment_preamble: bool, + ) -> Self { + self.output_tool_name = Some(name.into()); + self.output_tool_description = Some(description.into()); + self.augment_output_preamble = augment_preamble; + self + } + + /// Ignore invalid tool calls when every registered hook declines to act. + /// + /// This is an internal compatibility policy for extractors, whose legacy + /// transport treated every non-`submit` call as irrelevant response + /// content. Hooks still receive the invalid-call event first and retain + /// full control over recovery or termination. + pub(crate) fn ignore_unhandled_invalid_tool_calls(mut self) -> Self { + self.unhandled_invalid_tool_call_policy = + UnhandledInvalidToolCallPolicy::IgnoreForExtractor; + self + } + /// Opt in or out of recording sensitive request, response, and tool content /// on GenAI telemetry spans for this run. /// @@ -341,14 +476,18 @@ where /// memory-loaded history). Delegates to [`build_agent_run`] — the single /// construction site shared with the streaming driver. pub(crate) fn build_run(&self, history_override: Option>) -> AgentRun { - build_agent_run( + let run = build_agent_run( self.prompt.clone(), self.max_turns, self.max_invalid_tool_call_retries, self.output_schema.as_ref(), history_override.or_else(|| self.chat_history.clone()), self.tool_choice.clone(), - ) + ); + match &self.output_tool_name { + Some(name) => run.with_output_tool_name(name.clone()), + None => run, + } } } @@ -795,9 +934,18 @@ where let action = runner .hooks .on_invalid_tool_call(hook_ctx, &context) - .await - .unwrap_or_else(InvalidToolCallAction::fail); - outcome = match run.resolve_invalid_tool_call(action) { + .await; + let resolution = match action { + Some(action) => run.resolve_invalid_tool_call(action), + None + if runner.unhandled_invalid_tool_call_policy + == UnhandledInvalidToolCallPolicy::IgnoreForExtractor => + { + run.ignore_invalid_tool_call() + } + None => run.resolve_invalid_tool_call(InvalidToolCallAction::fail()), + }; + outcome = match resolution { Ok(outcome) => outcome, Err(err) => { yield Err(Box::new(err).into()); @@ -909,6 +1057,19 @@ impl AgentRunner where M: CompletionModel, { + pub(crate) async fn run_with_error_usage( + mut self, + ) -> (Result, Usage) { + let usage = Arc::new(Mutex::new(Usage::new())); + self.error_usage = Some(usage.clone()); + let result = self.run().await; + let observed = result.as_ref().map_or_else( + |_| *usage.lock().unwrap_or_else(|error| error.into_inner()), + |response| response.usage, + ); + (result, observed) + } + /// Drive the agent loop to completion, returning the aggregated /// [`PromptResponse`]. Hooks fire at every observable point; the first hook /// to terminate cancels the run. @@ -988,7 +1149,8 @@ mod tests { use crate::{ agent::{AgentBuilder, AgentHook, HookContext, ToolResultAction, ToolResultEvent}, - completion::CompletionModel, + completion::{CompletionModel, Document}, + message::ToolChoice, test_utils::{MockCompletionModel, MockStreamEvent, MockTurn}, tool::{Tool, ToolContext, ToolErrorKind, ToolExecutionError}, }; @@ -1113,6 +1275,195 @@ mod tests { } } + #[test] + fn agent_exposes_read_only_name_and_description() { + let named = AgentBuilder::new(MockCompletionModel::text("done")) + .name("researcher") + .description("Finds evidence") + .build(); + assert_eq!(named.name(), Some("researcher")); + assert_eq!(named.description(), Some("Finds evidence")); + + let unnamed = AgentBuilder::new(MockCompletionModel::text("done")).build(); + assert_eq!(unnamed.name(), None); + assert_eq!(unnamed.description(), None); + } + + #[tokio::test] + async fn runner_applies_per_run_request_overrides() { + let model = MockCompletionModel::text("done"); + AgentBuilder::new(model.clone()) + .preamble("baseline preamble") + .context("baseline document") + .temperature(0.1) + .max_tokens(10) + .additional_params(json!({"baseline": true})) + .build() + .runner("go") + .preamble("run preamble") + .document(Document { + id: "run-one".into(), + text: "first run document".into(), + additional_props: Default::default(), + }) + .documents([Document { + id: "run-two".into(), + text: "second run document".into(), + additional_props: Default::default(), + }]) + .temperature(0.7) + .max_tokens(42) + .replace_additional_params(json!({"override": true})) + .tool_choice(ToolChoice::None) + .run() + .await + .expect("runner request should succeed"); + + let requests = model.requests(); + let request = requests.first().expect("one request"); + assert!(request.chat_history.iter().any( + |message| matches!(message, crate::completion::Message::System { content } if content == "run preamble") + )); + assert!( + request + .documents + .iter() + .any(|document| document.text == "baseline document") + ); + assert!( + request + .documents + .iter() + .any(|document| document.id == "run-one") + ); + assert!( + request + .documents + .iter() + .any(|document| document.id == "run-two") + ); + assert_eq!(request.temperature, Some(0.7)); + assert_eq!(request.max_tokens, Some(42)); + assert_eq!(request.additional_params, Some(json!({"override": true}))); + assert_eq!(request.tool_choice, Some(ToolChoice::None)); + } + + #[tokio::test] + async fn runner_can_merge_additional_params_into_the_baseline() { + let model = MockCompletionModel::text("done"); + AgentBuilder::new(model.clone()) + .additional_params(json!({"baseline": true, "winner": "baseline"})) + .build() + .runner("go") + .merge_additional_params( + json!({"override": true, "winner": "runner"}) + .as_object() + .expect("object") + .clone(), + ) + .run() + .await + .expect("runner request should succeed"); + + assert_eq!( + model + .requests() + .first() + .expect("one request") + .additional_params, + Some(json!({"baseline": true, "override": true, "winner": "runner"})) + ); + } + + #[tokio::test] + async fn runner_can_replace_additional_params_wholesale() { + let model = MockCompletionModel::text("done"); + AgentBuilder::new(model.clone()) + .additional_params(json!({"baseline": true})) + .build() + .runner("go") + .replace_additional_params(json!({"replacement": true})) + .run() + .await + .expect("runner request should succeed"); + + let requests = model.requests(); + let request = requests.first().expect("one request"); + assert_eq!( + request.additional_params, + Some(json!({"replacement": true})) + ); + } + + #[tokio::test] + async fn runner_can_clear_configured_request_defaults() { + let model = MockCompletionModel::text("done"); + AgentBuilder::new(model.clone()) + .preamble("baseline") + .temperature(0.1) + .max_tokens(10) + .additional_params(json!({"baseline": true})) + .tool_choice(ToolChoice::Required) + .build() + .runner("go") + .without_preamble() + .without_temperature() + .without_max_tokens() + .without_additional_params() + .without_tool_choice() + .run() + .await + .expect("runner request should succeed"); + + let requests = model.requests(); + let request = requests.first().expect("one request"); + assert!( + !request + .chat_history + .iter() + .any(|message| matches!(message, crate::completion::Message::System { .. })) + ); + assert_eq!(request.temperature, None); + assert_eq!(request.max_tokens, None); + assert_eq!(request.additional_params, None); + assert_eq!(request.tool_choice, None); + } + + #[tokio::test] + async fn direct_completion_model_requests_are_intentionally_hook_free() { + #[derive(Clone)] + struct CountCompletionCalls(Arc); + + impl AgentHook for CountCompletionCalls + where + M: CompletionModel, + { + async fn on_completion_call( + &self, + _ctx: &HookContext, + _event: crate::agent::CompletionCallEvent<'_>, + ) -> crate::agent::CompletionCallAction { + self.0.fetch_add(1, Ordering::SeqCst); + crate::agent::CompletionCallAction::Continue + } + } + + let model = MockCompletionModel::text("raw response"); + let calls = Arc::new(AtomicUsize::new(0)); + let _agent = AgentBuilder::new(model.clone()) + .add_hook(CountCompletionCalls(calls.clone())) + .build(); + + model + .completion_request("raw request") + .send() + .await + .expect("direct model request should succeed"); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!(model.request_count(), 1); + } + #[tokio::test] async fn blocking_and_streaming_preserve_raw_failure_while_rewriting_presentation() { let blocking = Results::default(); @@ -5904,14 +6255,12 @@ mod migrated_tests { ["add"], "active_tools narrows the advertised set to `add` (drops `subtract`)" ); - // additional_params is shallow-merged: the agent baseline survives and - // the override's key is added. + // The runner replaces the agent baseline, then the hook shallow-merges + // last and therefore wins conflicts. let params = req.additional_params.as_ref().expect("additional_params"); - assert_eq!( - params.get("baseline").and_then(|v| v.as_str()), - Some("keep") - ); + assert_eq!(params.get("runner").and_then(|v| v.as_str()), Some("keep")); assert_eq!(params.get("injected").and_then(|v| v.as_bool()), Some(true)); + assert!(params.get("baseline").is_none()); } let blocking_model = MockCompletionModel::from_turns([MockTurn::text("done")]); @@ -5926,6 +6275,7 @@ mod migrated_tests { .add_hook(PatchRequestHook) .build() .runner("go") + .replace_additional_params(json!({"runner": "keep", "injected": false})) .max_turns(2) .run() .await @@ -5949,6 +6299,7 @@ mod migrated_tests { .add_hook(PatchRequestHook) .build() .runner("go") + .replace_additional_params(json!({"runner": "keep", "injected": false})) .max_turns(2) .stream() .await; diff --git a/crates/rig-core/src/completion/mod.rs b/crates/rig-core/src/completion/mod.rs index d06f35348..04fba7b3f 100644 --- a/crates/rig-core/src/completion/mod.rs +++ b/crates/rig-core/src/completion/mod.rs @@ -7,9 +7,12 @@ //! - [`Prompt`] sends one user prompt and returns assistant text. //! - [`Chat`] sends a prompt with existing history and returns assistant text. //! - [`TypedPrompt`] requests structured output and deserializes it into a Rust type. -//! - [`Completion`] exposes a request builder for call-site overrides. //! - [`CompletionModel`] is the provider-facing trait implemented by completion models. //! +//! Agent execution always goes through [`AgentRunner`](crate::agent::AgentRunner), +//! including the high-level traits above. Call [`CompletionModel::completion`] +//! directly only when intentionally making a raw, hook-free provider request. +//! //! `CompletionRequest` is Rig's canonical request representation. Provider modules //! translate it into provider-specific request bodies and convert responses back into //! [`CompletionResponse`]. diff --git a/crates/rig-core/src/completion/request.rs b/crates/rig-core/src/completion/request.rs index 29542cacf..74a7db3ab 100644 --- a/crates/rig-core/src/completion/request.rs +++ b/crates/rig-core/src/completion/request.rs @@ -445,30 +445,6 @@ pub trait TypedPrompt: WasmCompatSend + WasmCompatSync { T: schemars::JsonSchema + DeserializeOwned + WasmCompatSend; } -/// Trait defining a low-level LLM completion interface -pub trait Completion { - /// Generates a completion request builder for the given `prompt` and `chat_history`. - /// This function is meant to be called by the user to further customize the - /// request at prompt time before sending it. - /// - /// ❗IMPORTANT: The type that implements this trait might have already - /// populated fields in the builder (the exact fields depend on the type). - /// For fields that have already been set by the model, calling the corresponding - /// method on the builder will overwrite the value set by the model. - /// - /// For example, the request builder returned by [`Agent::completion`](crate::agent::Agent::completion) will already - /// contain the `preamble` provided when creating the agent. - fn completion( - &self, - prompt: impl Into + WasmCompatSend, - chat_history: I, - ) -> impl std::future::Future, CompletionError>> - + WasmCompatSend - where - I: IntoIterator + WasmCompatSend, - T: Into; -} - /// General completion response struct that contains the high-level completion choice /// and the raw response. The completion choice contains one or more assistant content. #[derive(Debug)] diff --git a/crates/rig-core/src/extractor.rs b/crates/rig-core/src/extractor.rs index 19a20d145..e7e2c422a 100644 --- a/crates/rig-core/src/extractor.rs +++ b/crates/rig-core/src/extractor.rs @@ -31,15 +31,13 @@ use std::marker::PhantomData; -use schemars::{JsonSchema, schema_for}; +use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use serde_json::json; use crate::{ - agent::{Agent, AgentBuilder, WithBuilderTools}, - completion::{Completion, CompletionError, CompletionModel, Usage}, - message::{AssistantContent, Message, ToolCall, ToolChoice, ToolFunction}, - tool::Tool, + agent::{Agent, AgentBuilder, AgentHook, OutputMode}, + completion::{CompletionError, CompletionModel, PromptError, Usage}, + message::{Message, ToolChoice}, vector_store::VectorStoreIndexDyn, wasm_compat::{WasmCompatSend, WasmCompatSync}, }; @@ -65,6 +63,9 @@ pub enum ExtractionError { #[error("CompletionError: {0}")] CompletionError(#[from] CompletionError), + + #[error("PromptError: {0}")] + PromptError(#[from] PromptError), } /// Extractor for structured data from text @@ -201,62 +202,43 @@ where text: &Message, messages: &[Message], ) -> (Result, Usage) { - let completion = async { self.agent.completion(text, messages).await?.send().await }; - let response = match completion.await { + let (result, error_usage) = self + .agent + .runner(text.clone()) + .history(messages.iter().cloned()) + .max_turns(1) + .output_tool( + SUBMIT_TOOL_NAME, + "Submit the structured data you extracted from the provided text.", + false, + ) + .ignore_unhandled_invalid_tool_calls() + .run_with_error_usage() + .await; + let response = match result { Ok(response) => response, - Err(e) => return (Err(e.into()), Usage::new()), + Err(PromptError::CompletionError(e)) => { + return (Err(ExtractionError::CompletionError(e)), error_usage); + } + Err(e) => return (Err(e.into()), error_usage), }; let usage = response.usage; - if !response.choice.iter().any(|x| { - let AssistantContent::ToolCall(ToolCall { - function: ToolFunction { name, .. }, - .. - }) = x - else { - return false; - }; - - name == SUBMIT_TOOL_NAME - }) { + let submissions = response.output_tool_calls(); + if submissions == 0 { tracing::warn!( "The submit tool was not called. If this happens more than once, please ensure the model you are using is powerful enough to reliably call tools." ); + return (Err(ExtractionError::NoData), usage); } - - let arguments = response - .choice - .into_iter() - // We filter tool calls to look for submit tool calls - .filter_map(|content| { - if let AssistantContent::ToolCall(ToolCall { - function: ToolFunction { arguments, name }, - .. - }) = content - { - if name == SUBMIT_TOOL_NAME { - Some(arguments) - } else { - None - } - } else { - None - } - }) - .collect::>(); - - if arguments.len() > 1 { + if submissions > 1 { tracing::warn!( "Multiple submit calls detected, using the first one. Providers / agents should only ensure one submit call." ); } - let Some(raw_data) = arguments.into_iter().next() else { - return (Err(ExtractionError::NoData), usage); - }; - ( - serde_json::from_value(raw_data).map_err(ExtractionError::from), + serde_json::from_str(&response.output).map_err(ExtractionError::from), usage, ) } @@ -268,7 +250,7 @@ where M: CompletionModel, T: JsonSchema + for<'a> Deserialize<'a> + Serialize + WasmCompatSend + WasmCompatSync + 'static, { - agent_builder: AgentBuilder, + agent_builder: AgentBuilder, _t: PhantomData, retries: Option, } @@ -287,11 +269,9 @@ where Use the `submit` function to submit the structured data.\n\ Be sure to fill out every field and ALWAYS CALL THE `submit` function, even with default values!!!. ") - .tool(SubmitTool:: {_t: PhantomData}) + .output_schema::() .tool_choice(ToolChoice::Required) - // The extractor already implements its own tool-based output via - // `SubmitTool`; opt out of the agent's OutputMode routing (#1928). - .output_mode(crate::agent::OutputMode::Native), + .output_mode(OutputMode::Tool), retries: None, _t: PhantomData, } @@ -334,6 +314,15 @@ where self } + /// Add a lifecycle hook to every extraction attempt. + pub fn add_hook(mut self, hook: H) -> Self + where + H: AgentHook + 'static, + { + self.agent_builder = self.agent_builder.add_hook(hook); + self + } + /// Build the Extractor pub fn build(self) -> Extractor { Extractor { @@ -357,49 +346,18 @@ where } } -#[derive(Deserialize, Serialize)] -struct SubmitTool -where - T: JsonSchema + for<'a> Deserialize<'a> + Serialize + WasmCompatSend + WasmCompatSync, -{ - _t: PhantomData, -} - -#[derive(Debug, thiserror::Error)] -#[error("SubmitError")] -struct SubmitError; - -impl Tool for SubmitTool -where - T: JsonSchema + for<'a> Deserialize<'a> + Serialize + WasmCompatSend + WasmCompatSync + 'static, -{ - const NAME: &'static str = SUBMIT_TOOL_NAME; - type Error = SubmitError; - type Args = T; - type Output = T; - - fn description(&self) -> String { - "Submit the structured data you extracted from the provided text.".to_string() - } - - fn parameters(&self) -> serde_json::Value { - json!(schema_for!(T)) - } - - async fn call( - &self, - _context: &mut crate::tool::ToolContext, - data: Self::Args, - ) -> Result { - Ok(data) - } -} - #[cfg(test)] mod tests { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + use serde_json::json; use super::*; + use crate::agent::{CompletionResponseEvent, HookContext, ObservationAction}; + use crate::message::{AssistantContent, ToolCall, ToolFunction}; use crate::test_utils::{MockCompletionModel, MockTurn}; #[derive(Debug, PartialEq, Deserialize, Serialize, JsonSchema)] @@ -425,6 +383,209 @@ mod tests { MockTurn::tool_call("id1", SUBMIT_TOOL_NAME, json!({ "name": name })) } + fn tool_call(id: &str, name: &str, arguments: serde_json::Value) -> AssistantContent { + AssistantContent::ToolCall(ToolCall::new( + id.to_string(), + ToolFunction::new(name.to_string(), arguments), + )) + } + + #[derive(Clone, Default)] + struct LifecycleCounts { + completion_calls: Arc, + completion_responses: Arc, + model_turns: Arc, + invalid_tool_calls: Arc, + } + + impl AgentHook for LifecycleCounts + where + M: CompletionModel, + { + async fn on_completion_call( + &self, + _ctx: &HookContext, + _event: crate::agent::CompletionCallEvent<'_>, + ) -> crate::agent::CompletionCallAction { + self.completion_calls.fetch_add(1, Ordering::SeqCst); + crate::agent::CompletionCallAction::Continue + } + + async fn on_completion_response( + &self, + _ctx: &HookContext, + _event: CompletionResponseEvent<'_, M>, + ) -> ObservationAction { + self.completion_responses.fetch_add(1, Ordering::SeqCst); + ObservationAction::Continue + } + + async fn on_model_turn_finished( + &self, + _ctx: &HookContext, + _event: crate::agent::ModelTurnFinished<'_>, + ) -> ObservationAction { + self.model_turns.fetch_add(1, Ordering::SeqCst); + ObservationAction::Continue + } + + async fn on_invalid_tool_call( + &self, + _ctx: &HookContext, + _event: &crate::agent::InvalidToolCallContext, + ) -> Option { + self.invalid_tool_calls.fetch_add(1, Ordering::SeqCst); + None + } + } + + struct StopBeforeCompletion; + + impl AgentHook for StopBeforeCompletion + where + M: CompletionModel, + { + async fn on_completion_call( + &self, + _ctx: &HookContext, + _event: crate::agent::CompletionCallEvent<'_>, + ) -> crate::agent::CompletionCallAction { + crate::agent::CompletionCallAction::stop("extractor stopped") + } + } + + #[derive(Clone, Copy)] + enum StopFirstBilledResponseAt { + CompletionResponse, + ModelTurnFinished, + } + + #[derive(Clone)] + struct StopFirstBilledResponse { + phase: StopFirstBilledResponseAt, + calls: Arc, + } + + impl AgentHook for StopFirstBilledResponse + where + M: CompletionModel, + { + async fn on_completion_response( + &self, + _ctx: &HookContext, + _event: CompletionResponseEvent<'_, M>, + ) -> ObservationAction { + if matches!(self.phase, StopFirstBilledResponseAt::CompletionResponse) + && self.calls.fetch_add(1, Ordering::SeqCst) == 0 + { + ObservationAction::stop("stop first billed response") + } else { + ObservationAction::continue_run() + } + } + + async fn on_model_turn_finished( + &self, + _ctx: &HookContext, + _event: crate::agent::ModelTurnFinished<'_>, + ) -> ObservationAction { + if matches!(self.phase, StopFirstBilledResponseAt::ModelTurnFinished) + && self.calls.fetch_add(1, Ordering::SeqCst) == 0 + { + ObservationAction::stop("stop first billed model turn") + } else { + ObservationAction::continue_run() + } + } + } + + struct StopOnInvalidToolCall; + + impl AgentHook for StopOnInvalidToolCall + where + M: CompletionModel, + { + async fn on_invalid_tool_call( + &self, + _ctx: &HookContext, + _event: &crate::agent::InvalidToolCallContext, + ) -> Option { + Some(crate::agent::InvalidToolCallAction::stop( + "unexpected extractor tool call", + )) + } + } + + struct RepairUnexpectedAsSubmit; + + impl AgentHook for RepairUnexpectedAsSubmit + where + M: CompletionModel, + { + async fn on_invalid_tool_call( + &self, + _ctx: &HookContext, + _event: &crate::agent::InvalidToolCallContext, + ) -> Option { + Some(crate::agent::InvalidToolCallAction::repair( + SUBMIT_TOOL_NAME, + )) + } + } + + struct SkipUnexpected; + + impl AgentHook for SkipUnexpected + where + M: CompletionModel, + { + async fn on_invalid_tool_call( + &self, + _ctx: &HookContext, + _event: &crate::agent::InvalidToolCallContext, + ) -> Option { + Some(crate::agent::InvalidToolCallAction::skip( + "ignored by extractor hook", + )) + } + } + + #[tokio::test] + async fn extractor_runs_through_full_response_lifecycle() { + let model = MockCompletionModel::new([submit_turn("John")]); + let counts = LifecycleCounts::default(); + let response = ExtractorBuilder::<_, Person>::new(model.clone()) + .add_hook(counts.clone()) + .build() + .extract("John") + .await + .expect("extraction should succeed"); + + assert_eq!(response.name, "John"); + assert_eq!(model.request_count(), 1); + assert_eq!(counts.completion_calls.load(Ordering::SeqCst), 1); + assert_eq!(counts.completion_responses.load(Ordering::SeqCst), 1); + assert_eq!(counts.model_turns.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn extractor_completion_call_stop_prevents_provider_io() { + let model = MockCompletionModel::new([submit_turn("John")]); + let error = ExtractorBuilder::<_, Person>::new(model.clone()) + .add_hook(StopBeforeCompletion) + .build() + .extract("John") + .await + .expect_err("terminating hook should cancel extraction"); + + assert!(matches!( + error, + ExtractionError::PromptError(PromptError::PromptCancelled { reason, .. }) + if reason == "extractor stopped" + )); + assert_eq!(model.request_count(), 0); + } + #[tokio::test] async fn usage_accumulates_across_failed_attempts() { let model = MockCompletionModel::new([ @@ -446,6 +607,183 @@ mod tests { assert_eq!(response.usage.total_tokens, 15); } + async fn assert_billed_hook_termination_usage(phase: StopFirstBilledResponseAt) { + let model = MockCompletionModel::new([ + submit_turn("ignored").with_usage(usage(10)), + submit_turn("John").with_usage(usage(5)), + ]); + let response = ExtractorBuilder::<_, Person>::new(model) + .retries(1) + .add_hook(StopFirstBilledResponse { + phase, + calls: Arc::new(AtomicUsize::new(0)), + }) + .build() + .extract_with_usage("John") + .await + .expect("second attempt should succeed"); + + assert_eq!(response.data.name, "John"); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn completion_response_hook_termination_preserves_billed_usage() { + assert_billed_hook_termination_usage(StopFirstBilledResponseAt::CompletionResponse).await; + } + + #[tokio::test] + async fn model_turn_finished_hook_termination_preserves_billed_usage() { + assert_billed_hook_termination_usage(StopFirstBilledResponseAt::ModelTurnFinished).await; + } + + #[tokio::test] + async fn unexpected_tool_call_preserves_usage_and_retries() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("unknown", "unexpected", json!({})).with_usage(usage(10)), + submit_turn("John").with_usage(usage(5)), + ]); + + let response = extractor(model, 1) + .extract_with_usage("John") + .await + .expect("second attempt should succeed"); + + assert_eq!(response.data.name, "John"); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn unexpected_tool_call_runs_hooks_before_extractor_fallback() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("unknown", "unexpected", json!({})).with_usage(usage(10)), + submit_turn("John").with_usage(usage(5)), + ]); + let counts = LifecycleCounts::default(); + + let response = ExtractorBuilder::<_, Person>::new(model) + .retries(1) + .add_hook(counts.clone()) + .build() + .extract_with_usage("John") + .await + .expect("deferred invalid call should use extractor fallback"); + + assert_eq!(response.data.name, "John"); + assert_eq!(response.usage.total_tokens, 15); + assert_eq!(counts.invalid_tool_calls.load(Ordering::SeqCst), 1); + assert_eq!(counts.completion_responses.load(Ordering::SeqCst), 2); + assert_eq!(counts.model_turns.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn unexpected_tool_call_hook_can_stop_extraction() { + let model = + MockCompletionModel::new([MockTurn::tool_call("unknown", "unexpected", json!({}))]); + + let error = ExtractorBuilder::<_, Person>::new(model) + .add_hook(StopOnInvalidToolCall) + .build() + .extract("John") + .await + .expect_err("invalid-tool hook should retain control"); + + assert!(matches!( + error, + ExtractionError::PromptError(PromptError::PromptCancelled { reason, .. }) + if reason == "unexpected extractor tool call" + )); + } + + #[tokio::test] + async fn unexpected_tool_call_hook_can_repair_to_submit() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "unknown", + "unexpected", + json!({ "name": "John" }), + )]); + + let response = ExtractorBuilder::<_, Person>::new(model) + .add_hook(RepairUnexpectedAsSubmit) + .build() + .extract("John") + .await + .expect("repaired output-tool call should finalize extraction"); + + assert_eq!(response.name, "John"); + } + + #[tokio::test] + async fn skip_hook_preserves_valid_submit_sibling() { + let turn = MockTurn::from_contents([ + tool_call("unknown", "unexpected", json!({})), + tool_call("submit", SUBMIT_TOOL_NAME, json!({ "name": "John" })), + ]) + .expect("two tool calls"); + let model = MockCompletionModel::new([turn]); + + let response = ExtractorBuilder::<_, Person>::new(model) + .add_hook(SkipUnexpected) + .build() + .extract("John") + .await + .expect("skipping an invalid sibling should preserve submit"); + + assert_eq!(response.name, "John"); + } + + #[tokio::test] + async fn submit_call_wins_over_unexpected_sibling_call() { + let turn = MockTurn::from_contents([ + tool_call("unknown", "unexpected", json!({})), + tool_call("submit", SUBMIT_TOOL_NAME, json!({ "name": "John" })), + ]) + .expect("two tool calls") + .with_usage(usage(7)); + let model = MockCompletionModel::new([turn]); + + let response = extractor(model, 0) + .extract_with_usage("John") + .await + .expect("submit should remain authoritative"); + + assert_eq!(response.data.name, "John"); + assert_eq!(response.usage.total_tokens, 7); + } + + #[tokio::test] + async fn submit_call_wins_before_unexpected_sibling_call() { + let turn = MockTurn::from_contents([ + tool_call("submit", SUBMIT_TOOL_NAME, json!({ "name": "John" })), + tool_call("unknown", "unexpected", json!({})), + ]) + .expect("two tool calls"); + + let response = extractor(MockCompletionModel::new([turn]), 0) + .extract("John") + .await + .expect("an earlier submit should remain authoritative"); + + assert_eq!(response.name, "John"); + } + + #[tokio::test] + async fn multiple_unexpected_calls_surrounding_submit_are_ignored() { + let turn = MockTurn::from_contents([ + tool_call("unknown-before", "unexpected_before", json!({})), + tool_call("submit", SUBMIT_TOOL_NAME, json!({ "name": "John" })), + tool_call("unknown-after", "unexpected_after", json!({})), + ]) + .expect("three tool calls"); + + let response = extractor(MockCompletionModel::new([turn]), 0) + .extract("John") + .await + .expect("unexpected siblings should not displace submit"); + + assert_eq!(response.name, "John"); + } + #[tokio::test] async fn transport_errors_contribute_no_usage() { let model = MockCompletionModel::new([ diff --git a/crates/rig-core/src/prelude.rs b/crates/rig-core/src/prelude.rs index 2ef4617f2..4d5fff05d 100644 --- a/crates/rig-core/src/prelude.rs +++ b/crates/rig-core/src/prelude.rs @@ -30,8 +30,8 @@ pub use crate::agent::Agent; // Completion: prompting, chatting, the model trait, and the core types. pub use crate::completion::{ - Chat, Completion, CompletionError, CompletionModel, Message, Prompt, PromptError, - StructuredOutputError, TypedPrompt, + Chat, CompletionError, CompletionModel, Message, Prompt, PromptError, StructuredOutputError, + TypedPrompt, }; // Streaming counterparts of the blocking `Prompt`/`Chat` traits, plus the items diff --git a/crates/rig-core/src/streaming.rs b/crates/rig-core/src/streaming.rs index fe934936d..128a4b04d 100644 --- a/crates/rig-core/src/streaming.rs +++ b/crates/rig-core/src/streaming.rs @@ -5,14 +5,12 @@ //! The main traits defined in this module are: //! - [StreamingPrompt]: Defines a high-level streaming LLM one-shot prompt interface //! - [StreamingChat]: Defines a high-level streaming LLM chat interface with history -//! - [StreamingCompletion]: Defines a low-level streaming LLM completion interface //! use crate::OneOrMany; use crate::agent::prompt_request::streaming::StreamingPromptRequest; use crate::completion::{ - CompletionError, CompletionModel, CompletionRequestBuilder, CompletionResponse, GetTokenUsage, - Message, Usage, + CompletionError, CompletionModel, CompletionResponse, GetTokenUsage, Message, Usage, }; use crate::message::{ AssistantContent, Reasoning, ReasoningContent, Text, ToolCall, ToolFunction, ToolResult, @@ -21,7 +19,6 @@ use crate::wasm_compat::{WasmCompatSend, WasmCompatSync}; use futures::stream::{AbortHandle, Abortable}; use futures::{Stream, StreamExt}; use serde::{Deserialize, Serialize}; -use std::future::Future; use std::pin::Pin; use std::sync::atomic::AtomicBool; use std::task::{Context, Poll}; @@ -626,19 +623,6 @@ where T: Into; } -/// Trait for low-level streaming completion interface -pub trait StreamingCompletion { - /// Generate a streaming completion from a request - fn stream_completion( - &self, - prompt: impl Into + WasmCompatSend, - chat_history: I, - ) -> impl Future, CompletionError>> - where - I: IntoIterator + WasmCompatSend, - T: Into; -} - // Test module #[cfg(test)] mod tests { diff --git a/crates/rig-core/src/test_utils/model_conformance.rs b/crates/rig-core/src/test_utils/model_conformance.rs index 6d3159bf6..1a70f0b22 100644 --- a/crates/rig-core/src/test_utils/model_conformance.rs +++ b/crates/rig-core/src/test_utils/model_conformance.rs @@ -19,15 +19,15 @@ use serde::{Deserialize, Serialize}; use crate::{ agent::{ - AgentBuilder, AgentHook, CompletionCallAction, CompletionCallEvent, HookContext, - InvalidToolCallAction, MultiTurnStreamItem, NoToolConfig, OutputMode, RequestPatch, - StreamingError, ToolCall as ToolCallEvent, ToolCallAction, ToolResultAction, - ToolResultEvent, + AgentBuilder, AgentHook, CompletionCallAction, CompletionCallEvent, + CompletionResponseEvent, HookContext, InvalidToolCallAction, MultiTurnStreamItem, + NoToolConfig, ObservationAction, OutputMode, RequestPatch, StreamingError, + ToolCall as ToolCallEvent, ToolCallAction, ToolResultAction, ToolResultEvent, run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}, }, completion::{ - AssistantContent, Completion, CompletionError, CompletionModel, Message, Prompt, - PromptError, ToolDefinition, + AssistantContent, CompletionError, CompletionModel, Message, Prompt, PromptError, + ToolDefinition, }, message::{ToolChoice, UserContent}, streaming::StreamingPrompt, @@ -1408,11 +1408,44 @@ where .tool(CountingSum(sum_calls.clone())) .tool_choice(ToolChoice::Required) .build(); - let response = agent - .completion(PROMPT, Vec::::new()) - .await? - .send() - .await?; + #[derive(Clone)] + struct CaptureTurn(Arc>>); + + impl AgentHook for CaptureTurn + where + M: CompletionModel, + { + async fn on_completion_response( + &self, + _ctx: &HookContext, + event: CompletionResponseEvent<'_, M>, + ) -> ObservationAction { + *lock_recover(&self.0) = Some(ModelTurn::new( + event.response.message_id.clone(), + event.response.choice.clone(), + event.response.usage, + BTreeSet::new(), + BTreeSet::new(), + )); + ObservationAction::stop("captured conformance model turn") + } + } + + let captured = Arc::new(Mutex::new(None)); + let stopped = agent + .runner(PROMPT) + .add_hook(CaptureTurn(captured.clone())) + .run() + .await; + if !matches!(stopped, Err(PromptError::PromptCancelled { .. })) { + return Err(ScenarioError::contract( + SCENARIO, + format!("capture hook did not stop after the model response: {stopped:?}"), + )); + } + let response = lock_recover(&captured).take().ok_or_else(|| { + ScenarioError::contract(SCENARIO, "capture hook observed no model response") + })?; let emitted = response .choice .iter() diff --git a/examples/README.md b/examples/README.md index e6cbb7de8..80a278b46 100644 --- a/examples/README.md +++ b/examples/README.md @@ -47,7 +47,7 @@ Most examples expect provider API keys in the environment (e.g. `OPENAI_API_KEY` | `gemini_nanobanana_image_generation` | See source. | | `gemini_stream_kill_token_count` | Live Gemini example: obtaining a token-count estimate when a streaming | | `gemini_video_understanding` | Demonstrates Gemini video understanding with provider-specific request parameters. | -| `manual_tool_calls` | Demonstrates manual tool-call handling with `Agent::completion()`. | +| `manual_tool_calls` | Demonstrates manual tool-call handling with an explicit raw `CompletionModel` request. | | `multi_agent` | See source. | | `multi_extract` | Demonstrates fan-out structured extraction with `try_parallel!`. | | `multi_turn_agent_extended` | See source. | diff --git a/examples/agent_run_stepping/src/main.rs b/examples/agent_run_stepping/src/main.rs index ccacdcd89..7777a2753 100644 --- a/examples/agent_run_stepping/src/main.rs +++ b/examples/agent_run_stepping/src/main.rs @@ -25,10 +25,10 @@ use rig::agent::{ AgentHook, HookContext, InvalidToolCallAction, ToolCall as ToolCallEvent, ToolCallAction, }; use rig::client::{CompletionClient, ProviderClient}; -use rig::completion::{Completion, CompletionModel}; +use rig::completion::CompletionModel; use rig::message::UserContent; use rig::providers::openai; -use rig::tool::Tool; +use rig::tool::{Tool, ToolSet}; use serde::Deserialize; use serde_json::json; @@ -91,11 +91,13 @@ impl AgentHook for ToolLoggerHook { #[tokio::main] async fn main() -> Result<()> { let openai = openai::Client::from_env()?; - let agent = openai - .agent(openai::GPT_4O) + let model = openai.completion_model(openai::GPT_4O); + let agent = rig::agent::AgentBuilder::new(model.clone()) .preamble("You are a calculator. Always use the provided tools to compute results.") .tool(Add) .build(); + let local_tools = ToolSet::builder().static_tool(Add).build(); + let tool_definitions = local_tools.get_tool_definitions(); let mut run = AgentRun::new("What is 2 + 5?").max_turns(2); @@ -107,17 +109,26 @@ async fn main() -> Result<()> { turn, } => { println!("→ model call #{turn}"); - let response = agent.completion(prompt, history).await?.send().await?; + // A hand-driven `AgentRun` is a sans-IO protocol primitive, not + // execution of the configured `Agent`. Its transport is an + // explicit raw model request and therefore has no agent hooks. + let response = model + .completion_request(prompt) + .messages(history) + .preamble( + "You are a calculator. Always use the provided tools to compute results." + .to_string(), + ) + .tools(tool_definitions.clone()) + .send() + .await?; // The tools advertised to the provider for this turn. With // static tools these are the agent's registered tools; agents // with dynamic (RAG) tools would resolve them per turn. - let tool_names: BTreeSet = agent - .tool_server_handle - .get_tool_defs(None) - .await? - .into_iter() - .map(|def| def.name) + let tool_names: BTreeSet = tool_definitions + .iter() + .map(|def| def.name.clone()) .collect(); let mut outcome = run.model_response(ModelTurn::new( @@ -157,10 +168,7 @@ async fn main() -> Result<()> { let args = call.tool_call.function.arguments.to_string(); println!("→ executing {name}({args})"); let mut context = rig::tool::ToolContext::new(); - let result = agent - .tool_server_handle - .execute(name, &args, &mut context) - .await; + let result = local_tools.execute(name, args, &mut context).await; results.push(UserContent::tool_result( call.tool_call.id.clone(), result.output().clone().into_content(), diff --git a/examples/agent_with_durable_approval/src/main.rs b/examples/agent_with_durable_approval/src/main.rs index 49ae5d9b0..44c0553d4 100644 --- a/examples/agent_with_durable_approval/src/main.rs +++ b/examples/agent_with_durable_approval/src/main.rs @@ -31,10 +31,10 @@ use anyhow::Result; use rig::agent::InvalidToolCallAction; use rig::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; use rig::client::{CompletionClient, ProviderClient}; -use rig::completion::Completion; +use rig::completion::CompletionModel; use rig::message::{ToolResultContent, UserContent}; use rig::providers::openai; -use rig::tool::Tool; +use rig::tool::{Tool, ToolSet}; use serde::Deserialize; use serde_json::json; use std::collections::BTreeSet; @@ -145,15 +145,17 @@ async fn ask(prompt: &str) -> Option { #[tokio::main] async fn main() -> Result<()> { - let agent = openai::Client::from_env()? - .agent(openai::GPT_4O) - .preamble( - "You are a banking assistant. Use the tools to carry out the user's request. \ - Call one tool at a time.", - ) - .tool(GetBalance) - .tool(TransferFunds) + // A serializable `AgentRun` is a sans-IO protocol primitive. This example + // intentionally supplies raw model transport and tool dispatch explicitly; + // configured `Agent` execution instead always goes through `AgentRunner`. + let model = openai::Client::from_env()?.completion_model(openai::GPT_4O); + let preamble = "You are a banking assistant. Use the tools to carry out the user's request. \ + Call one tool at a time."; + let tools = ToolSet::builder() + .static_tool(GetBalance) + .static_tool(TransferFunds) .build(); + let tool_definitions = tools.get_tool_definitions(); let prompt = "Check the balance of account A-1, then transfer $500 to account B-2."; println!("User: {prompt}"); @@ -172,13 +174,16 @@ async fn main() -> Result<()> { turn, } => { println!("\n→ model call #{turn}"); - let response = agent.completion(prompt, history).await?.send().await?; - let tool_names: BTreeSet = agent - .tool_server_handle - .get_tool_defs(None) - .await? - .into_iter() - .map(|def| def.name) + let response = model + .completion_request(prompt) + .messages(history) + .preamble(preamble.to_string()) + .tools(tool_definitions.clone()) + .send() + .await?; + let tool_names: BTreeSet = tool_definitions + .iter() + .map(|def| def.name.clone()) .collect(); let mut outcome = run.model_response(ModelTurn::new( response.message_id.clone(), @@ -226,9 +231,8 @@ async fn main() -> Result<()> { .as_deref() { Some("a") | Some("approve") => { - let execution = agent - .tool_server_handle - .execute(&name, &args, &mut rig::tool::ToolContext::new()) + let execution = tools + .execute(&name, args, &mut rig::tool::ToolContext::new()) .await; results.push(UserContent::tool_result( id, @@ -242,11 +246,10 @@ async fn main() -> Result<()> { .map(serde_json::from_str::) { Some(Ok(value)) => { - let execution = agent - .tool_server_handle + let execution = tools .execute( &name, - &value.to_string(), + value.to_string(), &mut rig::tool::ToolContext::new(), ) .await; diff --git a/examples/candle_local/src/main.rs b/examples/candle_local/src/main.rs index 0a407e7fa..4706f7e53 100644 --- a/examples/candle_local/src/main.rs +++ b/examples/candle_local/src/main.rs @@ -3,8 +3,8 @@ use std::io::Write; use anyhow::Context; use futures::StreamExt; use rig::candle::{CandleModel, ModelData}; -use rig::streaming::{StreamedAssistantContent, StreamingCompletion}; -use rig::{agent::AgentBuilder, message::Message}; +use rig::completion::CompletionModel; +use rig::streaming::StreamedAssistantContent; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -25,15 +25,11 @@ async fn main() -> anyhow::Result<()> { tokenizer: std::fs::read(model_dir.join("tokenizer.json"))?, weights: std::fs::read(model_dir.join("model.gguf"))?, })?; - let agent = AgentBuilder::new(model) - .preamble("You are a concise and helpful assistant.") + let mut response = model + .completion_request(prompt) + .preamble("You are a concise and helpful assistant.".to_string()) .temperature(0.0) .max_tokens(64) - .build(); - - let mut response = agent - .stream_completion(prompt, std::iter::empty::()) - .await? .stream() .await?; let mut final_response = None; diff --git a/examples/manual_tool_calls/src/main.rs b/examples/manual_tool_calls/src/main.rs index 03a4684f2..bebedb5be 100644 --- a/examples/manual_tool_calls/src/main.rs +++ b/examples/manual_tool_calls/src/main.rs @@ -1,4 +1,4 @@ -//! Demonstrates manual tool-call handling with `Agent::completion()`. +//! Demonstrates manual tool-call handling with a raw `CompletionModel` request. //! Requires `OPENAI_API_KEY`. //! //! Unlike `agent.prompt(...)`, this example never lets Rig execute tools automatically. @@ -12,7 +12,7 @@ use anyhow::{Result, bail}; use rig::OneOrMany; use rig::client::{CompletionClient, ProviderClient}; -use rig::completion::Completion; +use rig::completion::CompletionModel; use rig::message::{AssistantContent, Message, ToolCall, ToolChoice, UserContent}; use rig::providers::openai; use rig::tool::{Tool, ToolOutput, ToolSet}; @@ -133,17 +133,11 @@ fn tool_result_message(tool_call: &ToolCall, output: ToolOutput) -> Message { async fn main() -> Result<()> { const MAX_ROUNDS: usize = 8; - let agent = openai::Client::from_env()? - .agent(openai::GPT_4O_MINI) - .preamble( - "You are a calculator. Never do arithmetic from memory. \ - Use the provided tools for every intermediate step. \ - You may emit one or multiple tool calls in a single turn. \ - Once all tool results are available, give a short final answer.", - ) - .tool(Add) - .tool(Subtract) - .build(); + let model = openai::Client::from_env()?.completion_model(openai::GPT_4O_MINI); + let preamble = "You are a calculator. Never do arithmetic from memory. \ + Use the provided tools for every intermediate step. \ + You may emit one or multiple tool calls in a single turn. \ + Once all tool results are available, give a short final answer."; let local_tools = ToolSet::builder() .static_tool(Add) @@ -156,9 +150,13 @@ async fn main() -> Result<()> { ); for round in 1..=MAX_ROUNDS { - let mut request = agent - .completion(current_prompt.clone(), history.clone()) - .await?; + // This example intentionally operates below the Agent abstraction. Raw + // model requests have no agent lifecycle or hooks. + let mut request = model + .completion_request(current_prompt.clone()) + .preamble(preamble.to_string()) + .messages(history.clone()) + .tools(local_tools.get_tool_definitions()); if round == 1 { // Force the first turn through the tool path so the example always demonstrates it. request = request.tool_choice(ToolChoice::Required); diff --git a/tests/providers/gemini/agent_run_support.rs b/tests/providers/gemini/agent_run_support.rs index 6117a883e..f0225162f 100644 --- a/tests/providers/gemini/agent_run_support.rs +++ b/tests/providers/gemini/agent_run_support.rs @@ -7,14 +7,62 @@ use std::collections::BTreeSet; use rig::agent::CompletionCall; use rig::agent::run::{ModelTurn, PendingToolCall}; -use rig::completion::{Completion, ToolDefinition, Usage}; -use rig::message::{AssistantContent, Message, ToolResultContent, UserContent}; +use rig::completion::{CompletionModel, CompletionRequestBuilder, ToolDefinition, Usage}; +use rig::message::{AssistantContent, Message, ToolChoice, ToolResultContent, UserContent}; use rig::providers::gemini; use rig::tool::Tool; use serde::Deserialize; use serde_json::json; -pub(crate) type GeminiAgent = rig::agent::Agent; +pub(crate) struct GeminiAgent { + model: gemini::completion::CompletionModel, + preamble: String, + tools: Vec, + tool_choice: Option, +} + +impl GeminiAgent { + pub(crate) fn new( + model: gemini::completion::CompletionModel, + preamble: impl Into, + tool_names: &[&str], + tool_choice: Option, + ) -> Self { + Self { + model, + preamble: preamble.into(), + tools: tool_names + .iter() + .map(|name| match *name { + "add" => operation_definition("add", "Add x and y together"), + "sum" => operation_definition("sum", "Add x and y together (alias of add)"), + "subtract" => { + operation_definition("subtract", "Subtract y from x (i.e. x - y)") + } + other => panic!("unsupported raw harness tool `{other}`"), + }) + .collect(), + tool_choice, + } + } + + pub(crate) fn request( + &self, + prompt: Message, + history: Vec, + ) -> CompletionRequestBuilder { + let mut request = self + .model + .completion_request(prompt) + .messages(history) + .preamble(self.preamble.clone()) + .tools(self.tools.clone()); + if let Some(tool_choice) = &self.tool_choice { + request = request.tool_choice(tool_choice.clone()); + } + request + } +} pub(crate) const FORCE_TOOLS_PREAMBLE: &str = "You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing results yourself. Once you have all the tool results you need, reply with the final numeric answer in plain text."; @@ -166,9 +214,9 @@ pub(crate) fn execute_pending_calls(calls: &[PendingToolCall]) -> Vec, ) -> ModelTurn { let response = agent - .completion(prompt, history) - .await - .expect("completion request should build") + .request(prompt, history) .send() .await .expect("gemini completion should succeed"); diff --git a/tests/providers/gemini/cassette/agent_run_recovery.rs b/tests/providers/gemini/cassette/agent_run_recovery.rs index 563a0820c..547c9e4bf 100644 --- a/tests/providers/gemini/cassette/agent_run_recovery.rs +++ b/tests/providers/gemini/cassette/agent_run_recovery.rs @@ -12,7 +12,7 @@ use rig::providers::gemini; use rig::test_utils::validate_unknown_tool_failure; use super::super::agent_run_support::{ - Add, FORCE_TOOLS_PREAMBLE, Subtract, Sum, assistant_tool_call_names, call_model, + FORCE_TOOLS_PREAMBLE, GeminiAgent, assistant_tool_call_names, call_model, execute_pending_calls, history_has_assistant_tool_call, tool_names, user_content_tool_result_texts, }; @@ -53,12 +53,12 @@ async fn fail_resolution_returns_unknown_tool_call() { with_gemini_cassette( "agent_run_recovery/fail_resolution_returns_unknown_tool_call", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool_choice(ToolChoice::Required) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + Some(ToolChoice::Required), + ); let mut run = run_until_invalid_add_call(&agent, &tool_names(&[]), 0).await; let error = run @@ -96,12 +96,12 @@ async fn repair_renames_tool_call_and_executes_it() { |client| async move { // `sum` is registered alongside `add` so the post-repair wire // history references a tool Gemini saw advertised. - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Sum) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "sum"], + None, + ); let machine_names = tool_names(&["sum"]); let mut run = @@ -181,12 +181,12 @@ async fn skip_suppresses_every_call_in_the_turn() { with_gemini_cassette( "agent_run_recovery/skip_suppresses_every_call_in_the_turn", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Subtract) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "subtract"], + None, + ); let executable = tool_names(&["add", "subtract"]); // `add` is disallowed for the first turn. let restricted = tool_names(&["subtract"]); @@ -273,12 +273,12 @@ async fn retry_with_exhausted_budget_fails_with_unknown_tool_call() { with_gemini_cassette( "agent_run_recovery/retry_with_exhausted_budget_fails_with_unknown_tool_call", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool_choice(ToolChoice::Required) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + Some(ToolChoice::Required), + ); // Zero retry budget: the first retry resolution must fail. let mut run = run_until_invalid_add_call(&agent, &tool_names(&[]), 0).await; @@ -303,12 +303,12 @@ async fn repair_to_disallowed_name_fails_with_unknown_tool_call() { with_gemini_cassette( "agent_run_recovery/repair_to_disallowed_name_fails_with_unknown_tool_call", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool_choice(ToolChoice::Required) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + Some(ToolChoice::Required), + ); let mut run = run_until_invalid_add_call(&agent, &tool_names(&["subtract"]), 0).await; let error = run diff --git a/tests/providers/gemini/cassette/agent_run_resume.rs b/tests/providers/gemini/cassette/agent_run_resume.rs index 831fe7bb7..f98ab380d 100644 --- a/tests/providers/gemini/cassette/agent_run_resume.rs +++ b/tests/providers/gemini/cassette/agent_run_resume.rs @@ -8,8 +8,8 @@ use rig::client::CompletionClient; use rig::providers::gemini; use super::super::agent_run_support::{ - Add, FORCE_TOOLS_PREAMBLE, call_model, execute_pending_calls, history_has_assistant_tool_call, - is_tool_result_user_message, tool_names, tool_result_texts, + FORCE_TOOLS_PREAMBLE, GeminiAgent, call_model, execute_pending_calls, + history_has_assistant_tool_call, is_tool_result_user_message, tool_names, tool_result_texts, }; use super::super::support::with_gemini_cassette; use crate::support::{assert_mentions_expected_number, assert_nonempty_response}; @@ -26,11 +26,12 @@ async fn resume_from_serialized_state_mid_tool_execution() { with_gemini_cassette( "agent_run_resume/resume_from_serialized_state_mid_tool_execution", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + None, + ); let names = tool_names(&["add"]); let mut run = AgentRun::new("What is 21 + 21? Use the add tool.").max_turns(2); @@ -131,11 +132,12 @@ async fn resume_while_invalid_tool_call_awaits_resolution() { with_gemini_cassette( "agent_run_resume/resume_while_invalid_tool_call_awaits_resolution", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + None, + ); let executable = tool_names(&["add"]); // Machine-side restriction: the model's `add` call is disallowed // for this turn even though it was advertised on the wire. @@ -229,11 +231,12 @@ async fn resume_after_invalid_tool_call_retry_rollback() { with_gemini_cassette( "agent_run_resume/resume_after_invalid_tool_call_retry_rollback", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + None, + ); let executable = tool_names(&["add"]); let nothing_allowed = tool_names(&[]); const FEEDBACK: &str = "Tools are temporarily unavailable. Answer the question directly in plain text without calling any tools."; diff --git a/tests/providers/gemini/cassette/agent_run_stepping.rs b/tests/providers/gemini/cassette/agent_run_stepping.rs index 1d2f240c1..6214ecc43 100644 --- a/tests/providers/gemini/cassette/agent_run_stepping.rs +++ b/tests/providers/gemini/cassette/agent_run_stepping.rs @@ -9,7 +9,7 @@ use rig::message::{Message, ToolChoice, UserContent}; use rig::providers::gemini; use super::super::agent_run_support::{ - Add, FORCE_TOOLS_PREAMBLE, Subtract, call_model, execute_pending_calls, + FORCE_TOOLS_PREAMBLE, GeminiAgent, call_model, execute_pending_calls, history_has_assistant_tool_call, is_tool_result_user_message, sum_completion_call_usage, tool_names, }; @@ -23,10 +23,12 @@ async fn hand_driven_single_turn_completes() { with_gemini_cassette( "agent_run_stepping/hand_driven_single_turn_completes", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(BASIC_PREAMBLE) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + BASIC_PREAMBLE, + &[], + None, + ); let names = tool_names(&[]); let mut run = AgentRun::new(BASIC_PROMPT); @@ -104,12 +106,12 @@ async fn hand_driven_multi_turn_tool_run_completes() { with_gemini_cassette( "agent_run_stepping/hand_driven_multi_turn_tool_run_completes", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Subtract) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "subtract"], + None, + ); let names = tool_names(&["add", "subtract"]); let mut run = AgentRun::new( @@ -196,12 +198,12 @@ async fn hand_driven_parallel_tool_calls_arrive_in_one_step() { with_gemini_cassette( "agent_run_stepping/hand_driven_parallel_tool_calls_arrive_in_one_step", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Subtract) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "subtract"], + None, + ); let names = tool_names(&["add", "subtract"]); let mut run = AgentRun::new( @@ -260,12 +262,12 @@ async fn max_turns_error_carries_pending_tool_results_message() { with_gemini_cassette( "agent_run_stepping/max_turns_error_carries_pending_tool_results_message", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool_choice(ToolChoice::Required) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + Some(ToolChoice::Required), + ); let names = tool_names(&["add"]); // max_turns 2: the run errors once tool results demand a third diff --git a/tests/providers/gemini/cassette/agent_run_streamed.rs b/tests/providers/gemini/cassette/agent_run_streamed.rs index a5eb1d191..ed3e98200 100644 --- a/tests/providers/gemini/cassette/agent_run_streamed.rs +++ b/tests/providers/gemini/cassette/agent_run_streamed.rs @@ -18,11 +18,11 @@ use rig::client::CompletionClient; use rig::completion::{GetTokenUsage, PromptError, Usage}; use rig::message::{Message, ToolChoice, ToolResult}; use rig::providers::gemini; -use rig::streaming::{StreamedAssistantContent, StreamingCompletion, StreamingPrompt}; +use rig::streaming::{StreamedAssistantContent, StreamingPrompt}; use rig::test_utils::{validate_cancelled_failure, validate_max_turns_failure}; use super::super::agent_run_support::{ - Add, FORCE_TOOLS_PREAMBLE, GeminiAgent, Subtract, Sum, assert_canonical_assistant_order, + Add, FORCE_TOOLS_PREAMBLE, GeminiAgent, assert_canonical_assistant_order, assistant_tool_call_names, execute_pending_calls, history_has_assistant_tool_call, is_tool_result_user_message, sum_completion_call_usage, tool_names, }; @@ -56,9 +56,7 @@ async fn run_streamed_turn( collected_text: &mut String, ) -> Result { let mut stream = agent - .stream_completion(prompt, &history) - .await - .expect("stream request should build") + .request(prompt, history) .stream() .await .expect("gemini stream should open"); @@ -163,12 +161,12 @@ async fn streamed_hand_driven_multi_turn_run_completes() { "a phantom completion call must be rejected on a fresh run" ); - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Subtract) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "subtract"], + None, + ); let names = tool_names(&["add", "subtract"]); let mut run = AgentRun::new( @@ -244,12 +242,12 @@ async fn streamed_invalid_tool_call_fails_fast_mid_stream() { with_gemini_cassette( "agent_run_streamed/streamed_invalid_tool_call_fails_fast_mid_stream", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool_choice(ToolChoice::Required) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + Some(ToolChoice::Required), + ); let executable = tool_names(&["add"]); let nothing_allowed = tool_names(&[]); @@ -303,12 +301,12 @@ async fn streamed_repair_continues_the_same_stream() { with_gemini_cassette( "agent_run_streamed/streamed_repair_continues_the_same_stream", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .tool(Sum) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add", "sum"], + None, + ); let machine_names = tool_names(&["sum"]); let mut run = @@ -375,11 +373,12 @@ async fn streamed_skip_abandons_the_turn_and_recovers() { with_gemini_cassette( "agent_run_streamed/streamed_skip_abandons_the_turn_and_recovers", |client| async move { - let agent = client - .agent(gemini::completion::GEMINI_2_5_FLASH) - .preamble(FORCE_TOOLS_PREAMBLE) - .tool(Add) - .build(); + let agent = GeminiAgent::new( + client.completion_model(gemini::completion::GEMINI_2_5_FLASH), + FORCE_TOOLS_PREAMBLE, + &["add"], + None, + ); let executable = tool_names(&["add"]); let nothing_allowed = tool_names(&[]); const SKIP_REASON: &str = "The add tool is disabled for this request."; diff --git a/tests/providers/gemini/cassette/dynamic_tools.rs b/tests/providers/gemini/cassette/dynamic_tools.rs index ee96883a3..7725490c4 100644 --- a/tests/providers/gemini/cassette/dynamic_tools.rs +++ b/tests/providers/gemini/cassette/dynamic_tools.rs @@ -150,8 +150,7 @@ async fn sample_caps_retrieved_definitions() { .build(); let defs = agent - .tool_server_handle - .get_tool_defs(Some( + .tool_definitions(Some( "Multiply two numbers together to get their product.".to_string(), )) .await diff --git a/tests/providers/gemini/cassette/multi_turn_streaming.rs b/tests/providers/gemini/cassette/multi_turn_streaming.rs index 0a128fd4c..27ffc9b1a 100644 --- a/tests/providers/gemini/cassette/multi_turn_streaming.rs +++ b/tests/providers/gemini/cassette/multi_turn_streaming.rs @@ -1,41 +1,24 @@ //! Migrated from `examples/multi_turn_streaming_gemini.rs`. -use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use futures::{Stream, StreamExt}; -use rig::OneOrMany; -use rig::agent::Agent; +use futures::StreamExt; +use rig::agent::MultiTurnStreamItem; use rig::client::CompletionClient; -use rig::completion::{self, CompletionError, CompletionModel, PromptError}; -use rig::message::{AssistantContent, Message, Text, UserContent}; use rig::providers::gemini; -use rig::streaming::{StreamedAssistantContent, StreamingCompletion}; -use rig::tool::{Tool, ToolContext, ToolExecutionError}; +use rig::streaming::StreamingPrompt; +use rig::tool::Tool; use schemars::{JsonSchema, schema_for}; use serde::Deserialize; -use thiserror::Error; use crate::support::{ MULTI_TURN_STREAMING_EXPECTED_RESULT, MULTI_TURN_STREAMING_PROMPT, assert_mentions_expected_number, assert_nonempty_response, }; -#[derive(Debug, Error)] -enum StreamingError { - #[error("CompletionError: {0}")] - Completion(#[from] CompletionError), - #[error("PromptError: {0}")] - Prompt(#[from] Box), - #[error("ToolExecutionError: {0}")] - Tool(#[from] ToolExecutionError), -} - -type StreamingResult = Pin> + Send>>; - #[tokio::test] -async fn manual_multi_turn_streaming_loop() { +async fn runner_driven_multi_turn_streaming_loop() { let add_calls = Arc::new(AtomicUsize::new(0)); let subtract_calls = Arc::new(AtomicUsize::new(0)); let multiply_calls = Arc::new(AtomicUsize::new(0)); @@ -53,11 +36,19 @@ async fn manual_multi_turn_streaming_loop() { .tool(Divide::new(divide_calls.clone())) .build(); - let mut stream = - multi_turn_prompt(agent, MULTI_TURN_STREAMING_PROMPT, Vec::new()).await; - let response = collect_text(&mut stream) - .await - .expect("manual multi-turn streaming should succeed"); + let mut stream = agent + .stream_prompt(MULTI_TURN_STREAMING_PROMPT) + .max_turns(10) + .await; + let mut response = None; + while let Some(item) = stream.next().await { + if let MultiTurnStreamItem::FinalResponse(final_response) = + item.expect("runner-driven multi-turn streaming should succeed") + { + response = Some(final_response.output); + } + } + let response = response.expect("stream should emit a final response"); assert_nonempty_response(&response); assert!( @@ -87,110 +78,6 @@ async fn manual_multi_turn_streaming_loop() { .await; } -async fn multi_turn_prompt( - agent: Agent, - prompt: impl Into + Send, - mut chat_history: Vec, -) -> StreamingResult -where - M: CompletionModel + 'static, - M::StreamingResponse: Send, -{ - let prompt: Message = prompt.into(); - - Box::pin(async_stream::stream! { - let mut current_prompt = prompt; - let mut did_call_tool = false; - - 'outer: loop { - let mut stream = agent - .stream_completion(current_prompt.clone(), &chat_history) - .await? - .stream() - .await?; - - chat_history.push(current_prompt.clone()); - let mut tool_calls = vec![]; - let mut tool_results = vec![]; - - while let Some(content) = stream.next().await { - match content { - Ok(StreamedAssistantContent::Text(text)) => { - yield Ok(Text::new(text.text)); - did_call_tool = false; - } - Ok(StreamedAssistantContent::ToolCall { tool_call, .. }) => { - let execution = agent - .tool_server_handle - .execute( - &tool_call.function.name, - &tool_call.function.arguments.to_string(), - &mut ToolContext::new(), - ) - .await; - if let Some(error) = execution.error() { - Err(StreamingError::Tool(error.clone()))?; - } - let tool_result = execution.output().clone(); - - tool_calls.push(AssistantContent::ToolCall(tool_call.clone())); - tool_results.push((tool_call.id, tool_call.call_id, tool_result)); - did_call_tool = true; - } - Ok(StreamedAssistantContent::Reasoning(reasoning)) => { - let rendered = reasoning.display_text(); - if !rendered.is_empty() { - yield Ok(Text::new(rendered)); - } - did_call_tool = false; - } - Ok(_) => {} - Err(error) => { - yield Err(error.into()); - break 'outer; - } - } - } - - if !tool_calls.is_empty() { - chat_history.push(Message::Assistant { - id: None, - content: OneOrMany::many(tool_calls).expect("tool calls should be non-empty"), - }); - } - - for (id, call_id, tool_result) in tool_results { - let tool_content = tool_result.into_content(); - let user_content = if let Some(call_id) = call_id { - UserContent::tool_result_with_call_id(id, call_id, tool_content) - } else { - UserContent::tool_result(id, tool_content) - }; - chat_history.push(Message::User { - content: OneOrMany::one(user_content), - }); - } - - current_prompt = match chat_history.pop() { - Some(prompt) => prompt, - None => unreachable!("chat history should not be empty"), - }; - - if !did_call_tool { - break; - } - } - }) -} - -async fn collect_text(stream: &mut StreamingResult) -> Result { - let mut text = String::new(); - while let Some(content) = stream.next().await { - text.push_str(&content?.text); - } - Ok(text) -} - #[derive(Deserialize, JsonSchema)] struct OperationArgs { x: i32,