Skip to main content

dot/provider/
mod.rs

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}