ai.rs

 1pub mod assistant;
 2mod assistant_settings;
 3
 4pub use assistant::AssistantPanel;
 5use gpui::AppContext;
 6use serde::{Deserialize, Serialize};
 7use std::fmt::{self, Display};
 8
 9// Data types for chat completion requests
10#[derive(Serialize)]
11struct OpenAIRequest {
12    model: String,
13    messages: Vec<RequestMessage>,
14    stream: bool,
15}
16
17#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
18struct RequestMessage {
19    role: Role,
20    content: String,
21}
22
23#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
24struct ResponseMessage {
25    role: Option<Role>,
26    content: Option<String>,
27}
28
29#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
30#[serde(rename_all = "lowercase")]
31enum Role {
32    User,
33    Assistant,
34    System,
35}
36
37impl Role {
38    pub fn cycle(&mut self) {
39        *self = match self {
40            Role::User => Role::Assistant,
41            Role::Assistant => Role::System,
42            Role::System => Role::User,
43        }
44    }
45}
46
47impl Display for Role {
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
49        match self {
50            Role::User => write!(f, "User"),
51            Role::Assistant => write!(f, "Assistant"),
52            Role::System => write!(f, "System"),
53        }
54    }
55}
56
57#[derive(Deserialize, Debug)]
58struct OpenAIResponseStreamEvent {
59    pub id: Option<String>,
60    pub object: String,
61    pub created: u32,
62    pub model: String,
63    pub choices: Vec<ChatChoiceDelta>,
64    pub usage: Option<Usage>,
65}
66
67#[derive(Deserialize, Debug)]
68struct Usage {
69    pub prompt_tokens: u32,
70    pub completion_tokens: u32,
71    pub total_tokens: u32,
72}
73
74#[derive(Deserialize, Debug)]
75struct ChatChoiceDelta {
76    pub index: u32,
77    pub delta: ResponseMessage,
78    pub finish_reason: Option<String>,
79}
80
81#[derive(Deserialize, Debug)]
82struct OpenAIUsage {
83    prompt_tokens: u64,
84    completion_tokens: u64,
85    total_tokens: u64,
86}
87
88#[derive(Deserialize, Debug)]
89struct OpenAIChoice {
90    text: String,
91    index: u32,
92    logprobs: Option<serde_json::Value>,
93    finish_reason: Option<String>,
94}
95
96pub fn init(cx: &mut AppContext) {
97    assistant::init(cx);
98}