spider_browser/ai/
llm_provider.rs1use crate::errors::{Result, SpiderError};
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "lowercase")]
11pub enum LLMProviderKind {
12 OpenAI,
13 Anthropic,
14 OpenRouter,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct LLMConfig {
20 pub provider: LLMProviderKind,
22 pub model: String,
24 pub api_key: String,
26 #[serde(default)]
28 pub base_url: Option<String>,
29 #[serde(default)]
31 pub max_tokens: Option<u32>,
32 #[serde(default)]
34 pub temperature: Option<f64>,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum LLMRole {
41 System,
42 User,
43 Assistant,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
48#[serde(tag = "type")]
49pub enum LLMContentPart {
50 #[serde(rename = "text")]
51 Text { text: String },
52 #[serde(rename = "image_url")]
53 ImageUrl { image_url: ImageUrlValue },
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct ImageUrlValue {
59 pub url: String,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(untagged)]
65pub enum LLMContent {
66 Text(String),
67 Parts(Vec<LLMContentPart>),
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct LLMMessage {
73 pub role: LLMRole,
74 pub content: LLMContent,
75}
76
77impl LLMMessage {
78 pub fn system(text: impl Into<String>) -> Self {
80 Self {
81 role: LLMRole::System,
82 content: LLMContent::Text(text.into()),
83 }
84 }
85
86 pub fn user(text: impl Into<String>) -> Self {
88 Self {
89 role: LLMRole::User,
90 content: LLMContent::Text(text.into()),
91 }
92 }
93
94 pub fn user_parts(parts: Vec<LLMContentPart>) -> Self {
96 Self {
97 role: LLMRole::User,
98 content: LLMContent::Parts(parts),
99 }
100 }
101
102 pub fn assistant(text: impl Into<String>) -> Self {
104 Self {
105 role: LLMRole::Assistant,
106 content: LLMContent::Text(text.into()),
107 }
108 }
109}
110
111#[derive(Debug, Clone, Default)]
113pub struct ChatOptions {
114 pub json_mode: bool,
116}
117
118#[async_trait::async_trait]
122pub trait LLMProvider: Send + Sync {
123 async fn chat(&self, messages: &[LLMMessage], options: Option<ChatOptions>) -> Result<String>;
125}
126
127pub async fn chat_json<T: serde::de::DeserializeOwned>(
132 llm: &dyn LLMProvider,
133 messages: &[LLMMessage],
134) -> Result<T> {
135 let text = llm
136 .chat(messages, Some(ChatOptions { json_mode: true }))
137 .await?;
138 parse_json_response(&text)
139}
140
141pub fn parse_json_response<T: serde::de::DeserializeOwned>(text: &str) -> Result<T> {
143 if let Ok(val) = serde_json::from_str::<T>(text) {
145 return Ok(val);
146 }
147 if let Some(start) = text.find("```") {
149 let after_fence = &text[start + 3..];
150 let json_start = after_fence
152 .find('\n')
153 .map(|i| i + 1)
154 .unwrap_or(0);
155 if let Some(end) = after_fence[json_start..].find("```") {
156 let json_str = &after_fence[json_start..json_start + end];
157 if let Ok(val) = serde_json::from_str::<T>(json_str.trim()) {
158 return Ok(val);
159 }
160 }
161 }
162 Err(SpiderError::Llm(format!(
163 "LLM response is not valid JSON: {}",
164 &text[..text.len().min(200)]
165 )))
166}
167
168pub fn create_provider(config: LLMConfig) -> Box<dyn LLMProvider> {
170 match config.provider {
171 LLMProviderKind::OpenAI | LLMProviderKind::OpenRouter => {
172 Box::new(crate::ai::providers::openai::OpenAICompatibleProvider::new(config))
173 }
174 LLMProviderKind::Anthropic => {
175 Box::new(crate::ai::providers::anthropic::AnthropicProvider::new(config))
176 }
177 }
178}