Skip to main content

spider_browser/ai/
llm_provider.rs

1//! LLM provider abstraction for AI methods.
2//!
3//! Ported from TypeScript `ai/llm-provider.ts`.
4
5use crate::errors::{Result, SpiderError};
6use serde::{Deserialize, Serialize};
7
8/// Supported LLM provider backends.
9#[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/// LLM configuration for AI methods.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct LLMConfig {
20    /// Provider: OpenAI, Anthropic, or OpenRouter.
21    pub provider: LLMProviderKind,
22    /// Model name (e.g. "gpt-4o", "claude-sonnet-4-5-20250929").
23    pub model: String,
24    /// API key for the provider.
25    pub api_key: String,
26    /// Base URL override (e.g. for OpenRouter or local vLLM).
27    #[serde(default)]
28    pub base_url: Option<String>,
29    /// Max tokens (default: 4096).
30    #[serde(default)]
31    pub max_tokens: Option<u32>,
32    /// Temperature (default: 0.1).
33    #[serde(default)]
34    pub temperature: Option<f64>,
35}
36
37/// Role in a conversation message.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum LLMRole {
41    System,
42    User,
43    Assistant,
44}
45
46/// A content part within a message (text or image).
47#[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/// URL value for an image content part.
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct ImageUrlValue {
59    pub url: String,
60}
61
62/// Message content: either a plain string or structured parts.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(untagged)]
65pub enum LLMContent {
66    Text(String),
67    Parts(Vec<LLMContentPart>),
68}
69
70/// Message format for LLM calls.
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct LLMMessage {
73    pub role: LLMRole,
74    pub content: LLMContent,
75}
76
77impl LLMMessage {
78    /// Create a system message with plain text content.
79    pub fn system(text: impl Into<String>) -> Self {
80        Self {
81            role: LLMRole::System,
82            content: LLMContent::Text(text.into()),
83        }
84    }
85
86    /// Create a user message with plain text content.
87    pub fn user(text: impl Into<String>) -> Self {
88        Self {
89            role: LLMRole::User,
90            content: LLMContent::Text(text.into()),
91        }
92    }
93
94    /// Create a user message with structured content parts.
95    pub fn user_parts(parts: Vec<LLMContentPart>) -> Self {
96        Self {
97            role: LLMRole::User,
98            content: LLMContent::Parts(parts),
99        }
100    }
101
102    /// Create an assistant message with plain text content.
103    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/// Options for chat calls.
112#[derive(Debug, Clone, Default)]
113pub struct ChatOptions {
114    /// When true, request JSON-formatted output from the model.
115    pub json_mode: bool,
116}
117
118/// Pluggable LLM provider interface.
119///
120/// All methods are async and return `Result<_>` via [`SpiderError::Llm`].
121#[async_trait::async_trait]
122pub trait LLMProvider: Send + Sync {
123    /// Call the LLM with messages and get a text response.
124    async fn chat(&self, messages: &[LLMMessage], options: Option<ChatOptions>) -> Result<String>;
125}
126
127/// Call the LLM with messages and get a parsed JSON response.
128///
129/// This is a standalone function (not a trait method) so that `LLMProvider`
130/// remains dyn-compatible (no generic methods on the trait).
131pub 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
141/// Parse a JSON response, handling markdown code fences.
142pub fn parse_json_response<T: serde::de::DeserializeOwned>(text: &str) -> Result<T> {
143    // Try direct parse first
144    if let Ok(val) = serde_json::from_str::<T>(text) {
145        return Ok(val);
146    }
147    // Try extracting JSON from markdown code blocks
148    if let Some(start) = text.find("```") {
149        let after_fence = &text[start + 3..];
150        // Skip optional language tag (e.g. "json")
151        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
168/// Create an LLM provider from config.
169pub 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}