1use std::collections::BTreeMap;
2
3use schemars::{JsonSchema, schema_for};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::{ExtensionMap, Message};
8
9#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
11pub struct ModelRef {
12 pub provider: String,
14 pub name: String,
16}
17
18impl ModelRef {
19 pub fn new(provider: impl Into<String>, name: impl Into<String>) -> Self {
21 Self {
22 provider: provider.into(),
23 name: name.into(),
24 }
25 }
26}
27
28#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
30#[serde(rename_all = "snake_case")]
31#[non_exhaustive]
32pub enum FeaturePolicy {
33 #[default]
35 Strict,
36 AllowEmulation,
38 BestEffort,
40}
41
42#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
44pub struct GenerationOptions {
45 pub temperature: Option<f64>,
47 pub top_p: Option<f64>,
49 pub max_output_tokens: Option<u64>,
51 pub seed: Option<u64>,
53 pub stop: Vec<String>,
55}
56
57#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
59#[serde(tag = "type", rename_all = "snake_case")]
60#[non_exhaustive]
61pub enum OutputFormat {
62 #[default]
64 Text,
65 Json,
67 JsonSchema {
69 name: String,
71 schema: Value,
73 strict: bool,
75 },
76}
77
78impl OutputFormat {
79 pub fn typed<T>(name: impl Into<String>) -> Self
84 where
85 T: JsonSchema,
86 {
87 Self::JsonSchema {
88 name: name.into(),
89 schema: schema_for!(T).to_value(),
90 strict: true,
91 }
92 }
93
94 pub fn typed_with_strictness<T>(name: impl Into<String>, strict: bool) -> Self
97 where
98 T: JsonSchema,
99 {
100 Self::JsonSchema {
101 name: name.into(),
102 schema: schema_for!(T).to_value(),
103 strict,
104 }
105 }
106}
107
108#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
110pub struct ToolSpec {
111 pub name: String,
113 pub description: String,
115 pub input_schema: Value,
117 pub output_schema: Option<Value>,
119 pub metadata: ExtensionMap,
121}
122
123#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
125#[serde(tag = "type", rename_all = "snake_case")]
126#[non_exhaustive]
127pub enum ToolChoice {
128 #[default]
130 Auto,
131 None,
133 Required,
135 Named {
137 name: String,
139 },
140}
141
142#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
144pub struct ModelRequest {
145 pub model: ModelRef,
147 pub messages: Vec<Message>,
149 pub tools: Vec<ToolSpec>,
151 pub tool_choice: ToolChoice,
153 pub output_format: OutputFormat,
155 pub generation: GenerationOptions,
157 pub feature_policy: FeaturePolicy,
159 pub provider_options: BTreeMap<String, Value>,
161 pub metadata: ExtensionMap,
163}
164
165impl ModelRequest {
166 pub fn new(model: ModelRef, message: Message) -> Self {
168 Self {
169 model,
170 messages: vec![message],
171 tools: Vec::new(),
172 tool_choice: ToolChoice::Auto,
173 output_format: OutputFormat::Text,
174 generation: GenerationOptions::default(),
175 feature_policy: FeaturePolicy::Strict,
176 provider_options: BTreeMap::new(),
177 metadata: BTreeMap::new(),
178 }
179 }
180
181 #[must_use]
183 pub fn message(mut self, message: Message) -> Self {
184 self.messages.push(message);
185 self
186 }
187
188 #[must_use]
190 pub fn tool(mut self, tool: ToolSpec) -> Self {
191 self.tools.push(tool);
192 self
193 }
194
195 #[must_use]
197 pub fn output_format(mut self, output_format: OutputFormat) -> Self {
198 self.output_format = output_format;
199 self
200 }
201
202 #[must_use]
204 pub fn structured_output<T>(self, name: impl Into<String>) -> Self
205 where
206 T: JsonSchema,
207 {
208 self.output_format(OutputFormat::typed::<T>(name))
209 }
210
211 #[must_use]
213 pub const fn feature_policy(mut self, feature_policy: FeaturePolicy) -> Self {
214 self.feature_policy = feature_policy;
215 self
216 }
217}