1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

#![allow(dead_code)]
#![feature(type_name_of_val)]
#![feature(associated_type_defaults)]
#![feature(async_fn_in_trait)]
#![allow(async_fn_in_trait)]

mod chat_completion;
mod chat_completion_delta;
mod chat_completion_request;
mod error;

use lazy_static::lazy_static;
#[allow(unused_imports)]
use log::{debug, error, info, trace, warn};
use std::{
    any::type_name_of_val,
    sync::{Arc, RwLock},
};

use schemars::{schema_for, JsonSchema};
use serde_derive::{Deserialize, Serialize};
use serde_json::Value;
pub use {
    chat_completion::ChatCompletion as Chat,
    chat_completion_delta::ChatCompletionDelta as ChatDelta, chat_completion_delta::DeltaReceiver,
    chat_completion_request::ChatCompletionRequest as ChatRequest,
    chat_completion_request::AiAgent as AiAgent,
};

lazy_static! {
    static ref OPENAI_API_KEY: Arc<RwLock<Option<String>>> = Arc::new(RwLock::new(None));
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub function_call: Option<FunctionCall>,
}

impl Message {
    pub fn new(role: impl Into<String>) -> Self {
        Self {
            role: role.into(),
            content: None,
            name: None,
            function_call: None,
        }
    }

    pub fn with_content(mut self, content: impl Into<String>) -> Self {
        self.content = Some(content.into());
        self
    }

    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    pub name: String,
    pub arguments: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Function {
    pub name: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub parameters: Value,
}

impl Function {
    pub fn from<FunctionArgs, Func, T>(function: Func) -> Self
    where
        FunctionArgs: JsonSchema,
        Func: Fn(FunctionArgs) -> T + 'static,
    {
        let schema = schema_for!(FunctionArgs);
        let fn_type_name = type_name_of_val(&function);
        let parameters = serde_json::to_value(schema)
            .unwrap_or_else(|_| panic!("Failed to serialize schema for function {}", fn_type_name));

        let fn_name = fn_type_name.split("::").last().unwrap_or("");
        Self {
            name: fn_name.to_string(),
            description: match parameters.get("description") {
                Some(Value::String(s)) => Some(s.clone()),
                _ => None,
            },
            parameters,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Choice {
    pub index: i64,
    pub message: Message,
    pub finish_reason: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChoiceDelta {
    pub index: i64,
    pub delta: Delta,

    #[serde(default)]
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Delta {
    pub role: Option<String>,

    pub content: Option<String>,

    pub function_call: Option<FunctionCallDelta>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCallDelta {
    pub name: Option<String>,
    pub arguments: Option<String>,
}

#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Usage {
    pub prompt_tokens: i64,
    pub completion_tokens: i64,
    pub total_tokens: i64,
}

pub fn api_key(api_key: String) {
    let mut key = OPENAI_API_KEY.write().unwrap();
    *key = Some(api_key);
}

#[derive(Default, Debug, Clone, JsonSchema)]
#[schemars(description = "this function takes no arguments")]
pub struct NoArgs {
    _unused: (),
}