llmclient 0.3.2

use regex::Regex;
use serde_derive::{Deserialize, Serialize};
use crate::gpt::*;

// Input structures
// Function Calls

#[derive(Debug, Serialize, Clone)]
pub struct Functions {
    pub model: String,
    pub tools: Vec<FunctionCall>,
}

impl Functions {
    pub fn new(model: &str, tools: Vec<FunctionCall>) -> Self {
        Functions { model: model.to_string(), tools }
    }
}

#[derive(Debug, Serialize, Clone)]
pub struct FunctionCall {
    pub r#type: String,
    pub function: Function,
}

#[derive(Debug, Serialize, Clone)]
pub struct Function {
    pub name: String,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<Parameters>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub input_schema: Option<Parameters>,
    pub required: Vec<String>
}

#[derive(Debug, Serialize, Clone)]
pub struct Parameters {
    pub r#type: String,
    pub properties: Properties,
    pub required: Vec<String>
}

#[derive(Debug, Serialize, Clone)]
pub struct Properties {
    #[serde(rename = "function")]
    pub function: Vec<ParameterType>
}

#[derive(Debug, Serialize, Clone)]
pub struct ParameterType {
    pub r#type: String,
    //pub r#enum: String,
    pub description: String,
}

#[derive(Debug, Deserialize)]
pub struct Choices {
    index: u32,
    message: Message,
    #[serde(skip_serializing_if = "Option::is_none")]
    logprobs: Option<f32>,
    finish_reason: String,
}

#[derive(Debug, Deserialize)]
pub struct Message {
    role: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    content: Option<String>,
    tool_calls: Vec<Calls>,
}

#[derive(Debug, Deserialize)]
pub struct Calls {
    id: String,
    r#type: String,
    function: FuncionSummary,
}

#[derive(Debug, Deserialize)]
pub struct FuncionSummary {
    name: String,
    arguments: FunctionArgs,
}

#[derive(Debug, Deserialize)]
pub struct FunctionArgs {
    #[serde(rename = "arg")]
    pub arg: String,
}

#[derive(Debug, Deserialize)]
pub struct Usage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

/*
#[derive(Debug, Deserialize)]
pub struct FunctionResponse {
    id: String,
    object: String,
    created: u64,
    model: String,
    choices: Choices,
    usage: Usage,
    system_fingerprint: String,
}
*/

#[derive(Debug, Deserialize)]
pub struct FunctionResponse {
    pub r#type: String,
    pub id: String,
    pub name: String,
    pub input: Vec<ParameterValues>
}

#[derive(Debug, Deserialize)]
pub struct ParameterValues {
    #[serde(rename = "parameter")]
    pub parameter: String,
}

/// Takes function definition of the form 'fn <fn_name>(arg1, ... argn)'
/// All areguments are assumed to be strings as is the return value
fn get_func(expr: &str) -> Option<(String, Vec<String>)> {
    let re = Regex::new(r#"fn ([a-zA-Z]+)\(([a-zA-Z,\s]+)\)"#).ok()?;
    let Some(bits) = re.captures(expr) else { return None };
    let args: Vec<&str> = bits[2].split(',').collect();
    let args = args.iter().map(|s| s.trim().to_string()).collect::<Vec<String>>();

    Some((bits[1].to_string(), args))
}

fn get_func_json(exprs: Vec<&str>) -> String {
     "".to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_call_function() {
        let messages: Vec<GptMessage> = vec![GptMessage { role: "user".into(), content: "What is the weather like in San Francisco?".into()}];
        let model: String = std::env::var("GPT_MODEL")
            .expect("GPT_MODEL not found in enviroment variables");
        match call_gpt(messages).await {
            Ok(ret) => { println!("{ret}"); assert!(true) },
            Err(e) => { println!("{e}"); assert!(false) },
        }
    }
}