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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum Role {
10 User,
11 Assistant,
12}
13
14#[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#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct Message {
25 pub role: Role,
26 pub content: String,
27}
28
29impl Message {
30 pub fn user(content: impl Into<String>) -> Self {
32 Self {
33 role: Role::User,
34 content: content.into(),
35 }
36 }
37
38 pub fn assistant(content: impl Into<String>) -> Self {
40 Self {
41 role: Role::Assistant,
42 content: content.into(),
43 }
44 }
45}
46
47#[async_trait::async_trait]
49pub trait LLMClient: Send + Sync {
50 async fn generate(&self, request: LLMRequest) -> DoumResult<String>;
52
53 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
68pub 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}