Skip to main content

tiny_agent/providers/
mod.rs

1pub use super::shared::{Message, ToolDefinition};
2use async_trait::async_trait;
3use futures_util::stream::BoxStream;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::fmt;
7pub mod anthropic;
8pub mod openai_compatible;
9
10#[derive(Debug, Serialize, Deserialize, Clone)]
11pub struct LLMRequest {
12    pub model: String,
13    pub system: Option<String>,
14    pub messages: Vec<Message>,
15    pub tools: Vec<ToolDefinition>,
16    pub temperature: Option<f32>,
17    pub max_tokens: Option<u32>,
18}
19
20#[derive(Debug, Clone)]
21pub enum ProviderStreamChunk {
22    ContentDelta(String),
23    ThinkingDelta(String),
24    ToolCallDelta {
25        index: usize,
26        id: Option<String>,
27        name: Option<String>,
28        arguments_delta: String,
29    },
30    Done {
31        finish_reason: FinishReason,
32        usage: Option<TokenUsage>,
33    },
34}
35
36#[derive(Debug, Clone)]
37pub enum FinishReason {
38    Stop,
39    ToolCalls,
40    Length,
41    ContentFilter,
42    Error,
43    Unknown(String),
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct TokenUsage {
48    pub input_tokens: Option<u32>,
49    pub output_tokens: Option<u32>,
50    pub total_tokens: Option<u32>,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub enum ProviderErrorKind {
55    Transient,
56    Permanent,
57    RateLimited { retry_after_secs: Option<u64> },
58}
59
60#[derive(Debug, Clone)]
61pub struct ProviderError {
62    pub kind: ProviderErrorKind,
63    pub message: String,
64}
65
66impl ProviderError {
67    pub fn transient(message: impl Into<String>) -> Self {
68        Self {
69            kind: ProviderErrorKind::Transient,
70            message: message.into(),
71        }
72    }
73
74    pub fn permanent(message: impl Into<String>) -> Self {
75        Self {
76            kind: ProviderErrorKind::Permanent,
77            message: message.into(),
78        }
79    }
80
81    pub fn rate_limited(message: impl Into<String>, retry_after_secs: Option<u64>) -> Self {
82        Self {
83            kind: ProviderErrorKind::RateLimited { retry_after_secs },
84            message: message.into(),
85        }
86    }
87}
88
89impl fmt::Display for ProviderError {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        match &self.kind {
92            ProviderErrorKind::Transient | ProviderErrorKind::Permanent => {
93                f.write_str(&self.message)
94            }
95            ProviderErrorKind::RateLimited { retry_after_secs } => {
96                if let Some(seconds) = retry_after_secs {
97                    write!(f, "{} (retry_after={}s)", self.message, seconds)
98                } else {
99                    f.write_str(&self.message)
100                }
101            }
102        }
103    }
104}
105
106impl std::error::Error for ProviderError {}
107
108#[async_trait]
109pub trait LlmProvider: Send + Sync {
110    fn name(&self) -> &str;
111    fn request_snapshot(&self, req: &LLMRequest) -> Value {
112        serde_json::to_value(req).unwrap_or(Value::Null)
113    }
114    async fn chat_completion(
115        &self,
116        req: LLMRequest,
117    ) -> Result<BoxStream<'static, Result<ProviderStreamChunk, ProviderError>>, ProviderError>;
118}