rig_ecs/policy/
mod.rs

1//! Pure request assembly, output selection, and tool-result shaping over plain data.
2//!
3//! ```
4//! use rig_ecs::{agent::OutputKind, policy::resolve_output};
5//! let mode = resolve_output(OutputKind::Auto, true, 0, true, false);
6//! assert_eq!(mode, OutputKind::Native);
7//! ```
8
9use rig_core::{
10    completion::{
11        CompletionRequest, Document, ToolDefinition,
12        message::{CallId, Message, ToolChoice, ToolName, UserContent},
13    },
14    effect::{HandlerDescriptor, Outcome},
15    error::{ErrorKind, ErrorReport},
16    json_utils::to_canonical_string,
17    structured_output::{
18        AUGMENTATION_SEPARATOR, OUTPUT_TOOL_DESCRIPTION, output_tool_augmentation,
19        prompted_augmentation,
20    },
21    tool::{ToolExecutionError, ToolResult},
22    transcript::tool_result_output,
23};
24
25use crate::agent::{
26    Failure, MessageParts, OutputKind, OutputToolConfig,
27    content::parts::{ToolResultLimit, ToolResultStatus},
28};
29
30/// The output mode a turn runs under once `Auto` is resolved, never `Auto`:
31/// no schema is `Native`; an explicit `Tool` the choice forbids degrades to
32/// `Native` (the constraint is still enforced, natively); `Auto` is `Tool`
33/// only with a real tool of the program's own, a permitting choice, and a
34/// provider that does not compose native output with tools; otherwise `Native`.
35pub fn resolve_output(
36    mode: OutputKind,
37    has_schema: bool,
38    granted_tools: usize,
39    callable: bool,
40    provider_composes_native: bool,
41) -> OutputKind {
42    if !has_schema {
43        return OutputKind::Native;
44    }
45    match mode {
46        OutputKind::Native => OutputKind::Native,
47        OutputKind::Prompted => OutputKind::Prompted,
48        OutputKind::Tool => {
49            if callable {
50                OutputKind::Tool
51            } else {
52                OutputKind::Native
53            }
54        }
55        OutputKind::Auto => {
56            if granted_tools > 0 && callable && !provider_composes_native {
57                OutputKind::Tool
58            } else {
59                OutputKind::Native
60            }
61        }
62    }
63}
64
65/// Ordered request inputs gathered by assembly, borrowing message parts and
66/// handler descriptors while owning the attached document list.
67pub struct RequestGraph<'a> {
68    /// The preamble, if the program has one.
69    pub preamble: Option<&'a str>,
70    /// The utterances in order.
71    pub utterances: Vec<&'a MessageParts>,
72    /// The documents attached to the turn, in order.
73    pub documents: Vec<Document>,
74    /// The tools granted, in advertisement order, by their descriptors.
75    pub tools: Vec<&'a HandlerDescriptor>,
76    /// Sampling.
77    pub temperature: Option<f64>,
78    /// The token budget.
79    pub max_tokens: Option<u64>,
80    /// Whether unknown finish reasons stop normally.
81    pub accept_unknown_finish_reasons: bool,
82    /// Provider parameters.
83    pub additional_params: Option<&'a serde_json::Value>,
84    /// Portable generation options, the run's overlaid on the agent's.
85    pub options: &'a rig_core::completion::GenerationOptions,
86    /// Typed per-provider options, the run's entries over the agent's.
87    pub provider_options: &'a rig_core::completion::ProviderOptions,
88    /// The program's tool choice.
89    pub tool_choice: Option<&'a ToolChoice>,
90    /// The output mode, resolved, and its schema.
91    pub output: OutputKind,
92    /// The schema, if any.
93    pub schema: Option<&'a serde_json::Value>,
94    /// The output tool's name, when the mode is `Tool`.
95    pub output_tool: Option<&'a str>,
96    /// Optional custom description and preamble behavior of the output tool.
97    pub output_tool_config: Option<&'a OutputToolConfig>,
98}
99
100/// Construct a completion request from ordered graph inputs and resolved output
101/// policy. Non-tool descriptors are omitted; an invalid native schema is omitted.
102/// A graph with neither a preamble nor an utterance holds no conversation and
103/// is [`ContentError::Missing`](crate::agent::content::parts::ContentError::Missing),
104/// as is an empty output tool name.
105pub fn fold_request(
106    graph: &RequestGraph<'_>,
107) -> Result<CompletionRequest, crate::agent::content::parts::ContentError> {
108    let mut chat_history: Vec<Message> = Vec::with_capacity(graph.utterances.len() + 2);
109    let system = system_message(graph);
110    if let Some(content) = system {
111        chat_history.push(Message::System { content });
112    }
113    chat_history.extend(graph.utterances.iter().map(|parts| parts.to_message()));
114
115    let mut tools: Vec<ToolDefinition> = graph
116        .tools
117        .iter()
118        .filter_map(|descriptor| tool_definition(descriptor))
119        .collect();
120    if graph.output == OutputKind::Tool
121        && let (Some(name), Some(schema)) = (graph.output_tool, graph.schema)
122    {
123        tools.push(ToolDefinition {
124            name: ToolName::new(name)
125                .map_err(|_| crate::agent::content::parts::ContentError::Missing)?,
126            description: graph
127                .output_tool_config
128                .and_then(|config| config.description.as_deref())
129                .unwrap_or(OUTPUT_TOOL_DESCRIPTION)
130                .to_owned(),
131            parameters: schema.clone(),
132        });
133    }
134
135    let output_schema = match (graph.output, graph.schema) {
136        (OutputKind::Native, Some(schema)) => {
137            rig_core::schemars::Schema::try_from(schema.clone()).ok()
138        }
139        (OutputKind::Native, None)
140        | (OutputKind::Auto | OutputKind::Tool | OutputKind::Prompted, _) => None,
141    };
142
143    if chat_history.is_empty() {
144        return Err(crate::agent::content::parts::ContentError::Missing);
145    }
146    // Fields are assigned, not set through the builder, whose setters log
147    // shadowed `additional_params` keys the program already declared.
148    let mut request = CompletionRequest::from(chat_history);
149    request.documents = graph.documents.clone();
150    request.tools = tools;
151    request.temperature = graph.temperature;
152    request.max_tokens = graph.max_tokens;
153    request.tool_choice = graph.tool_choice.cloned();
154    request.additional_params = graph.additional_params.cloned();
155    request.options = graph.options.clone();
156    request.provider_options = graph.provider_options.clone();
157    request.output_schema = output_schema;
158    request.accept_unknown_finish_reasons = graph.accept_unknown_finish_reasons;
159    Ok(request)
160}
161
162/// The system message: the preamble with the output mode's augmentation,
163/// or none when the program has no preamble and nothing to add.
164fn system_message(graph: &RequestGraph<'_>) -> Option<String> {
165    let augmentation = match (graph.output, graph.output_tool, graph.schema) {
166        (OutputKind::Tool, Some(name), _) => graph
167            .output_tool_config
168            .is_none_or(|config| config.augment_preamble)
169            .then(|| output_tool_augmentation(name)),
170        (OutputKind::Prompted, _, Some(schema)) => {
171            Some(prompted_augmentation(&to_canonical_string(schema)))
172        }
173        (OutputKind::Tool, None, _)
174        | (OutputKind::Prompted, _, None)
175        | (OutputKind::Auto | OutputKind::Native, _, _) => None,
176    };
177    match (graph.preamble, augmentation) {
178        (None, None) => None,
179        (Some(preamble), None) => Some(preamble.to_owned()),
180        (None, Some(augmentation)) => Some(augmentation),
181        (Some(preamble), Some(augmentation)) => Some(format!(
182            "{preamble}{}{augmentation}",
183            AUGMENTATION_SEPARATOR
184        )),
185    }
186}
187
188/// The name a tool descriptor is called by; `None` for other effect families.
189pub(crate) fn tool_name(descriptor: &HandlerDescriptor) -> Option<&str> {
190    match &descriptor.family {
191        rig_core::effect::FamilyDescriptor::Tool { name, .. } => Some(name),
192        rig_core::effect::FamilyDescriptor::Completion { .. }
193        | rig_core::effect::FamilyDescriptor::Embed { .. }
194        | rig_core::effect::FamilyDescriptor::Rerank { .. }
195        | rig_core::effect::FamilyDescriptor::Memory { .. }
196        | rig_core::effect::FamilyDescriptor::Retrieve { .. }
197        | rig_core::effect::FamilyDescriptor::Custom { .. } => None,
198    }
199}
200
201/// Convert a tool descriptor to its model-facing definition; return `None` for
202/// other effect families and for a tool with an empty name.
203pub fn tool_definition(descriptor: &HandlerDescriptor) -> Option<ToolDefinition> {
204    match &descriptor.family {
205        rig_core::effect::FamilyDescriptor::Tool {
206            name,
207            description,
208            parameters,
209            ..
210        } => Some(ToolDefinition {
211            name: ToolName::new(name.as_str()).ok()?,
212            description: description.clone(),
213            parameters: parameters.clone(),
214        }),
215        rig_core::effect::FamilyDescriptor::Completion { .. }
216        | rig_core::effect::FamilyDescriptor::Embed { .. }
217        | rig_core::effect::FamilyDescriptor::Rerank { .. }
218        | rig_core::effect::FamilyDescriptor::Memory { .. }
219        | rig_core::effect::FamilyDescriptor::Retrieve { .. }
220        | rig_core::effect::FamilyDescriptor::Custom { .. } => None,
221    }
222}
223
224/// Convert a tool outcome to model-visible content and a separate graph status.
225/// Denials become skipped results; other nonterminal reports become error results.
226/// Returns a run failure for cancellation, unavailable handlers, a closed bus,
227/// or replay divergence.
228pub fn tool_result_part(
229    id: CallId,
230    name: ToolName,
231    outcome: &Result<Outcome, ErrorReport>,
232) -> Result<(UserContent, ToolResultStatus), Failure> {
233    if let Some(failure) = tool_failure(outcome) {
234        return Err(failure);
235    }
236    let (result, status) = match outcome {
237        Ok(Outcome::ToolResult { result }) => (
238            result.clone(),
239            if result.is_success() {
240                ToolResultStatus::Ok
241            } else if result.is_refused() {
242                ToolResultStatus::Refused
243            } else if result.is_skipped() {
244                ToolResultStatus::Skipped
245            } else {
246                ToolResultStatus::Error
247            },
248        ),
249        Ok(other) => (
250            ToolResult::failed(ToolExecutionError::other(format!(
251                "the tool handler answered with a {} outcome",
252                other.family()
253            ))),
254            ToolResultStatus::WrongFamily,
255        ),
256        Err(report) if report.kind == ErrorKind::Denied => (
257            ToolResult::skipped(report.message.clone()),
258            ToolResultStatus::Denied,
259        ),
260        Err(report) => {
261            let result = ToolResult::failed(report.clone().into());
262            let status = if result.is_refused() {
263                ToolResultStatus::Refused
264            } else {
265                ToolResultStatus::Error
266            };
267            (result, status)
268        }
269    };
270    Ok((tool_result_output(id, name, &result), status))
271}
272
273/// Return `None` if text fits the limit; otherwise retain head and tail totaling
274/// at most `limit.max_bytes`, cut at UTF-8 boundaries around the omission marker.
275pub fn limit_tool_result_text(text: &str, limit: &ToolResultLimit) -> Option<String> {
276    if text.len() <= limit.max_bytes {
277        return None;
278    }
279    let mut head = limit.max_bytes / 2;
280    while !text.is_char_boundary(head) {
281        head -= 1;
282    }
283    let mut tail = text.len() - (limit.max_bytes - limit.max_bytes / 2);
284    while !text.is_char_boundary(tail) {
285        tail += 1;
286    }
287    let omitted = tail - head;
288    let marker = limit.marker.replace("{omitted}", &omitted.to_string());
289    let head = text.get(..head).unwrap_or_default();
290    let tail = text.get(tail..).unwrap_or_default();
291    let mut cut = String::with_capacity(head.len() + marker.len() + tail.len());
292    cut.push_str(head);
293    cut.push_str(&marker);
294    cut.push_str(tail);
295    Some(cut)
296}
297
298/// Whether `limit_tool_results` would cut anything of `parts`: a text item
299/// of a tool-result part longer than the limit.
300#[must_use]
301pub fn tool_results_exceed(parts: &MessageParts, limit: &ToolResultLimit) -> bool {
302    let MessageParts::User { content } = parts else {
303        return false;
304    };
305    content.iter().any(|part| match part {
306        UserContent::ToolResult(result) => result.content.iter().any(|item| {
307            matches!(item, rig_core::message::ToolResultContent::Text(text) if text.text.len() > limit.max_bytes)
308        }),
309        _ => false,
310    })
311}
312
313/// `limit_tool_result_text` over every text item of every tool-result part
314/// of a message; JSON and image items and every other part are untouched.
315pub fn limit_tool_results(parts: &mut MessageParts, limit: &ToolResultLimit) {
316    let MessageParts::User { content } = parts else {
317        return;
318    };
319    for part in content {
320        let UserContent::ToolResult(result) = part else {
321            continue;
322        };
323        for item in &mut result.content {
324            if let rig_core::message::ToolResultContent::Text(text) = item
325                && let Some(cut) = limit_tool_result_text(&text.text, limit)
326            {
327                text.text = cut;
328            }
329        }
330    }
331}
332
333/// The failure a tool call's outcome ends the run in, if any: a cancel,
334/// or a report the bus could not serve the call with (closed, no handler,
335/// a replay divergence). Every other outcome is a result the model sees.
336pub fn tool_failure(outcome: &Result<Outcome, ErrorReport>) -> Option<Failure> {
337    match outcome {
338        Ok(_) => None,
339        Err(report) if report.kind == ErrorKind::Cancelled => {
340            Some(Failure::Cancelled(report.clone()))
341        }
342        Err(report)
343            if matches!(
344                report.kind,
345                ErrorKind::BusClosed | ErrorKind::HandlerUnavailable | ErrorKind::Divergence
346            ) =>
347        {
348            Some(Failure::Tool(report.clone()))
349        }
350        Err(_) => None,
351    }
352}
353
354/// Return retrieval text from the last utterance with usable text, or an empty
355/// string when none exists. Tool-result-only utterances do not replace the query.
356pub fn retrieval_query(utterances: &[MessageParts]) -> String {
357    utterances
358        .iter()
359        .rev()
360        .find_map(|parts| parts.to_message().rag_text())
361        .unwrap_or_default()
362}
363
364#[cfg(test)]
365mod tests;
rig_ecs/policy/mod.rs

rig_ecs/policy/
mod.rs
