Skip to main content

origin_ai/
request.rs

1use serde::{Deserialize, Serialize};
2
3/// What the application asks a model to do.
4///
5/// Deliberately not a chat transcript: most product features are one instruction over
6/// one piece of content, and modelling them as conversations invites state nobody
7/// needs.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct Prompt {
10    /// What the model should do.
11    pub instruction: String,
12    /// What it should do it to.
13    pub input: String,
14    /// Upper bound on the answer. A product that does not cap this is one bad prompt
15    /// away from a surprising bill.
16    pub max_output_tokens: u32,
17    /// `0.0` for extraction and classification, higher for drafting.
18    pub temperature: f32,
19}
20
21impl Prompt {
22    /// A deterministic prompt: temperature zero, short answer.
23    ///
24    /// The right default for classification and extraction, which is most of what a
25    /// desktop application actually needs.
26    pub fn new(instruction: impl Into<String>, input: impl Into<String>) -> Self {
27        Self {
28            instruction: instruction.into(),
29            input: input.into(),
30            max_output_tokens: 512,
31            temperature: 0.0,
32        }
33    }
34
35    pub fn with_max_output_tokens(mut self, max_output_tokens: u32) -> Self {
36        self.max_output_tokens = max_output_tokens;
37        self
38    }
39
40    /// Allow variation. Use for drafting, never for extraction.
41    pub fn with_temperature(mut self, temperature: f32) -> Self {
42        self.temperature = temperature.clamp(0.0, 2.0);
43        self
44    }
45}
46
47/// What a provider reports about cost.
48///
49/// Present so a product can show the user what its AI features consumed. A feature
50/// that cannot account for itself is one users learn to distrust.
51#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
52pub struct Usage {
53    pub input_tokens: u32,
54    pub output_tokens: u32,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct Completion {
59    pub text: String,
60    pub usage: Usage,
61    /// Which model answered. Recorded because the same prompt behaves differently
62    /// across models and versions, and "it used to work" needs an answer.
63    pub model: String,
64    /// `true` when the answer hit `max_output_tokens` and was cut off.
65    pub truncated: bool,
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn the_default_prompt_is_deterministic() {
74        let prompt = Prompt::new("Summarise", "text");
75
76        assert_eq!(prompt.temperature, 0.0);
77        assert_eq!(prompt.max_output_tokens, 512);
78    }
79
80    #[test]
81    fn temperature_is_clamped_rather_than_trusted() {
82        assert_eq!(Prompt::new("x", "y").with_temperature(9.0).temperature, 2.0);
83        assert_eq!(
84            Prompt::new("x", "y").with_temperature(-1.0).temperature,
85            0.0
86        );
87    }
88}