Break apart common_types mod (#334)

Signed-off-by: José Ulises Niño Rivera <junr03@users.noreply.github.com>

crates/common/src/api/open_ai.rs

@@ -0,0 +1,653 @@
+use crate::consts::{ARCH_FC_MODEL_NAME, ASSISTANT_ROLE};
+use serde::{ser::SerializeMap, Deserialize, Serialize};
+use serde_yaml::Value;
+use std::{
+    collections::{HashMap, VecDeque},
+    fmt::Display,
+};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionsRequest {
+    #[serde(default)]
+    pub model: String,
+    pub messages: Vec<Message>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools: Option<Vec<ChatCompletionTool>>,
+    #[serde(default)]
+    pub stream: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream_options: Option<StreamOptions>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ToolType {
+    #[serde(rename = "function")]
+    Function,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionTool {
+    #[serde(rename = "type")]
+    pub tool_type: ToolType,
+    pub function: FunctionDefinition,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FunctionDefinition {
+    pub name: String,
+    pub description: String,
+    pub parameters: FunctionParameters,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct FunctionParameters {
+    pub properties: HashMap<String, FunctionParameter>,
+}
+
+impl Serialize for FunctionParameters {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        // select all requried parameters
+        let required: Vec<&String> = self
+            .properties
+            .iter()
+            .filter(|(_, v)| v.required.unwrap_or(false))
+            .map(|(k, _)| k)
+            .collect();
+        let mut map = serializer.serialize_map(Some(2))?;
+        map.serialize_entry("properties", &self.properties)?;
+        if !required.is_empty() {
+            map.serialize_entry("required", &required)?;
+        }
+        map.end()
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct FunctionParameter {
+    #[serde(rename = "type")]
+    #[serde(default = "ParameterType::string")]
+    pub parameter_type: ParameterType,
+    pub description: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub required: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(rename = "enum")]
+    pub enum_values: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub default: Option<String>,
+}
+
+impl Serialize for FunctionParameter {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        let mut map = serializer.serialize_map(Some(5))?;
+        map.serialize_entry("type", &self.parameter_type)?;
+        map.serialize_entry("description", &self.description)?;
+        if let Some(enum_values) = &self.enum_values {
+            map.serialize_entry("enum", enum_values)?;
+        }
+        if let Some(default) = &self.default {
+            map.serialize_entry("default", default)?;
+        }
+        map.end()
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub enum ParameterType {
+    #[serde(rename = "int")]
+    Int,
+    #[serde(rename = "float")]
+    Float,
+    #[serde(rename = "bool")]
+    Bool,
+    #[serde(rename = "str")]
+    String,
+    #[serde(rename = "list")]
+    List,
+    #[serde(rename = "dict")]
+    Dict,
+}
+
+impl From<String> for ParameterType {
+    fn from(s: String) -> Self {
+        match s.as_str() {
+            "int" => ParameterType::Int,
+            "integer" => ParameterType::Int,
+            "float" => ParameterType::Float,
+            "bool" => ParameterType::Bool,
+            "boolean" => ParameterType::Bool,
+            "str" => ParameterType::String,
+            "string" => ParameterType::String,
+            "list" => ParameterType::List,
+            "array" => ParameterType::List,
+            "dict" => ParameterType::Dict,
+            "dictionary" => ParameterType::Dict,
+            _ => ParameterType::String,
+        }
+    }
+}
+
+impl ParameterType {
+    pub fn string() -> ParameterType {
+        ParameterType::String
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StreamOptions {
+    pub include_usage: bool,
+}
+
+#[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 model: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<ToolCall>>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_call_id: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Choice {
+    pub finish_reason: String,
+    pub index: usize,
+    pub message: Message,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ToolCall {
+    pub id: String,
+    #[serde(rename = "type")]
+    pub tool_type: ToolType,
+    pub function: FunctionCallDetail,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FunctionCallDetail {
+    pub name: String,
+    pub arguments: HashMap<String, Value>,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct ToolCallState {
+    pub key: String,
+    pub message: Option<Message>,
+    pub tool_call: FunctionCallDetail,
+    pub tool_response: String,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(untagged)]
+pub enum ArchState {
+    ToolCall(Vec<ToolCallState>),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionsResponse {
+    pub usage: Option<Usage>,
+    pub choices: Vec<Choice>,
+    pub model: String,
+    pub metadata: Option<HashMap<String, String>>,
+}
+
+impl ChatCompletionsResponse {
+    pub fn new(message: String) -> Self {
+        ChatCompletionsResponse {
+            choices: vec![Choice {
+                message: Message {
+                    role: ASSISTANT_ROLE.to_string(),
+                    content: Some(message),
+                    model: Some(ARCH_FC_MODEL_NAME.to_string()),
+                    tool_calls: None,
+                    tool_call_id: None,
+                },
+                index: 0,
+                finish_reason: "done".to_string(),
+            }],
+            usage: None,
+            model: ARCH_FC_MODEL_NAME.to_string(),
+            metadata: None,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Usage {
+    pub completion_tokens: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatCompletionStreamResponse {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model: Option<String>,
+    pub choices: Vec<ChunkChoice>,
+}
+
+impl ChatCompletionStreamResponse {
+    pub fn new(
+        response: Option<String>,
+        role: Option<String>,
+        model: Option<String>,
+        tool_calls: Option<Vec<ToolCall>>,
+    ) -> Self {
+        ChatCompletionStreamResponse {
+            model,
+            choices: vec![ChunkChoice {
+                delta: Delta {
+                    role,
+                    content: response,
+                    tool_calls,
+                    model: None,
+                    tool_call_id: None,
+                },
+                finish_reason: None,
+            }],
+        }
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ChatCompletionChunkResponseError {
+    #[error("failed to deserialize")]
+    Deserialization(#[from] serde_json::Error),
+    #[error("empty content in data chunk")]
+    EmptyContent,
+    #[error("no chunks present")]
+    NoChunks,
+}
+
+pub struct ChatCompletionStreamResponseServerEvents {
+    pub events: Vec<ChatCompletionStreamResponse>,
+}
+
+impl Display for ChatCompletionStreamResponseServerEvents {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let tokens_str = self
+            .events
+            .iter()
+            .map(|response_chunk| {
+                if response_chunk.choices.is_empty() {
+                    return "".to_string();
+                }
+                response_chunk.choices[0]
+                    .delta
+                    .content
+                    .clone()
+                    .unwrap_or("".to_string())
+            })
+            .collect::<Vec<String>>()
+            .join("");
+
+        write!(f, "{}", tokens_str)
+    }
+}
+
+impl TryFrom<&str> for ChatCompletionStreamResponseServerEvents {
+    type Error = ChatCompletionChunkResponseError;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let response_chunks: VecDeque<ChatCompletionStreamResponse> = value
+            .lines()
+            .filter(|line| line.starts_with("data: "))
+            .map(|line| line.get(6..).unwrap())
+            .filter(|data_chunk| *data_chunk != "[DONE]")
+            .map(serde_json::from_str::<ChatCompletionStreamResponse>)
+            .collect::<Result<VecDeque<ChatCompletionStreamResponse>, _>>()?;
+
+        Ok(ChatCompletionStreamResponseServerEvents {
+            events: response_chunks.into(),
+        })
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChunkChoice {
+    pub delta: Delta,
+    // TODO: could this be an enum?
+    pub finish_reason: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Delta {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub role: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub content: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<ToolCall>>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_call_id: Option<String>,
+}
+
+pub fn to_server_events(chunks: Vec<ChatCompletionStreamResponse>) -> String {
+    let mut response_str = String::new();
+    for chunk in chunks.iter() {
+        response_str.push_str("data: ");
+        response_str.push_str(&serde_json::to_string(&chunk).unwrap());
+        response_str.push_str("\n\n");
+    }
+    response_str
+}
+
+#[cfg(test)]
+mod test {
+    use super::{ChatCompletionStreamResponseServerEvents, Message};
+    use pretty_assertions::assert_eq;
+    use std::collections::HashMap;
+
+    const TOOL_SERIALIZED: &str = r#"{
+  "model": "gpt-3.5-turbo",
+  "messages": [
+    {
+      "role": "user",
+      "content": "What city do you want to know the weather for?"
+    }
+  ],
+  "tools": [
+    {
+      "type": "function",
+      "function": {
+        "name": "weather_forecast",
+        "description": "function to retrieve weather forecast",
+        "parameters": {
+          "properties": {
+            "city": {
+              "type": "str",
+              "description": "city for weather forecast",
+              "default": "test"
+            }
+          },
+          "required": [
+            "city"
+          ]
+        }
+      }
+    }
+  ],
+  "stream": true,
+  "stream_options": {
+    "include_usage": true
+  }
+}"#;
+
+    #[test]
+    fn test_tool_type_request() {
+        use super::{
+            ChatCompletionTool, ChatCompletionsRequest, FunctionDefinition, FunctionParameter,
+            FunctionParameters, ParameterType, StreamOptions, ToolType,
+        };
+
+        let mut properties = HashMap::new();
+        properties.insert(
+            "city".to_string(),
+            FunctionParameter {
+                parameter_type: ParameterType::String,
+                description: "city for weather forecast".to_string(),
+                required: Some(true),
+                enum_values: None,
+                default: Some("test".to_string()),
+            },
+        );
+
+        let function_definition = FunctionDefinition {
+            name: "weather_forecast".to_string(),
+            description: "function to retrieve weather forecast".to_string(),
+            parameters: FunctionParameters { properties },
+        };
+
+        let chat_completions_request = ChatCompletionsRequest {
+            model: "gpt-3.5-turbo".to_string(),
+            messages: vec![Message {
+                role: "user".to_string(),
+                content: Some("What city do you want to know the weather for?".to_string()),
+                model: None,
+                tool_calls: None,
+                tool_call_id: None,
+            }],
+            tools: Some(vec![ChatCompletionTool {
+                tool_type: ToolType::Function,
+                function: function_definition,
+            }]),
+            stream: true,
+            stream_options: Some(StreamOptions {
+                include_usage: true,
+            }),
+            metadata: None,
+        };
+
+        let serialized = serde_json::to_string_pretty(&chat_completions_request).unwrap();
+        println!("{}", serialized);
+        assert_eq!(TOOL_SERIALIZED, serialized);
+    }
+
+    #[test]
+    fn test_parameter_types() {
+        use super::{FunctionParameter, ParameterType};
+
+        const PARAMETER_SERIALZIED: &str = r#"{
+  "city": {
+    "type": "str",
+    "description": "city for weather forecast",
+    "default": "test"
+  }
+}"#;
+
+        let properties = HashMap::from([(
+            "city".to_string(),
+            FunctionParameter {
+                parameter_type: ParameterType::String,
+                description: "city for weather forecast".to_string(),
+                required: Some(true),
+                enum_values: None,
+                default: Some("test".to_string()),
+            },
+        )]);
+
+        let serialized = serde_json::to_string_pretty(&properties).unwrap();
+        assert_eq!(PARAMETER_SERIALZIED, serialized);
+
+        // ensure that if type is missing it is set to string
+        const PARAMETER_SERIALZIED_MISSING_TYPE: &str = r#"
+        {
+          "city": {
+            "description": "city for weather forecast"
+          }
+        }"#;
+
+        let missing_type_deserialized: HashMap<String, FunctionParameter> =
+            serde_json::from_str(PARAMETER_SERIALZIED_MISSING_TYPE).unwrap();
+        println!("{:?}", missing_type_deserialized);
+        assert_eq!(
+            missing_type_deserialized
+                .get("city")
+                .unwrap()
+                .parameter_type,
+            ParameterType::String
+        );
+    }
+
+    #[test]
+    fn stream_chunk_parse() {
+        const CHUNK_RESPONSE: &str = r#"data: {"id":"chatcmpl-ALmdmtKulBMEq3fRLbrnxJwcKOqvS","object":"chat.completion.chunk","created":1729755226,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALmdmtKulBMEq3fRLbrnxJwcKOqvS","object":"chat.completion.chunk","created":1729755226,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"Hello"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALmdmtKulBMEq3fRLbrnxJwcKOqvS","object":"chat.completion.chunk","created":1729755226,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"!"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALmdmtKulBMEq3fRLbrnxJwcKOqvS","object":"chat.completion.chunk","created":1729755226,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" How"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALmdmtKulBMEq3fRLbrnxJwcKOqvS","object":"chat.completion.chunk","created":1729755226,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" can"},"logprobs":null,"finish_reason":null}]}
+
+
+"#;
+
+        let sever_events =
+            ChatCompletionStreamResponseServerEvents::try_from(CHUNK_RESPONSE).unwrap();
+        assert_eq!(sever_events.events.len(), 5);
+        assert_eq!(
+            sever_events.events[0].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            ""
+        );
+        assert_eq!(
+            sever_events.events[1].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            "Hello"
+        );
+        assert_eq!(
+            sever_events.events[2].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            "!"
+        );
+        assert_eq!(
+            sever_events.events[3].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " How"
+        );
+        assert_eq!(
+            sever_events.events[4].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " can"
+        );
+        assert_eq!(sever_events.to_string(), "Hello! How can");
+    }
+
+    #[test]
+    fn stream_chunk_parse_done() {
+        const CHUNK_RESPONSE: &str = r#"data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" I"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" assist"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" you"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":" today"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{"content":"?"},"logprobs":null,"finish_reason":null}]}
+
+data: {"id":"chatcmpl-ALn2KTfmrIpYd9N3Un4Kyg08WIIP6","object":"chat.completion.chunk","created":1729756748,"model":"gpt-3.5-turbo-0125","system_fingerprint":null,"choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}
+
+data: [DONE]
+"#;
+
+        let sever_events: ChatCompletionStreamResponseServerEvents =
+            ChatCompletionStreamResponseServerEvents::try_from(CHUNK_RESPONSE).unwrap();
+        assert_eq!(sever_events.events.len(), 6);
+        assert_eq!(
+            sever_events.events[0].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " I"
+        );
+        assert_eq!(
+            sever_events.events[1].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " assist"
+        );
+        assert_eq!(
+            sever_events.events[2].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " you"
+        );
+        assert_eq!(
+            sever_events.events[3].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            " today"
+        );
+        assert_eq!(
+            sever_events.events[4].choices[0]
+                .delta
+                .content
+                .as_ref()
+                .unwrap(),
+            "?"
+        );
+        assert_eq!(sever_events.events[5].choices[0].delta.content, None);
+
+        assert_eq!(sever_events.to_string(), " I assist you today?");
+    }
+
+    #[test]
+    fn stream_chunk_parse_mistral() {
+        const CHUNK_RESPONSE: &str = r#"data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" How"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" can"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" I"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" assist"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" you"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":" today"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":"?"},"finish_reason":null}]}
+
+data: {"id":"e1ebce16de5443b79613512c2d757936","object":"chat.completion.chunk","created":1729805261,"model":"ministral-8b-latest","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"total_tokens":13,"completion_tokens":9}}
+
+data: [DONE]
+"#;
+
+        let sever_events: ChatCompletionStreamResponseServerEvents =
+            ChatCompletionStreamResponseServerEvents::try_from(CHUNK_RESPONSE).unwrap();
+        assert_eq!(sever_events.events.len(), 11);
+
+        assert_eq!(
+            sever_events.to_string(),
+            "Hello! How can I assist you today?"
+        );
+    }
+}