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

ferrin_core/output/
strategies.rs
