rig_agent/
completion.rs

1//! High-level prompting traits and runtime errors for the classic agent runtime.
2
3use serde::de::DeserializeOwned;
4use thiserror::Error;
5
6use rig_core::{
7    memory::MemoryError,
8    wasm_compat::{WasmCompatSend, WasmCompatSync},
9};
10
11pub use rig_core::completion::*;
12
13/// Errors from classic agent prompting.
14#[derive(Debug, Error)]
15#[non_exhaustive]
16pub enum PromptError {
17    /// A provider completion failed.
18    #[error("CompletionError: {0}")]
19    CompletionError(#[from] CompletionError),
20
21    /// Conversation memory failed to load or persist history.
22    #[error("MemoryError: {0}")]
23    MemoryError(#[from] MemoryError),
24
25    /// The run exhausted its total model-call budget.
26    #[error("MaxTurnsError: reached max turns limit: {max_turns}")]
27    MaxTurnsError {
28        /// Configured total model-call budget.
29        max_turns: usize,
30        /// Canonical history available when the budget was exhausted.
31        chat_history: Box<Vec<Message>>,
32        /// Prompt for the call that could not be dispatched.
33        prompt: Box<Message>,
34    },
35
36    /// A prompting loop was cancelled.
37    #[error("PromptCancelled: {reason}")]
38    PromptCancelled {
39        /// Canonical history available at cancellation.
40        chat_history: Vec<Message>,
41        /// Human-readable cancellation reason.
42        reason: String,
43    },
44
45    /// The model attempted to call a tool unavailable for the current turn.
46    #[error(
47        "UnknownToolCall: model attempted to call unknown or disallowed tool `{tool_name}`. Available tools: {available_tools:?}. Allowed tools for this turn: {allowed_tools:?}"
48    )]
49    UnknownToolCall {
50        /// Tool name emitted by the model.
51        tool_name: String,
52        /// Tools registered on the runtime.
53        available_tools: Vec<String>,
54        /// Exact immutable set allowed for this turn.
55        allowed_tools: Vec<String>,
56        /// Canonical history available at failure.
57        chat_history: Box<Vec<Message>>,
58    },
59}
60
61impl PromptError {
62    /// Returns the provider response body exposed by a wrapped completion error.
63    pub fn provider_response_body(&self) -> Option<&str> {
64        match self {
65            Self::CompletionError(error) => error.provider_response_body(),
66            _ => None,
67        }
68    }
69
70    /// Parses a wrapped provider response body as JSON when present.
71    pub fn provider_response_json(&self) -> Result<Option<serde_json::Value>, serde_json::Error> {
72        match self {
73            Self::CompletionError(error) => error.provider_response_json(),
74            _ => Ok(None),
75        }
76    }
77
78    /// Returns the HTTP status exposed by a wrapped completion error.
79    pub fn provider_response_status(&self) -> Option<http::StatusCode> {
80        match self {
81            Self::CompletionError(error) => error.provider_response_status(),
82            _ => None,
83        }
84    }
85
86    pub(crate) fn prompt_cancelled(
87        chat_history: impl IntoIterator<Item = Message>,
88        reason: impl Into<String>,
89    ) -> Self {
90        Self::PromptCancelled {
91            chat_history: chat_history.into_iter().collect(),
92            reason: reason.into(),
93        }
94    }
95}
96
97/// Errors returned by typed structured prompting.
98#[derive(Debug, Error)]
99#[non_exhaustive]
100pub enum StructuredOutputError {
101    /// The underlying classic run failed.
102    #[error("PromptError: {0}")]
103    PromptError(#[from] Box<PromptError>),
104    /// The accepted response could not be deserialized.
105    #[error("DeserializationError: {0}")]
106    DeserializationError(#[from] serde_json::Error),
107    /// The model returned no accepted content.
108    #[error("EmptyResponse: model returned no content")]
109    EmptyResponse,
110}
111
112impl StructuredOutputError {
113    /// Returns the provider response body exposed through the wrapped prompt error.
114    pub fn provider_response_body(&self) -> Option<&str> {
115        match self {
116            Self::PromptError(error) => error.provider_response_body(),
117            _ => None,
118        }
119    }
120
121    /// Parses the wrapped provider response body as JSON when present.
122    pub fn provider_response_json(&self) -> Result<Option<serde_json::Value>, serde_json::Error> {
123        match self {
124            Self::PromptError(error) => error.provider_response_json(),
125            _ => Ok(None),
126        }
127    }
128
129    /// Returns the provider HTTP status exposed through the wrapped prompt error.
130    pub fn provider_response_status(&self) -> Option<http::StatusCode> {
131        match self {
132            Self::PromptError(error) => error.provider_response_status(),
133            _ => None,
134        }
135    }
136}
137
138/// High-level one-shot prompting for the classic runtime.
139pub trait Prompt: WasmCompatSend + WasmCompatSync {
140    /// Send a prompt and return accepted assistant text after runtime orchestration.
141    fn prompt(
142        &self,
143        prompt: impl Into<Message> + WasmCompatSend,
144    ) -> impl std::future::IntoFuture<Output = Result<String, PromptError>, IntoFuture: WasmCompatSend>;
145}
146
147/// High-level prompting with caller-owned canonical chat history.
148pub trait Chat: WasmCompatSend + WasmCompatSync {
149    /// Execute one turn and append only committed messages to `chat_history`.
150    fn chat(
151        &self,
152        prompt: impl Into<Message> + WasmCompatSend,
153        chat_history: &mut Vec<Message>,
154    ) -> impl std::future::Future<Output = Result<String, PromptError>> + WasmCompatSend;
155}
156
157/// High-level typed structured prompting for the classic runtime.
158pub trait TypedPrompt: WasmCompatSend + WasmCompatSync {
159    /// Request type returned for one target output type.
160    type TypedRequest<T>: std::future::IntoFuture<Output = Result<T, StructuredOutputError>>
161    where
162        T: schemars::JsonSchema + DeserializeOwned + WasmCompatSend + 'static;
163
164    /// Send a prompt and deserialize the accepted structured response as `T`.
165    fn prompt_typed<T>(&self, prompt: impl Into<Message> + WasmCompatSend) -> Self::TypedRequest<T>
166    where
167        T: schemars::JsonSchema + DeserializeOwned + WasmCompatSend;
168}
169
170#[cfg(test)]
171mod provider_response_tests {
172    use rig_core::{ProviderResponseError, http_client};
173
174    use super::*;
175
176    #[test]
177    fn prompt_error_forwards_provider_response_to_completion_error() {
178        let body = r#"{"error":{"message":"boom"}}"#;
179        let inner =
180            CompletionError::from_http_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
181        let error = PromptError::CompletionError(inner);
182
183        assert_eq!(
184            error.provider_response_status(),
185            Some(http::StatusCode::SERVICE_UNAVAILABLE),
186        );
187        assert_eq!(error.provider_response_body(), Some(body));
188        assert_eq!(
189            error
190                .provider_response_json()
191                .expect("valid json")
192                .expect("present json")["error"]["message"],
193            "boom",
194        );
195    }
196
197    #[test]
198    fn prompt_error_provider_response_helpers_forward_http_status_and_body() {
199        let body = r#"{"error":{"message":"unauthorized"}}"#;
200        let error = PromptError::CompletionError(CompletionError::HttpError(
201            http_client::Error::InvalidStatusCodeWithMessage(
202                http::StatusCode::UNAUTHORIZED,
203                body.to_string(),
204            ),
205        ));
206
207        assert_eq!(error.provider_response_body(), Some(body));
208        assert_eq!(
209            error.provider_response_status(),
210            Some(http::StatusCode::UNAUTHORIZED)
211        );
212        assert_eq!(
213            error.provider_response_json().expect("valid JSON body"),
214            Some(serde_json::json!({
215                "error": { "message": "unauthorized" }
216            }))
217        );
218    }
219
220    #[test]
221    fn prompt_error_provider_response_helpers_forward_wrapped_completion_error() {
222        let body = r#"{"error":{"code":"invalid_request","message":"bad input"}}"#;
223        let error = PromptError::CompletionError(CompletionError::ProviderResponse(
224            ProviderResponseError {
225                status: None,
226                body: body.to_string(),
227            },
228        ));
229
230        assert_eq!(error.provider_response_body(), Some(body));
231        assert_eq!(error.provider_response_status(), None);
232        assert_eq!(
233            error.provider_response_json().expect("valid JSON body"),
234            Some(serde_json::json!({
235                "error": {
236                    "code": "invalid_request",
237                    "message": "bad input"
238                }
239            }))
240        );
241    }
242
243    #[test]
244    fn prompt_error_provider_response_helpers_return_none_for_unrelated_variant() {
245        let error = PromptError::PromptCancelled {
246            chat_history: vec![Message::user("hi")],
247            reason: "cancelled".to_string(),
248        };
249
250        assert_eq!(error.provider_response_body(), None);
251        assert_eq!(error.provider_response_status(), None);
252        assert_eq!(
253            error
254                .provider_response_json()
255                .expect("no body is not an error"),
256            None
257        );
258    }
259
260    #[test]
261    fn structured_output_error_provider_response_helpers_forward_prompt_error() {
262        let body = r#"{"error":{"message":"bad input"}}"#;
263        let error = StructuredOutputError::PromptError(Box::new(PromptError::CompletionError(
264            CompletionError::ProviderResponse(ProviderResponseError {
265                status: Some(http::StatusCode::BAD_REQUEST),
266                body: body.to_string(),
267            }),
268        )));
269
270        assert_eq!(error.provider_response_body(), Some(body));
271        assert_eq!(
272            error.provider_response_status(),
273            Some(http::StatusCode::BAD_REQUEST)
274        );
275    }
276}
rig_agent/completion.rs

rig_agent/
completion.rs
