1use crate::client::DeepinfraClient;
2use bon::Builder;
3use serde::{Deserialize, Serialize};
4use serde_json;
5use tracing::instrument;
6
7const CHAT_COMPLETIONS_API_URL: &str = "https://api.deepinfra.com/v1/openai/chat/completions";
8
9#[derive(Debug, Deserialize, Serialize, Builder)]
10pub struct SystemMessage {
11 #[builder(into)]
12 content: String,
13 #[builder(into)]
14 name: Option<String>,
15}
16
17#[derive(Debug, Deserialize, Serialize, Builder)]
18pub struct UserMessage {
19 #[builder(into)]
20 content: String,
21 #[builder(into)]
22 name: Option<String>,
23}
24
25#[derive(Debug, Deserialize, Serialize, Builder)]
26pub struct AssistantMessage {
27 #[builder(into)]
28 pub content: String,
29 name: Option<String>,
30 tool_calls: Option<Vec<ToolCall>>,
31}
32
33#[derive(Debug, Deserialize, Serialize, Builder)]
34pub struct ToolMessage {
35 #[builder(into)]
36 content: String,
37 tool_call_id: String,
38}
39
40#[derive(Debug, Deserialize, Serialize)]
41#[serde(tag = "role", rename_all = "snake_case")]
42pub enum Message {
43 System(SystemMessage),
44 User(UserMessage),
45 Assistant(AssistantMessage),
46 Tool(ToolMessage),
47}
48
49#[derive(Debug, Serialize, Deserialize, Builder)]
52pub struct ChatCompletionRequest {
53 #[builder(default = 0.0)]
57 frequency_penalty: f64,
58
59 #[builder(default = 100000)]
62 max_tokens: u32,
63
64 #[builder(into)]
67 messages: Vec<Message>,
68
69 #[builder(default = 0.0)]
72 min_p: f64,
73
74 #[builder(default = "deepseek-ai/DeepSeek-V3".to_string(), into)]
77 model: String,
78
79 #[builder(default = 1)]
82 n: u32,
83
84 #[builder(default = 0.0)]
88 presence_penalty: f64,
89
90 #[builder(default = 1.0)]
93 repetition_penalty: f64,
94
95 #[builder(into)]
97 response_format: Option<ResponseFormat>,
98
99 #[builder(into)]
102 seed: Option<u64>,
103
104 #[builder(into)]
106 stop: Option<Vec<String>>,
107
108 #[builder(default = false)]
110 stream: bool,
111
112 #[builder(default = 1.0)]
115 temperature: f64,
116
117 #[builder(into)]
121 tool_choice: Option<String>,
122
123 #[builder(into)]
125 tools: Option<Vec<ChatTool>>,
126
127 #[builder(default = 0)]
129 top_k: u32,
130
131 #[builder(default = 1.0)]
134 top_p: f64,
135
136 #[builder(into)]
139 user: Option<String>,
140}
141
142#[derive(Debug, Serialize, Deserialize)]
145pub struct ChatTool {
146 #[serde(default = "default_tool_type", rename = "type")]
148 type_: String,
149
150 function: FunctionDefinition,
152}
153
154fn default_tool_type() -> String {
155 "function".to_string()
156}
157
158#[derive(Debug, Serialize, Deserialize)]
160pub struct FunctionDefinition {
161 name: String,
163
164 description: String,
166
167 parameters: serde_json::Value,
169}
170
171#[derive(Debug, Serialize, Deserialize)]
172#[serde(rename_all = "snake_case")]
173pub enum ResponseFormatType {
174 Text,
175 JsonObject,
176}
177
178#[derive(Debug, Serialize, Deserialize)]
180pub struct ResponseFormat {
181 #[serde(default = "default_response_format_type", rename = "type")]
183 pub response_type: ResponseFormatType,
184}
185
186fn default_response_format_type() -> ResponseFormatType {
187 ResponseFormatType::Text
188}
189
190#[derive(Debug, Serialize, Deserialize)]
192pub struct ToolCall {
193 id: String,
195
196 #[serde(rename = "type")]
198 type_: String,
199
200 function: FunctionCall,
202}
203
204#[derive(Debug, Serialize, Deserialize)]
206pub struct FunctionCall {
207 name: String,
209
210 arguments: String,
213}
214
215#[derive(Debug, Serialize, Deserialize)]
216pub struct Choice {
217 index: i32,
218 pub message: Message,
219 finish_reason: String,
220}
221
222#[derive(Debug, Serialize, Deserialize)]
223pub struct Usage {
224 prompt_tokens: i32,
225 total_tokens: i32,
226 completion_tokens: i32,
227}
228
229#[derive(Debug, Serialize, Deserialize)]
230pub struct ChatCompletionResponse {
231 id: Option<String>,
232 object: Option<String>,
233 created: Option<i64>,
234 model: Option<String>,
235 pub choices: Vec<Choice>,
236 usage: Option<Usage>,
237}
238
239#[derive(Debug, thiserror::Error)]
240pub enum ChatCompletionError {
241 #[error("Request errored {0}")]
242 ReqwestError(#[from] reqwest::Error),
243}
244
245type Result<T> = std::result::Result<T, ChatCompletionError>;
246
247impl DeepinfraClient {
248 #[instrument(skip(self))]
260 pub async fn chat_completition(
261 &self,
262 body: ChatCompletionRequest,
263 ) -> Result<ChatCompletionResponse> {
264 let response = self
265 .client
266 .post(CHAT_COMPLETIONS_API_URL)
267 .json(&body)
268 .send()
269 .await?
270 .json()
271 .await?;
272
273 Ok(response)
274 }
275}