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 validate_configuration(&self) -> Result<(), Error> {
241        if let (Some(min), Some(max)) = (self.min_items, self.max_items)
242            && min > max
243        {
244            return Err(Error::invalid_argument(
245                "min_items",
246                "min_items must not exceed max_items",
247            ));
248        }
249        Ok(())
250    }
251
252    fn max_elements(&self) -> Option<usize> {
253        self.max_items
254    }
255
256    fn response_format(&self) -> Option<ResponseFormat> {
257        Some(ResponseFormat::Json {
258            schema: Some(self.wrapper_schema()),
259            name: self.name.clone(),
260            description: self.description.clone(),
261        })
262    }
263
264    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<Vec<T>, Error> {
265        let value = parse_json(text, ctx)?;
266        let Some(elements) = Self::elements_of(&value) else {
267            return Err(no_object(
268                "response must be an object with an elements array",
269                text,
270                ctx,
271                None,
272            ));
273        };
274        if let Some(min) = self.min_items
275            && elements.len() < min
276        {
277            return Err(no_object(
278                &format!("elements array must contain at least {min} items"),
279                text,
280                ctx,
281                None,
282            ));
283        }
284        if let Some(max) = self.max_items
285            && elements.len() > max
286        {
287            return Err(no_object(
288                &format!("elements array must contain at most {max} items"),
289                text,
290                ctx,
291                None,
292            ));
293        }
294        elements
295            .iter()
296            .map(|element| {
297                self.element.validate(element.clone()).map_err(|error| {
298                    no_object(
299                        "response did not match schema",
300                        text,
301                        ctx,
302                        Some(Box::new(error)),
303                    )
304                })
305            })
306            .collect()
307    }
308
309    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
310        self.parse_elements(text).map(JsonValue::Array)
311    }
312
313    fn typed_partial(&self, value: &JsonValue) -> Option<Vec<T>> {
314        value
315            .as_array()?
316            .iter()
317            .map(|element| self.element.validate(element.clone()).ok())
318            .collect()
319    }
320
321    fn parse_elements(&self, text: &str) -> Option<Vec<JsonValue>> {
322        let (value, state) = partial_value(text)?;
323        let elements = Self::elements_of(&value)?;
324        let complete = match state {
325            PartialParseState::RepairedParse if !elements.is_empty() => {
326                &elements[..elements.len() - 1]
327            }
328            _ => elements.as_slice(),
329        };
330        let mut validated = Vec::with_capacity(complete.len());
331        for element in complete {
332            if self.element.validate(element.clone()).is_ok() {
333                validated.push(element.clone());
334            }
335        }
336        Some(validated)
337    }
338
339    fn parse_typed_elements(&self, text: &str) -> Option<Vec<T>> {
340        let (value, state) = partial_value(text)?;
341        let elements = Self::elements_of(&value)?;
342        let complete = match state {
343            PartialParseState::RepairedParse if !elements.is_empty() => {
344                &elements[..elements.len() - 1]
345            }
346            _ => elements.as_slice(),
347        };
348        Some(
349            complete
350                .iter()
351                .filter_map(|element| self.element.validate(element.clone()).ok())
352                .collect(),
353        )
354    }
355}
356
357/// One of a fixed set of string options.
358#[derive(Debug, Clone)]
359pub struct ChoiceOutput {
360    options: Vec<String>,
361    name: Option<String>,
362    description: Option<String>,
363}
364
365impl ChoiceOutput {
366    /// Creates the strategy.
367    #[must_use]
368    pub fn new(options: impl IntoIterator<Item = impl Into<String>>) -> Self {
369        Self {
370            options: options.into_iter().map(Into::into).collect(),
371            name: None,
372            description: None,
373        }
374    }
375
376    /// Names the output.
377    #[must_use]
378    pub fn with_name(mut self, name: impl Into<String>) -> Self {
379        self.name = Some(name.into());
380        self
381    }
382
383    /// Describes the output.
384    #[must_use]
385    pub fn with_description(mut self, description: impl Into<String>) -> Self {
386        self.description = Some(description.into());
387        self
388    }
389
390    fn result_of(value: &JsonValue) -> Option<&str> {
391        value.as_object()?.get("result")?.as_str()
392    }
393}
394
395impl OutputHandler<String> for ChoiceOutput {
396    fn response_format(&self) -> Option<ResponseFormat> {
397        Some(ResponseFormat::Json {
398            schema: Some(json!({
399                "$schema": "http://json-schema.org/draft-07/schema#",
400                "type": "object",
401                "properties": { "result": { "type": "string", "enum": self.options } },
402                "required": ["result"],
403                "additionalProperties": false,
404            })),
405            name: self.name.clone(),
406            description: self.description.clone(),
407        })
408    }
409
410    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<String, Error> {
411        let value = parse_json(text, ctx)?;
412        let Some(result) = Self::result_of(&value) else {
413            return Err(no_object(
414                "response must be an object with a result string",
415                text,
416                ctx,
417                None,
418            ));
419        };
420        if !self.options.iter().any(|option| option == result) {
421            return Err(no_object(
422                "response did not match one of the options",
423                text,
424                ctx,
425                None,
426            ));
427        }
428        Ok(result.to_owned())
429    }
430
431    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
432        let (value, state) = partial_value(text)?;
433        let partial = Self::result_of(&value)?;
434        let matches: Vec<&String> = self
435            .options
436            .iter()
437            .filter(|option| option.starts_with(partial))
438            .collect();
439        match state {
440            PartialParseState::SuccessfulParse => matches
441                .iter()
442                .any(|option| option.as_str() == partial)
443                .then(|| JsonValue::String(partial.to_owned())),
444            _ => (matches.len() == 1).then(|| JsonValue::String(matches[0].clone())),
445        }
446    }
447
448    fn typed_partial(&self, value: &JsonValue) -> Option<String> {
449        value.as_str().map(str::to_owned)
450    }
451}
452
453/// Unconstrained (or raw-schema-validated) JSON.
454#[derive(Debug, Clone)]
455pub struct JsonOutput {
456    schema: Option<Schema<JsonValue>>,
457    name: Option<String>,
458    description: Option<String>,
459}
460
461impl JsonOutput {
462    /// Creates the strategy; `schema` is a raw JSON schema.
463    #[must_use]
464    pub fn new(schema: Option<JsonValue>) -> Self {
465        Self {
466            schema: schema.map(Schema::from_json_schema),
467            name: None,
468            description: None,
469        }
470    }
471
472    /// Names the output.
473    #[must_use]
474    pub fn with_name(mut self, name: impl Into<String>) -> Self {
475        self.name = Some(name.into());
476        self
477    }
478
479    /// Describes the output.
480    #[must_use]
481    pub fn with_description(mut self, description: impl Into<String>) -> Self {
482        self.description = Some(description.into());
483        self
484    }
485}
486
487impl OutputHandler<JsonValue> for JsonOutput {
488    fn response_format(&self) -> Option<ResponseFormat> {
489        Some(ResponseFormat::Json {
490            schema: self
491                .schema
492                .as_ref()
493                .map(|schema| schema.json_schema().clone()),
494            name: self.name.clone(),
495            description: self.description.clone(),
496        })
497    }
498
499    fn parse_complete(&self, text: &str, ctx: &OutputContext) -> Result<JsonValue, Error> {
500        let value = parse_json(text, ctx)?;
501        match &self.schema {
502            Some(schema) => schema.validate(value).map_err(|error| {
503                no_object(
504                    "response did not match schema",
505                    text,
506                    ctx,
507                    Some(Box::new(error)),
508                )
509            }),
510            None => Ok(value),
511        }
512    }
513
514    fn parse_partial(&self, text: &str) -> Option<JsonValue> {
515        partial_value(text).map(|(value, _)| value)
516    }
517
518    fn typed_partial(&self, value: &JsonValue) -> Option<JsonValue> {
519        match &self.schema {
520            Some(schema) => schema.validate(value.clone()).ok(),
521            None => Some(value.clone()),
522        }
523    }
524}
ferrin_core/output/strategies.rs

ferrin_core/output/
strategies.rs
