1pub mod anthropic;
2pub mod openai;
3
4use std::{future::Future, pin::Pin};
5
6use serde::{Deserialize, Serialize};
7use tokio::sync::mpsc::UnboundedReceiver;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct Message {
11 pub role: Role,
12 pub content: Vec<ContentBlock>,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
16#[serde(rename_all = "lowercase")]
17pub enum Role {
18 User,
19 Assistant,
20 System,
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24#[serde(tag = "type", rename_all = "snake_case")]
25pub enum ContentBlock {
26 Text(String),
27 Image {
28 media_type: String,
29 data: String,
30 },
31 ToolUse {
32 id: String,
33 name: String,
34 input: serde_json::Value,
35 },
36 ToolResult {
37 tool_use_id: String,
38 content: String,
39 is_error: bool,
40 },
41 Thinking {
42 thinking: String,
43 signature: String,
44 },
45 Compaction {
46 content: String,
47 },
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct ToolDefinition {
52 pub name: String,
53 pub description: String,
54 pub input_schema: serde_json::Value,
55}
56
57#[derive(Debug, Clone)]
58pub struct StreamEvent {
59 pub event_type: StreamEventType,
60}
61
62#[derive(Debug, Clone)]
63pub enum StreamEventType {
64 TextDelta(String),
65 ThinkingDelta(String),
66 ThinkingComplete {
67 thinking: String,
68 signature: String,
69 },
70 ToolUseStart {
71 id: String,
72 name: String,
73 },
74 ToolUseInputDelta(String),
75 ToolUseEnd,
76 CompactionComplete(String),
77 MessageStart,
78 MessageEnd {
79 stop_reason: StopReason,
80 usage: Usage,
81 },
82 Error(String),
83}
84
85#[derive(Debug, Clone)]
86pub enum StopReason {
87 EndTurn,
88 MaxTokens,
89 ToolUse,
90 StopSequence,
91}
92
93#[derive(Debug, Clone, Default)]
94pub struct Usage {
95 pub input_tokens: u32,
96 pub output_tokens: u32,
97 pub cache_read_tokens: u32,
98 pub cache_write_tokens: u32,
99}
100
101pub trait Provider: Send + Sync {
102 fn name(&self) -> &str;
103 fn model(&self) -> &str;
104 fn set_model(&mut self, model: String);
105 fn available_models(&self) -> Vec<String>;
106 fn context_window(&self) -> u32;
107 fn supports_server_compaction(&self) -> bool {
108 false
109 }
110 fn fetch_context_window(
111 &self,
112 ) -> Pin<Box<dyn Future<Output = anyhow::Result<u32>> + Send + '_>>;
113 fn supports_vision(&self) -> bool {
114 true
115 }
116 fn fetch_models(
117 &self,
118 ) -> Pin<Box<dyn Future<Output = anyhow::Result<Vec<String>>> + Send + '_>>;
119 fn stream(
120 &self,
121 messages: &[Message],
122 system: Option<&str>,
123 tools: &[ToolDefinition],
124 max_tokens: u32,
125 thinking_budget: u32,
126 ) -> Pin<Box<dyn Future<Output = anyhow::Result<UnboundedReceiver<StreamEvent>>> + Send + '_>>;
127
128 fn stream_with_model(
129 &self,
130 model: &str,
131 messages: &[Message],
132 system: Option<&str>,
133 tools: &[ToolDefinition],
134 max_tokens: u32,
135 thinking_budget: u32,
136 ) -> Pin<Box<dyn Future<Output = anyhow::Result<UnboundedReceiver<StreamEvent>>> + Send + '_>>
137 {
138 let _ = model;
139 self.stream(messages, system, tools, max_tokens, thinking_budget)
140 }
141}