1use async_openai::config::OpenAIConfig;
4use async_openai::types::chat::{
5 ChatCompletionRequestMessage, ChatCompletionResponseStream, ChatCompletionTools,
6 CreateChatCompletionRequest, CreateChatCompletionResponse,
7};
8
9use crate::config::{resolve_profile, ResolvedModel, RobitConfig};
10use crate::error::LlmError;
11
12pub struct LlmClient {
13 client: async_openai::Client<OpenAIConfig>,
14 model: String,
15 resolved: ResolvedModel,
16}
17
18impl LlmClient {
19 pub fn from_config(
23 config: &RobitConfig,
24 profile_name: Option<&str>,
25 ) -> Result<Self, LlmError> {
26 let resolved = resolve_profile(config, profile_name)?;
27
28 let oc = OpenAIConfig::new()
29 .with_api_base(&resolved.base_url)
30 .with_api_key(&resolved.api_key);
31
32 let client = async_openai::Client::with_config(oc);
33
34 Ok(Self {
35 client,
36 model: resolved.model_id.clone(),
37 resolved,
38 })
39 }
40
41 pub async fn chat_stream(
43 &self,
44 messages: Vec<ChatCompletionRequestMessage>,
45 tools: Option<Vec<ChatCompletionTools>>,
46 ) -> Result<ChatCompletionResponseStream, LlmError> {
47 let request = CreateChatCompletionRequest {
48 model: self.model.clone(),
49 messages,
50 tools,
51 stream: Some(true),
52 max_completion_tokens: self.resolved.max_tokens,
53 temperature: self.resolved.temperature,
54 ..Default::default()
55 };
56
57 let stream = self.client.chat().create_stream(request).await?;
58 Ok(stream)
59 }
60
61 pub async fn chat(
63 &self,
64 messages: Vec<ChatCompletionRequestMessage>,
65 tools: Option<Vec<ChatCompletionTools>>,
66 ) -> Result<CreateChatCompletionResponse, LlmError> {
67 let request = CreateChatCompletionRequest {
68 model: self.model.clone(),
69 messages,
70 tools,
71 max_completion_tokens: self.resolved.max_tokens,
72 temperature: self.resolved.temperature,
73 ..Default::default()
74 };
75
76 let response = self.client.chat().create(request).await?;
77 Ok(response)
78 }
79
80 pub fn model(&self) -> &str {
82 &self.model
83 }
84
85 pub fn profile(&self) -> &str {
87 &self.resolved.profile_name
88 }
89
90 pub fn resolved(&self) -> &ResolvedModel {
92 &self.resolved
93 }
94
95 pub fn supports_images(&self) -> bool {
97 self.resolved.supports_images
98 }
99
100 pub fn supports_tools(&self) -> bool {
102 self.resolved.supports_tools
103 }
104}