1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use thiserror::Error;
5
6use crate::config::{
7 build_vv_llm_from_local_settings, resolve_model_endpoint, ResolvedModelConfig,
8};
9use crate::llm::{LlmClient, LlmRequest, ScriptStep, ScriptedLlmClient, VvLlmClient};
10use crate::model_settings::ModelSettings;
11use crate::types::LLMResponse;
12
13#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
14#[serde(tag = "type", rename_all = "snake_case")]
15pub enum ModelRef {
16 Named(String),
17 BackendModel { backend: String, model: String },
18 Resolved(ResolvedModelConfig),
19}
20
21impl ModelRef {
22 pub fn named(model: impl Into<String>) -> Self {
23 Self::Named(model.into())
24 }
25
26 pub fn backend(backend: impl Into<String>, model: impl Into<String>) -> Self {
27 Self::BackendModel {
28 backend: backend.into(),
29 model: model.into(),
30 }
31 }
32
33 pub fn resolved(resolved: ResolvedModelConfig) -> Self {
34 Self::Resolved(resolved)
35 }
36
37 pub fn model(&self) -> &str {
38 match self {
39 Self::Named(model) => model,
40 Self::BackendModel { model, .. } => model,
41 Self::Resolved(resolved) => resolved.selected_model.as_str(),
42 }
43 }
44
45 pub fn backend_name(&self) -> Option<&str> {
46 match self {
47 Self::Named(_) => None,
48 Self::BackendModel { backend, .. } => Some(backend),
49 Self::Resolved(resolved) => Some(resolved.backend.as_str()),
50 }
51 }
52}
53
54#[derive(Debug, Error)]
55pub enum ModelError {
56 #[error("model provider has no default backend for named model `{0}`")]
57 MissingDefaultBackend(String),
58 #[error("model backend mismatch: requested `{requested}`, provider `{provider}`")]
59 BackendMismatch { requested: String, provider: String },
60 #[error("{0}")]
61 Config(String),
62}
63
64pub trait ModelProvider: Send + Sync {
65 fn resolve(&self, model: &ModelRef) -> Result<ResolvedModelConfig, ModelError>;
66 fn client(&self, resolved: &ResolvedModelConfig) -> Result<Arc<dyn LlmClient>, ModelError>;
67
68 fn default_settings(&self, _resolved: &ResolvedModelConfig) -> ModelSettings {
69 ModelSettings::default()
70 }
71}
72
73#[derive(Clone)]
74pub struct ScriptedModelProvider {
75 backend: String,
76 _default_model: String,
77 llm: ScriptedLlmClient,
78 context_length: Option<u64>,
79 max_output_tokens: Option<u64>,
80 default_settings: ModelSettings,
81}
82
83impl ScriptedModelProvider {
84 pub fn new(
85 backend: impl Into<String>,
86 default_model: impl Into<String>,
87 responses: Vec<LLMResponse>,
88 ) -> Self {
89 Self::from_steps(
90 backend,
91 default_model,
92 responses.into_iter().map(ScriptStep::response).collect(),
93 )
94 }
95
96 pub fn from_steps(
97 backend: impl Into<String>,
98 default_model: impl Into<String>,
99 steps: Vec<ScriptStep>,
100 ) -> Self {
101 Self {
102 backend: backend.into(),
103 _default_model: default_model.into(),
104 llm: ScriptedLlmClient::from_steps(steps),
105 context_length: Some(128_000),
106 max_output_tokens: Some(16_384),
107 default_settings: ModelSettings::default(),
108 }
109 }
110
111 pub fn from_callback(
112 backend: impl Into<String>,
113 default_model: impl Into<String>,
114 callback: impl Fn(&LlmRequest) -> Result<LLMResponse, crate::llm::LlmError>
115 + Send
116 + Sync
117 + 'static,
118 ) -> Self {
119 Self::from_steps(backend, default_model, vec![ScriptStep::callback(callback)])
120 }
121
122 pub fn with_default_settings(mut self, settings: ModelSettings) -> Self {
123 self.default_settings = settings;
124 self
125 }
126
127 pub fn with_token_limits(
128 mut self,
129 context_length: Option<u64>,
130 max_output_tokens: Option<u64>,
131 ) -> Self {
132 self.context_length = context_length;
133 self.max_output_tokens = max_output_tokens;
134 self
135 }
136}
137
138impl ModelProvider for ScriptedModelProvider {
139 fn resolve(&self, model: &ModelRef) -> Result<ResolvedModelConfig, ModelError> {
140 match model {
141 ModelRef::Named(model) => Ok(ResolvedModelConfig::new(
142 self.backend.clone(),
143 model.clone(),
144 model.clone(),
145 model.clone(),
146 Vec::new(),
147 )
148 .with_token_limits(self.context_length, self.max_output_tokens)),
149 ModelRef::BackendModel { backend, model } => {
150 if backend != &self.backend {
151 return Err(ModelError::BackendMismatch {
152 requested: backend.clone(),
153 provider: self.backend.clone(),
154 });
155 }
156 Ok(ResolvedModelConfig::new(
157 backend.clone(),
158 model.clone(),
159 model.clone(),
160 model.clone(),
161 Vec::new(),
162 )
163 .with_token_limits(self.context_length, self.max_output_tokens))
164 }
165 ModelRef::Resolved(resolved) => Ok(resolved.clone()),
166 }
167 }
168
169 fn client(&self, _resolved: &ResolvedModelConfig) -> Result<Arc<dyn LlmClient>, ModelError> {
170 Ok(Arc::new(self.llm.clone()))
171 }
172
173 fn default_settings(&self, _resolved: &ResolvedModelConfig) -> ModelSettings {
174 self.default_settings.clone()
175 }
176}
177
178impl Default for ScriptedModelProvider {
179 fn default() -> Self {
180 Self::new("scripted", "demo-model", Vec::new())
181 }
182}
183
184#[derive(Debug, Clone)]
185pub struct VvLlmModelProvider {
186 settings_file: PathBuf,
187 default_backend: Option<String>,
188 timeout_seconds: f64,
189}
190
191impl VvLlmModelProvider {
192 pub fn from_settings_file(path: impl Into<PathBuf>) -> Self {
193 Self {
194 settings_file: path.into(),
195 default_backend: None,
196 timeout_seconds: 90.0,
197 }
198 }
199
200 pub fn with_default_backend(mut self, backend: impl Into<String>) -> Self {
201 self.default_backend = Some(backend.into());
202 self
203 }
204
205 pub fn with_timeout_seconds(mut self, timeout_seconds: f64) -> Self {
206 self.timeout_seconds = timeout_seconds.max(1.0);
207 self
208 }
209
210 pub fn settings_file(&self) -> &Path {
211 &self.settings_file
212 }
213}
214
215impl ModelProvider for VvLlmModelProvider {
216 fn resolve(&self, model: &ModelRef) -> Result<ResolvedModelConfig, ModelError> {
217 match model {
218 ModelRef::Named(model) => {
219 let Some(backend) = self.default_backend.as_deref() else {
220 return Err(ModelError::MissingDefaultBackend(model.clone()));
221 };
222 let settings = crate::config::load_llm_settings_from_file(&self.settings_file)
223 .map_err(|error| ModelError::Config(error.to_string()))?;
224 resolve_model_endpoint(&settings, backend, model)
225 .map_err(|error| ModelError::Config(error.to_string()))
226 }
227 ModelRef::BackendModel { backend, model } => {
228 let settings = crate::config::load_llm_settings_from_file(&self.settings_file)
229 .map_err(|error| ModelError::Config(error.to_string()))?;
230 resolve_model_endpoint(&settings, backend, model)
231 .map_err(|error| ModelError::Config(error.to_string()))
232 }
233 ModelRef::Resolved(resolved) => Ok(resolved.clone()),
234 }
235 }
236
237 fn client(&self, resolved: &ResolvedModelConfig) -> Result<Arc<dyn LlmClient>, ModelError> {
238 let (llm, _) = build_vv_llm_from_local_settings(
239 &self.settings_file,
240 &resolved.backend,
241 &resolved.selected_model,
242 self.timeout_seconds,
243 )
244 .map_err(|error| ModelError::Config(error.to_string()))?;
245 let llm: VvLlmClient = llm;
246 Ok(Arc::new(llm) as Arc<dyn LlmClient>)
247 }
248}