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
12/// Validate that all messages are valid before sending to LLM.
13/// Returns a filtered list of messages with invalid messages removed.
14fn validate_and_filter_messages(mut messages: Vec<ChatCompletionRequestMessage>) -> Vec<ChatCompletionRequestMessage> {
15    let original_len = messages.len();
16    messages.retain(|msg| {
17        match msg {
18            ChatCompletionRequestMessage::Assistant(assistant_msg) => {
19                // Assistant message must have either content or tool_calls
20                let has_content = assistant_msg.content.is_some();
21                let has_tool_calls = assistant_msg.tool_calls.is_some();
22                if !has_content && !has_tool_calls {
23                    tracing::warn!("Filtering out invalid assistant message (has neither content nor tool_calls)");
24                    false
25                } else {
26                    true
27                }
28            }
29            _ => true
30        }
31    });
32    let filtered_len = messages.len();
33    if filtered_len < original_len {
34        tracing::info!("Filtered {} invalid messages from history", original_len - filtered_len);
35    }
36    messages
37}
38
39pub struct LlmClient {
40    client: async_openai::Client<OpenAIConfig>,
41    model: String,
42    resolved: ResolvedModel,
43}
44
45impl LlmClient {
46    /// Create a new LlmClient from loaded configuration.
47    ///
48    /// `profile_name`: which profile to use. If `None`, uses the default profile.
49    pub fn from_config(
50        config: &RobitConfig,
51        profile_name: Option<&str>,
52    ) -> Result<Self, LlmError> {
53        let resolved = resolve_profile(config, profile_name)?;
54
55        let oc = OpenAIConfig::new()
56            .with_api_base(&resolved.base_url)
57            .with_api_key(&resolved.api_key);
58
59        let client = async_openai::Client::with_config(oc);
60
61        Ok(Self {
62            client,
63            model: resolved.model_id.clone(),
64            resolved,
65        })
66    }
67
68    /// Streaming chat completion. Returns an async stream of response chunks.
69    pub async fn chat_stream(
70        &self,
71        messages: Vec<ChatCompletionRequestMessage>,
72        tools: Option<Vec<ChatCompletionTools>>,
73    ) -> Result<ChatCompletionResponseStream, LlmError> {
74        // Validate and filter messages before sending to LLM
75        let messages = validate_and_filter_messages(messages);
76
77        let request = CreateChatCompletionRequest {
78            model: self.model.clone(),
79            messages,
80            tools,
81            stream: Some(true),
82            max_completion_tokens: self.resolved.max_tokens,
83            temperature: self.resolved.temperature,
84            ..Default::default()
85        };
86
87        let stream = self.client.chat().create_stream(request).await?;
88        Ok(stream)
89    }
90
91    /// Non-streaming chat completion. Returns the full response.
92    pub async fn chat(
93        &self,
94        messages: Vec<ChatCompletionRequestMessage>,
95        tools: Option<Vec<ChatCompletionTools>>,
96    ) -> Result<CreateChatCompletionResponse, LlmError> {
97        // Validate and filter messages before sending to LLM
98        let messages = validate_and_filter_messages(messages);
99
100        let request = CreateChatCompletionRequest {
101            model: self.model.clone(),
102            messages,
103            tools,
104            max_completion_tokens: self.resolved.max_tokens,
105            temperature: self.resolved.temperature,
106            ..Default::default()
107        };
108
109        let response = self.client.chat().create(request).await?;
110        Ok(response)
111    }
112
113    /// Get the current model ID (e.g. "deepseek-chat").
114    pub fn model(&self) -> &str {
115        &self.model
116    }
117
118    /// Get the profile name (e.g. "default").
119    pub fn profile(&self) -> &str {
120        &self.resolved.profile_name
121    }
122
123    /// Get the resolved model info.
124    pub fn resolved(&self) -> &ResolvedModel {
125        &self.resolved
126    }
127
128    /// Whether the current model supports image inputs.
129    pub fn supports_images(&self) -> bool {
130        self.resolved.supports_images
131    }
132
133    /// Whether the current model supports tool calling.
134    pub fn supports_tools(&self) -> bool {
135        self.resolved.supports_tools
136    }
137}
robit_ai/client.rs

robit_ai/
client.rs
