Skip to main content

deepinfra_client_rs/
chat_completition.rs

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/// Represents a request for generating chat completions.
50/// Includes all parameters as per the OpenAPI schema.
51#[derive(Debug, Serialize, Deserialize, Builder)]
52pub struct ChatCompletionRequest {
53    /// Penalizes new tokens based on their frequency in the text so far.
54    /// Increases the model's likelihood to talk about new topics.
55    /// Range: -2 to 2
56    #[builder(default = 0.0)]
57    frequency_penalty: f64,
58
59    /// Maximum number of tokens to generate in the chat completion.
60    /// Total length is limited by the model's context length.
61    #[builder(default = 100000)]
62    max_tokens: u32,
63
64    /// Conversation messages including user, assistant, and system messages.
65    /// Must include one system message anywhere.
66    #[builder(into)]
67    messages: Vec<Message>,
68
69    /// Minimum probability for a token to be considered, relative to the most likely token.
70    /// Must be between 0 and 1. Set to 0 to disable.
71    #[builder(default = 0.0)]
72    min_p: f64,
73
74    /// Model name to use for the chat completion.
75    /// Example: "meta-llama/Llama-2-70b-chat-hf"
76    #[builder(default = "deepseek-ai/DeepSeek-V3".to_string(), into)]
77    model: String,
78
79    /// Number of sequences to return.
80    /// Minimum: 1, Maximum: 4
81    #[builder(default = 1)]
82    n: u32,
83
84    /// Penalizes new tokens based on whether they appear in the text so far.
85    /// Increases the model's likelihood to talk about new topics.
86    /// Range: -2 to 2
87    #[builder(default = 0.0)]
88    presence_penalty: f64,
89
90    /// Penalty for repetition. Values >1 penalize, <1 encourage repetition.
91    /// Range: 0.01 to 5
92    #[builder(default = 1.0)]
93    repetition_penalty: f64,
94
95    /// The format of the response. Currently, only "text" or "json_object" are supported.
96    #[builder(into)]
97    response_format: Option<ResponseFormat>,
98
99    /// Seed for the random number generator.
100    /// If not provided, a random seed is used. Determinism is not guaranteed.
101    #[builder(into)]
102    seed: Option<u64>,
103
104    /// Up to 16 sequences where the API will stop generating further tokens.
105    #[builder(into)]
106    stop: Option<Vec<String>>,
107
108    /// Whether to stream the output via SSE or return the full response.
109    #[builder(default = false)]
110    stream: bool,
111
112    /// Sampling temperature to use, between 0 and 2.
113    /// Higher values make the output more random.
114    #[builder(default = 1.0)]
115    temperature: f64,
116
117    /// Controls which (if any) function is called by the model.
118    /// "none" means the model will not call a function.
119    /// "auto" means the model can choose to call a function or not.
120    #[builder(into)]
121    tool_choice: Option<String>,
122
123    /// A list of tools the model may call. Currently, only functions are supported.
124    #[builder(into)]
125    tools: Option<Vec<ChatTool>>,
126
127    /// Sample from the top_k number of tokens. 0 means off.
128    #[builder(default = 0)]
129    top_k: u32,
130
131    /// Nucleus sampling parameter between 0 and 1.
132    /// The model considers tokens with top_p probability mass.
133    #[builder(default = 1.0)]
134    top_p: f64,
135
136    /// A unique identifier representing your end-user.
137    /// Helps monitor and detect abuse. Avoid sending identifying information.
138    #[builder(into)]
139    user: Option<String>,
140}
141
142/// Represents a tool that the model may call during chat completion.
143/// Currently supports functions as tools.
144#[derive(Debug, Serialize, Deserialize)]
145pub struct ChatTool {
146    /// Type of the tool. Defaults to "function".
147    #[serde(default = "default_tool_type", rename = "type")]
148    type_: String,
149
150    /// The function definition of the tool.
151    function: FunctionDefinition,
152}
153
154fn default_tool_type() -> String {
155    "function".to_string()
156}
157
158/// Definition of a function that can be called as a tool.
159#[derive(Debug, Serialize, Deserialize)]
160pub struct FunctionDefinition {
161    /// The name of the function.
162    name: String,
163
164    /// A description of what the function does.
165    description: String,
166
167    /// Parameters for the function.
168    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/// Specifies the format of the response.
179#[derive(Debug, Serialize, Deserialize)]
180pub struct ResponseFormat {
181    /// Response type, such as "text" or "json_object".
182    #[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/// Details of a tool call made by the model.
191#[derive(Debug, Serialize, Deserialize)]
192pub struct ToolCall {
193    /// The ID of the tool call.
194    id: String,
195
196    /// The type of the tool call. Only "function" is supported currently.
197    #[serde(rename = "type")]
198    type_: String,
199
200    /// The function that the model called.
201    function: FunctionCall,
202}
203
204/// Represents a function call made by the model.
205#[derive(Debug, Serialize, Deserialize)]
206pub struct FunctionCall {
207    /// The name of the function to call.
208    name: String,
209
210    /// The function arguments in JSON format.
211    /// The model may not always generate valid JSON.
212    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    /// Sends a chat completion request to DeepInfra, returning a structured response.
249    ///
250    /// # Usage
251    /// ```
252    /// let request = ChatCompletionRequestBuilder::default()
253    ///     // Build your messages, model, temperature, etc.
254    ///     .build();
255    ///
256    /// let response = client.chat_completition(request).await?;
257    /// println!("Received chat response: {:?}", response);
258    /// ```
259    #[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}