Skip to main content

robit_ai/
client.rs

1//! LlmClient: a thin wrapper around async-openai with unified config support.
2
3use 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    /// Create a new LlmClient from loaded configuration.
20    ///
21    /// `profile_name`: which profile to use. If `None`, uses the default profile.
22    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    /// Streaming chat completion. Returns an async stream of response chunks.
42    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    /// Non-streaming chat completion. Returns the full response.
62    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    /// Get the current model ID (e.g. "deepseek-chat").
81    pub fn model(&self) -> &str {
82        &self.model
83    }
84
85    /// Get the profile name (e.g. "default").
86    pub fn profile(&self) -> &str {
87        &self.resolved.profile_name
88    }
89
90    /// Get the resolved model info.
91    pub fn resolved(&self) -> &ResolvedModel {
92        &self.resolved
93    }
94
95    /// Whether the current model supports image inputs.
96    pub fn supports_images(&self) -> bool {
97        self.resolved.supports_images
98    }
99
100    /// Whether the current model supports tool calling.
101    pub fn supports_tools(&self) -> bool {
102        self.resolved.supports_tools
103    }
104}