Skip to main content

systemprompt_models/ai/
models.rs

1//! Model configuration for tool-calling and chat requests.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
10pub struct ToolModelConfig {
11    #[serde(skip_serializing_if = "Option::is_none")]
12    pub provider: Option<String>,
13    #[serde(skip_serializing_if = "Option::is_none")]
14    pub model: Option<String>,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub max_output_tokens: Option<u32>,
17}
18
19impl ToolModelConfig {
20    pub fn new(provider: impl Into<String>, model: impl Into<String>) -> Self {
21        Self {
22            provider: Some(provider.into()),
23            model: Some(model.into()),
24            max_output_tokens: None,
25        }
26    }
27
28    pub const fn with_max_output_tokens(mut self, tokens: u32) -> Self {
29        self.max_output_tokens = Some(tokens);
30        self
31    }
32
33    pub const fn is_empty(&self) -> bool {
34        self.provider.is_none() && self.model.is_none() && self.max_output_tokens.is_none()
35    }
36
37    pub fn merge_with(&self, other: &Self) -> Self {
38        Self {
39            provider: other.provider.as_ref().or(self.provider.as_ref()).cloned(),
40            model: other.model.as_ref().or(self.model.as_ref()).cloned(),
41            max_output_tokens: other.max_output_tokens.or(self.max_output_tokens),
42        }
43    }
44}
45
46pub type ToolModelOverrides = HashMap<String, HashMap<String, ToolModelConfig>>;
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ModelConfig {
50    pub id: String,
51    pub max_tokens: u32,
52    pub supports_tools: bool,
53    #[serde(default)]
54    pub cost_per_1k_tokens: f32,
55}
56
57impl ModelConfig {
58    pub fn new(id: impl Into<String>, max_tokens: u32, supports_tools: bool) -> Self {
59        Self {
60            id: id.into(),
61            max_tokens,
62            supports_tools,
63            cost_per_1k_tokens: 0.0,
64        }
65    }
66
67    pub const fn with_cost(mut self, cost: f32) -> Self {
68        self.cost_per_1k_tokens = cost;
69        self
70    }
71}