doum_cli/llm/
client.rs

1use crate::llm::{AnthropicClient, AnthropicConfig, AnthropicSecret, OpenAIClient, OpenAIConfig};
2use crate::system::{DoumError, DoumResult};
3use crate::{llm::OpenAISecret, system::LLMConfig};
4use serde::{Deserialize, Serialize};
5
6/// LLM Message Role
7#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum Role {
10    User,
11    Assistant,
12}
13
14/// LLM Request
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct LLMRequest {
17    pub system: String,
18    pub messages: Vec<Message>,
19    pub use_websearch: bool,
20}
21
22/// LLM Message
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct Message {
25    pub role: Role,
26    pub content: String,
27}
28
29impl Message {
30    /// create user message
31    pub fn user(content: impl Into<String>) -> Self {
32        Self {
33            role: Role::User,
34            content: content.into(),
35        }
36    }
37
38    /// create assistant message
39    pub fn assistant(content: impl Into<String>) -> Self {
40        Self {
41            role: Role::Assistant,
42            content: content.into(),
43        }
44    }
45}
46
47/// LLM Client Trait
48#[async_trait::async_trait]
49pub trait LLMClient: Send + Sync {
50    /// Generate response from LLM
51    async fn generate(&self, request: LLMRequest) -> DoumResult<String>;
52
53    /// Verify LLM client connectivity
54    async fn verify(&self) -> DoumResult<bool> {
55        let test_request = LLMRequest {
56            system: "This is a test, please respond shortly.".to_string(),
57            messages: vec![Message::user("Hello")],
58            use_websearch: false,
59        };
60
61        match self.generate(test_request).await {
62            Ok(_) => Ok(true),
63            Err(_) => Ok(false),
64        }
65    }
66}
67
68/// Create LLM client based on configuration
69pub fn create_client(config: &LLMConfig) -> DoumResult<Box<dyn LLMClient>> {
70    let provider = &config.provider;
71
72    match provider.as_str() {
73        "openai" => {
74            let secret = OpenAISecret::load().map_err(|e| DoumError::Config(e.to_string()))?;
75
76            let openai_config = OpenAIConfig {
77                model: config.model.clone(),
78                api_key: secret.api_key,
79                organization: secret.organization,
80                project: secret.project,
81            };
82            let client = OpenAIClient::new(openai_config, config.timeout)?;
83            Ok(Box::new(client))
84        }
85        "anthropic" => {
86            let secret = AnthropicSecret::load().map_err(|e| DoumError::Config(e.to_string()))?;
87
88            let anthropic_config = AnthropicConfig {
89                model: config.model.clone(),
90                api_key: secret.api_key,
91            };
92            let client = AnthropicClient::new(anthropic_config, config.timeout)?;
93            Ok(Box::new(client))
94        }
95        _ => Err(crate::system::DoumError::Config(format!(
96            "Unknown provider: {}",
97            provider
98        ))),
99    }
100}