ferrin_core/output/
strategies.rs

1//! Built-in output strategies.
2//!
3//! Derived from the Vercel AI SDK (Apache-2.0, Copyright 2023 Vercel, Inc.),
4//! translated from TypeScript to Rust and modified; see `NOTICE`.
5
6use std::fmt;
7
8use ferrin_schema::Schema;
9use ferrin_schema::partial_json::PartialParseState;
10use ferrin_schema::partial_json::parse_partial;
11use ferrin_spec::JsonValue;
12use ferrin_spec::ResponseFormat;
13use serde::de::DeserializeOwned;
14use serde_json::json;
15
16use super::OutputContext;
17use super::OutputHandler;
18use crate::error::Error;
19use crate::error::NoObjectGeneratedDetails;
20
21fn no_object(
22    message: &str,
23    text: &str,
24    ctx: &OutputContext,
25    cause: Option<crate::error::BoxError>,
26) -> Error {
27    Error::no_object_generated(NoObjectGeneratedDetails {
28        message: message.to_owned(),
29        text: Some(text.to_owned()),
30        response: ctx.response.clone(),
31        usage: ctx.usage.clone(),
32        finish_reason: ctx.finish_reason.clone(),
33        cause,
34    })
35}
36
37fn parse_json(text: &str, ctx: &OutputContext) -> Result<JsonValue, Error> {
38    ferrin_schema::json::parse(text).map_err(|error| {
39        no_object(
40            "could not parse the response",
41            text,
42            ctx,
43            Some(Box::new(error)),
44        )
45    })
46}
47
48fn partial_value(text: &str) -> Option<(JsonValue, PartialParseState)> {
49    let parsed = parse_partial(text);
50    match parsed.state {
51        PartialParseState::FailedParse => None,
52        state => parsed.value.map(|value| (value, state)),
53    }
54}
55
56/// Plain text.
57#[derive(Debug, Clone, Copy, Default)]
58pub struct TextOutput;
59
60impl OutputHandler<String> for TextOutput {
61    fn response_format(&self) -> Option<ResponseFormat> {
62        Some(ResponseFormat::Text)
63    }
64
65    fn parse_complete(&self, text: &str, _ctx: &OutputContext) -> Result<String, Error> {
66        Ok(text.to_owned())
67    }
68
69    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
70        Some(JsonValue::String(text.to_owned()))
71    }
72
73    fn typed_partial(&self, value: &JsonValue) -> Option<String> {
74        value.as_str().map(str::to_owned)
75    }
76}
77
78/// A JSON object validated by a schema.
79pub struct ObjectOutput<T> {
80    schema: Schema<T>,
81    name: Option<String>,
82    description: Option<String>,
83}
84
85impl<T> ObjectOutput<T> {
86    /// Creates the strategy.
87    #[must_use]
88    pub fn new(schema: Schema<T>) -> Self {
89        Self {
90            schema,
91            name: None,
92            description: None,
93        }
94    }
95
96    /// Names the output for providers that accept a schema name.
97    #[must_use]
98    pub fn with_name(mut self, name: impl Into<String>) -> Self {
99        self.name = Some(name.into());
100        self
101    }
102
103    /// Describes the output for providers that accept a schema description.
104    #[must_use]
105    pub fn with_description(mut self, description: impl Into<String>) -> Self {
106        self.description = Some(description.into());
107        self
108    }
109}
110
111impl<T> fmt::Debug for ObjectOutput<T> {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        f.debug_struct("ObjectOutput")
114            .field("schema", self.schema.json_schema())
115            .field("name", &self.name)
116            .field("description", &self.description)
117            .finish()
118    }
119}
120
121impl<T: DeserializeOwned + Send + Sync + 'static> OutputHandler<T> for ObjectOutput<T> {
122    fn response_format(&self) -> Option<ResponseFormat> {
123        Some(ResponseFormat::Json {
124            schema: Some(self.schema.json_schema().clone()),
125            name: self.name.clone(),
126            description: self.description.clone(),
127        })
128    }
129
130    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<T, Error> {
131        let value = parse_json(text, ctx)?;
132        self.schema.validate(value).map_err(|error| {
133            no_object(
134                "response did not match schema",
135                text,
136                ctx,
137                Some(Box::new(error)),
138            )
139        })
140    }
141
142    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
143        partial_value(text).map(|(value, _)| value)
144    }
145
146    fn typed_partial(&self, value: &JsonValue) -> Option<T> {
147        self.schema.validate(value.clone()).ok()
148    }
149}
150
151/// An array of schema-validated elements, wrapped in `{ "elements": [...] }`
152/// on the wire.
153pub struct ArrayOutput<T> {
154    element: Schema<T>,
155    min_items: Option<usize>,
156    max_items: Option<usize>,
157    name: Option<String>,
158    description: Option<String>,
159}
160
161impl<T> ArrayOutput<T> {
162    /// Creates the strategy.
163    #[must_use]
164    pub fn new(element: Schema<T>) -> Self {
165        Self {
166            element,
167            min_items: None,
168            max_items: None,
169            name: None,
170            description: None,
171        }
172    }
173
174    /// Requires at least `n` elements.
175    #[must_use]
176    pub fn min_items(mut self, n: usize) -> Self {
177        self.min_items = Some(n);
178        self
179    }
180
181    /// Allows at most `n` elements.
182    #[must_use]
183    pub fn max_items(mut self, n: usize) -> Self {
184        self.max_items = Some(n);
185        self
186    }
187
188    /// Names the output.
189    #[must_use]
190    pub fn with_name(mut self, name: impl Into<String>) -> Self {
191        self.name = Some(name.into());
192        self
193    }
194
195    /// Describes the output.
196    #[must_use]
197    pub fn with_description(mut self, description: impl Into<String>) -> Self {
198        self.description = Some(description.into());
199        self
200    }
201
202    fn wrapper_schema(&self) -> JsonValue {
203        let mut element = self.element.json_schema().clone();
204        super::local_refs::relocate(&mut element, "#/properties/elements/items");
205        let mut elements = json!({
206            "type": "array",
207            "items": element,
208        });
209        if let Some(min) = self.min_items {
210            elements["minItems"] = json!(min);
211        }
212        if let Some(max) = self.max_items {
213            elements["maxItems"] = json!(max);
214        }
215        json!({
216            "$schema": "http://json-schema.org/draft-07/schema#",
217            "type": "object",
218            "properties": { "elements": elements },
219            "required": ["elements"],
220            "additionalProperties": false,
221        })
222    }
223
224    fn elements_of(value: &JsonValue) -> Option<&Vec<JsonValue>> {
225        value.as_object()?.get("elements")?.as_array()
226    }
227}
228
229impl<T> fmt::Debug for ArrayOutput<T> {
230    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
231        f.debug_struct("ArrayOutput")
232            .field("element", self.element.json_schema())
233            .field("min_items", &self.min_items)
234            .field("max_items", &self.max_items)
235            .finish_non_exhaustive()
236    }
237}
238
239impl<T: DeserializeOwned + Send + Sync + 'static> OutputHandler<Vec<T>> for ArrayOutput<T> {
240    fn response_format(&self) -> Option<ResponseFormat> {
241        Some(ResponseFormat::Json {
242            schema: Some(self.wrapper_schema()),
243            name: self.name.clone(),
244            description: self.description.clone(),
245        })
246    }
247
248    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<Vec<T>, Error> {
249        let value = parse_json(text, ctx)?;
250        let Some(elements) = Self::elements_of(&value) else {
251            return Err(no_object(
252                "response must be an object with an elements array",
253                text,
254                ctx,
255                None,
256            ));
257        };
258        if let Some(min) = self.min_items
259            && elements.len() < min
260        {
261            return Err(no_object(
262                &format!("elements array must contain at least {min} items"),
263                text,
264                ctx,
265                None,
266            ));
267        }
268        if let Some(max) = self.max_items
269            && elements.len() > max
270        {
271            return Err(no_object(
272                &format!("elements array must contain at most {max} items"),
273                text,
274                ctx,
275                None,
276            ));
277        }
278        elements
279            .iter()
280            .map(|element| {
281                self.element.validate(element.clone()).map_err(|error| {
282                    no_object(
283                        "response did not match schema",
284                        text,
285                        ctx,
286                        Some(Box::new(error)),
287                    )
288                })
289            })
290            .collect()
291    }
292
293    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
294        self.parse_elements(text).map(JsonValue::Array)
295    }
296
297    fn typed_partial(&self, value: &JsonValue) -> Option<Vec<T>> {
298        value
299            .as_array()?
300            .iter()
301            .map(|element| self.element.validate(element.clone()).ok())
302            .collect()
303    }
304
305    fn parse_elements(&self, text: &str) -> Option<Vec<JsonValue>> {
306        let (value, state) = partial_value(text)?;
307        let elements = Self::elements_of(&value)?;
308        let complete = match state {
309            PartialParseState::RepairedParse if !elements.is_empty() => {
310                &elements[..elements.len() - 1]
311            }
312            _ => elements.as_slice(),
313        };
314        let mut validated = Vec::with_capacity(complete.len());
315        for element in complete {
316            if self.element.validate(element.clone()).is_err() {
317                return None;
318            }
319            validated.push(element.clone());
320        }
321        Some(validated)
322    }
323}
324
325/// One of a fixed set of string options.
326#[derive(Debug, Clone)]
327pub struct ChoiceOutput {
328    options: Vec<String>,
329    name: Option<String>,
330    description: Option<String>,
331}
332
333impl ChoiceOutput {
334    /// Creates the strategy.
335    #[must_use]
336    pub fn new(options: impl IntoIterator<Item = impl Into<String>>) -> Self {
337        Self {
338            options: options.into_iter().map(Into::into).collect(),
339            name: None,
340            description: None,
341        }
342    }
343
344    /// Names the output.
345    #[must_use]
346    pub fn with_name(mut self, name: impl Into<String>) -> Self {
347        self.name = Some(name.into());
348        self
349    }
350
351    /// Describes the output.
352    #[must_use]
353    pub fn with_description(mut self, description: impl Into<String>) -> Self {
354        self.description = Some(description.into());
355        self
356    }
357
358    fn result_of(value: &JsonValue) -> Option<&str> {
359        value.as_object()?.get("result")?.as_str()
360    }
361}
362
363impl OutputHandler<String> for ChoiceOutput {
364    fn response_format(&self) -> Option<ResponseFormat> {
365        Some(ResponseFormat::Json {
366            schema: Some(json!({
367                "$schema": "http://json-schema.org/draft-07/schema#",
368                "type": "object",
369                "properties": { "result": { "type": "string", "enum": self.options } },
370                "required": ["result"],
371                "additionalProperties": false,
372            })),
373            name: self.name.clone(),
374            description: self.description.clone(),
375        })
376    }
377
378    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<String, Error> {
379        let value = parse_json(text, ctx)?;
380        let Some(result) = Self::result_of(&value) else {
381            return Err(no_object(
382                "response must be an object with a result string",
383                text,
384                ctx,
385                None,
386            ));
387        };
388        if !self.options.iter().any(|option| option == result) {
389            return Err(no_object(
390                "response did not match one of the options",
391                text,
392                ctx,
393                None,
394            ));
395        }
396        Ok(result.to_owned())
397    }
398
399    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
400        let (value, state) = partial_value(text)?;
401        let partial = Self::result_of(&value)?;
402        let matches: Vec<&String> = self
403            .options
404            .iter()
405            .filter(|option| option.starts_with(partial))
406            .collect();
407        match state {
408            PartialParseState::SuccessfulParse => matches
409                .iter()
410                .any(|option| option.as_str() == partial)
411                .then(|| JsonValue::String(partial.to_owned())),
412            _ => (matches.len() == 1).then(|| JsonValue::String(matches[0].clone())),
413        }
414    }
415
416    fn typed_partial(&self, value: &JsonValue) -> Option<String> {
417        value.as_str().map(str::to_owned)
418    }
419}
420
421/// Unconstrained (or raw-schema-validated) JSON.
422#[derive(Debug, Clone)]
423pub struct JsonOutput {
424    schema: Option<Schema<JsonValue>>,
425    name: Option<String>,
426    description: Option<String>,
427}
428
429impl JsonOutput {
430    /// Creates the strategy; `schema` is a raw JSON schema.
431    #[must_use]
432    pub fn new(schema: Option<JsonValue>) -> Self {
433        Self {
434            schema: schema.map(Schema::from_json_schema),
435            name: None,
436            description: None,
437        }
438    }
439
440    /// Names the output.
441    #[must_use]
442    pub fn with_name(mut self, name: impl Into<String>) -> Self {
443        self.name = Some(name.into());
444        self
445    }
446
447    /// Describes the output.
448    #[must_use]
449    pub fn with_description(mut self, description: impl Into<String>) -> Self {
450        self.description = Some(description.into());
451        self
452    }
453}
454
455impl OutputHandler<JsonValue> for JsonOutput {
456    fn response_format(&self) -> Option<ResponseFormat> {
457        Some(ResponseFormat::Json {
458            schema: self
459                .schema
460                .as_ref()
461                .map(|schema| schema.json_schema().clone()),
462            name: self.name.clone(),
463            description: self.description.clone(),
464        })
465    }
466
467    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<JsonValue, Error> {
468        let value = parse_json(text, ctx)?;
469        match &self.schema {
470            Some(schema) => schema.validate(value).map_err(|error| {
471                no_object(
472                    "response did not match schema",
473                    text,
474                    ctx,
475                    Some(Box::new(error)),
476                )
477            }),
478            None => Ok(value),
479        }
480    }
481
482    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
483        partial_value(text).map(|(value, _)| value)
484    }
485
486    fn typed_partial(&self, value: &JsonValue) -> Option<JsonValue> {
487        Some(value.clone())
488    }
489}
ferrin_core/output/strategies.rs

ferrin_core/output/
strategies.rs
