1use async_trait::async_trait;
2use futures::stream::Stream;
3use serde::{Deserialize, Serialize};
4use std::pin::Pin;
5
6use crate::event::{StopReason, TokenUsage};
7
8pub mod anthropic;
9pub mod detect;
10pub mod discovery;
11pub mod ollama;
12pub mod openai_compat;
13pub mod responses;
14pub mod sse_buffer;
15pub mod tool_markup;
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct ModelCaps {
21 pub context_window: u64,
23 pub max_output: u64,
25 pub tools: bool,
27 pub vision: bool,
29 pub cost_input_per_mtok: f64,
31 pub cost_output_per_mtok: f64,
33 pub latency: LatencyClass,
35}
36
37impl Default for ModelCaps {
38 fn default() -> Self {
39 Self {
40 context_window: 128_000,
41 max_output: 16_000,
42 tools: true,
43 vision: false,
44 cost_input_per_mtok: 0.0,
45 cost_output_per_mtok: 0.0,
46 latency: LatencyClass::Medium,
47 }
48 }
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
52pub enum LatencyClass {
53 Fast,
54 Medium,
55 Slow,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct Msg {
62 pub role: String,
63 pub content: Vec<ContentBlock>,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
67#[serde(tag = "type")]
68pub enum ContentBlock {
69 #[serde(rename = "text")]
70 Text { text: String },
71 #[serde(rename = "image")]
72 Image { source: ImageSource },
73 #[serde(rename = "tool_use")]
74 ToolUse {
75 id: String,
76 name: String,
77 input: serde_json::Value,
78 },
79 #[serde(rename = "tool_result")]
80 ToolResult {
81 tool_use_id: String,
82 content: Vec<ContentBlock>,
83 is_error: Option<bool>,
84 },
85 #[serde(rename = "reasoning")]
91 Reasoning { text: String },
92}
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
95#[serde(tag = "type")]
96pub enum ImageSource {
97 #[serde(rename = "base64")]
98 Base64 { media_type: String, data: String },
99 #[serde(rename = "url")]
100 Url { url: String },
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct ToolSpec {
107 pub name: String,
108 pub description: String,
109 pub input_schema: serde_json::Value,
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
115pub enum PromptCacheTtl {
116 FiveMinutes,
117 OneHour,
118}
119
120impl PromptCacheTtl {
121 pub fn anthropic_ttl(&self) -> &'static str {
122 match self {
123 Self::FiveMinutes => "5m",
124 Self::OneHour => "1h",
125 }
126 }
127
128 pub fn openai_retention(&self) -> &'static str {
129 "in_memory"
132 }
133}
134
135#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
136pub struct PromptCacheConfig {
137 pub enabled: bool,
138 pub ttl: PromptCacheTtl,
139 pub key: Option<String>,
140}
141
142impl PromptCacheConfig {
143 pub fn enabled(key: Option<String>) -> Self {
144 Self {
145 enabled: true,
146 ttl: PromptCacheTtl::OneHour,
147 key: key.into(),
148 }
149 }
150
151 pub fn disabled() -> Self {
152 Self {
153 enabled: false,
154 ttl: PromptCacheTtl::FiveMinutes,
155 key: None,
156 }
157 }
158}
159
160impl Default for PromptCacheConfig {
161 fn default() -> Self {
162 Self::enabled(None)
163 }
164}
165
166#[derive(Debug, Clone)]
169pub struct BrainRequest {
170 pub system: Option<String>,
171 pub messages: Vec<Msg>,
172 pub tools: Vec<ToolSpec>,
173 pub max_tokens: u32,
174 pub temperature: f32,
175 pub stop: Vec<String>,
176 pub cache: PromptCacheConfig,
177}
178
179impl Default for BrainRequest {
180 fn default() -> Self {
181 Self {
182 system: None,
183 messages: vec![],
184 tools: vec![],
185 max_tokens: 4096,
186 temperature: 0.0,
187 stop: vec![],
188 cache: PromptCacheConfig::default(),
189 }
190 }
191}
192
193#[derive(Debug, Clone)]
196pub enum BrainEvent {
197 TextDelta(String),
198 ReasoningDelta(String),
202 ToolUseStart {
203 id: String,
204 name: String,
205 },
206 ToolUseDelta {
207 id: String,
208 json: String,
209 },
210 ToolUseEnd {
211 id: String,
212 },
213 Usage(TokenUsage),
214 Done(StopReason),
215 Error(String),
216}
217
218pub type BrainStream = Pin<Box<dyn Stream<Item = BrainEvent> + Send>>;
221
222#[async_trait]
227pub trait Brain: Send + Sync {
228 fn id(&self) -> &str;
230 fn caps(&self) -> ModelCaps;
232 async fn complete(&self, req: BrainRequest) -> anyhow::Result<BrainStream>;
234}
235
236#[derive(Debug, Clone)]
239pub enum BrainError {
240 RateLimit { retry_after: Option<u64> },
241 ServerError { status: u16, body: String },
242 Timeout,
243 Refusal(String),
244 Unknown(String),
245}
246
247impl std::fmt::Display for BrainError {
248 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249 match self {
250 BrainError::RateLimit { retry_after } => {
251 write!(f, "rate limited (retry after {:?}s)", retry_after)
252 }
253 BrainError::ServerError { status, body } => {
254 write!(f, "server error {}: {}", status, body)
255 }
256 BrainError::Timeout => write!(f, "timeout"),
257 BrainError::Refusal(msg) => write!(f, "refusal: {}", msg),
258 BrainError::Unknown(msg) => write!(f, "unknown: {}", msg),
259 }
260 }
261}