Skip to main content

agent_base/llm/
mod.rs

1use async_trait::async_trait;
2use futures_core::Stream;
3use serde_json::Value;
4use std::pin::Pin;
5
6use crate::types::{AgentResult, ChatMessage, ResponseFormat};
7
8mod anthropic;
9mod openai;
10mod registry;
11
12pub use anthropic::AnthropicClient;
13pub use openai::OpenAiClient;
14pub use registry::{LlmClientBuilder, LlmProvider};
15
16#[derive(Clone, Debug)]
17pub enum StreamChunk {
18    Text(String),
19    Thought(String),
20    ToolCall(Value),
21    Usage(UsageInfo),
22    Stop,
23}
24
25#[derive(Clone, Debug, Default)]
26pub struct UsageInfo {
27    pub prompt_tokens: Option<u32>,
28    pub completion_tokens: Option<u32>,
29    pub total_tokens: Option<u32>,
30}
31
32#[derive(Clone, Debug, Default)]
33pub struct LlmCapabilities {
34    pub supports_streaming: bool,
35    pub supports_tools: bool,
36    pub supports_vision: bool,
37    pub supports_thinking: bool,
38    pub max_context_tokens: Option<u32>,
39    pub max_output_tokens: Option<u32>,
40}
41
42#[async_trait]
43pub trait LlmClient: Send + Sync {
44    async fn chat(
45        &self,
46        messages: &[ChatMessage],
47        tools: &[Value],
48        enable_thinking: Option<bool>,
49        response_format: Option<&ResponseFormat>,
50    ) -> AgentResult<Value>;
51
52    async fn chat_stream(
53        &self,
54        messages: &[ChatMessage],
55        tools: &[Value],
56        enable_thinking: Option<bool>,
57        response_format: Option<&ResponseFormat>,
58    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>>;
59
60    fn capabilities(&self) -> LlmCapabilities;
61}