llmclient 0.3.2

use std::collections::HashMap;
use serde_json::Value;
use serde_json::Value::*;
use std::string::String;
use regex::Regex;
use serde_derive::Deserialize;
use evalexpr::*;
use peg::*;
use peg::error::ParseError;
use peg::str::LineCol;
use stemplate::*;
//use llmclient::common::call_llm;

#[derive(Debug, Clone, Deserialize)]
pub struct Function {
    function: String,
    arguments: Vec<Argument>,
}

impl Function {
    fn new(function: &str, arguments: Vec<Argument>) -> Self {
        Function { function: function.to_string(), arguments }
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct Argument {
    name: String,
    value: String,
}

impl Argument {
    fn new(name: &str, value: &str) -> Self {
        Argument { name: name.to_string(), value: value.to_string() }
    }
}

fn json_filler(provider: &str, func_defs: &[&str]) -> Result<String, ParseError<LineCol>> {
    let func = match provider {
        "gpt" | "mistral" => 
r#"
"type": "function",
"function": {
    "name": "${func}",
    "description": "${func_desc}",
    "parameters": {
        "type": "object",
        "properties": {
        ${*,all_args}
        },
        "required": [${*,mand_args}]
    }
}
"#,
        _ => 
r#"
{
    "name": "${func}",
    "description": "${func_desc}",
    "input_schema": {
      "type": "object",
      "properties": {
      ${*,all_args}
      },
      "required": [${*,mand_args}]
    }
}"#,
   };

    let all_args = r#"
"${arg}": {
  "type": "string",
  "description": "${arg_desc}"
}
"#;

    let mand_args = r#""${marg}""#;

    let defs: Vec<String> = func_defs.iter()
        .map(|f| llmfunc::func(f, func, all_args, mand_args))
        .filter(|f| f.is_ok())
        .map(|f| f.unwrap())
        .collect();

    Ok(defs.join(","))
}

fn main() {
    let data = r#"{
  "choices": [
    {
      "message": {
        "tool_calls": [
          {
            "function": {
              "name": "${func}",
              "arguments": "${args}"
            }
          }
        ]
      }
    }
  ]
}"#;
    let data2 = r#"{
  "id": "chatcmpl-9KmDSgvDNpdWzZMjNikOFzWyZRi8D",
  "object": "chat.completion",
  "created": 1714738930,
  "model": "gpt-4-turbo-2024-04-09",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": null,
        "tool_calls": [
          {
            "id": "call_lPmkbGFiJkvuXvzs09uUYdzD",
            "type": "function",
            "function": {
              "name": "calc",
              "arguments": "{\"expr\":\"(60 * 24) * 265.251\"}"
            }
          },
          {
            "id": "call_lPmkbGFiJkvuXvzs09uUYdzD",
            "type": "function",
            "function": {
              "name": "calc2",
              "arguments": "{\"expr\":\"(60 * 24)\", \"unit\": \"hours\"}"
            }
          }
        ]
      },
      "logprobs": null,
      "finish_reason": "tool_calls"
    }
  ],
  "usage": {
    "prompt_tokens": 80,
    "completion_tokens": 23,
    "total_tokens": 103
  },
  "system_fingerprint": "fp_3450ce39d5"
}"#;
    let func_def =
r#"
// this is a func
// arg0: The first arg
// arg1: The Second arg
fn func(arg0, *arg1)
"#;

    let func_def2 =
r#"
// this is another func
// arg10: The first arg
// arg11: The Second arg
fn func(arg10, arg11)
"#;

    let func_call = json_filler("gpt", &[func_def, func_def2]);

    println!("{}", func_call.unwrap());

    // Convert to a string of JSON and print it out
    let v: serde_json::Value = serde_json::from_str(data).unwrap();

    let found = find_function(&v);
    println!("{found:?}");

    let f: serde_json::Value = serde_json::from_str(data2).unwrap();

    let h = get_functions(&f, &found);

    println!("{h:?}");

    let funcs = unpack_functions(h);

    println!("{funcs:?}");

    if let Some(fun) = funcs {
        for f in fun {
            let exprs: Vec<&Argument> = f.arguments.iter()
                .filter(|p| p.name == "expr")
                .collect();
            let value = eval(&exprs[0].value);

            if let Ok(v) = value {
                println!("{}: {:?}", &f.function, v.to_string());
            }
        }
    }
}

fn unpack_functions(h: HashMap<String, Vec<String>>) -> Option<Vec<Function>> {
    let func = h.get("func");
    let args = h.get("args");

    if let Some(func) = func {
        if let Some(args) = args {
            let funcs = func.iter().zip(args.iter())
                .map(|(f, a)| {
                    if a.starts_with('{') && a.ends_with("}") {
                        let fh: HashMap<String, String> = serde_json::from_str(a).unwrap();
                        let args: Vec<Argument> = fh.iter()
                            .map(|(pn, pv)| Argument::new(pn, pv))
                            .collect();

                        Function::new(f, args)
                    } else {
                        Function::new(f, vec![])
                    }
                })
                .collect();

            return Some(funcs);
        }
    }

    None
}

fn get_functions(val: &Value, found: &Vec<String>) -> HashMap<String, Vec<String>> {
    fn getter(val: &Value, places: &str, found: &mut HashMap<String, Vec<String>>) -> HashMap<String, Vec<String>> {
        let mut v = val;
        let items: Vec<&str> = places.split(':').collect();
        let var = *items.last().unwrap();

        for (pos, i) in items.iter().enumerate() {
            if *i == var {
                if let Value::String(it) = v {
                    let key = (&i[2..var.len()-1]).to_string();

                    found.entry(key)
                        .and_modify(|v| v.push(it.to_string()))
                        .or_insert(vec![it.to_string()]);
                }
            } else {
                if let Value::Array(it) = v {
                    for a in it {
                        getter(a, &items[pos..].join(":"), found);
                    }
                } else {
                    v = &v[i];
                }
            }
        }

        found.clone()
    }

    let mut res: HashMap<String, Vec<String>> = HashMap::new();

    for i in found {
        getter(val, &i, &mut res);
    }

    res
}

fn find_function(v: &Value) -> Vec<String> {
    fn finder(v: &Value, res: String, found: &mut Vec<String>) -> Vec<String> {
        match v {
            Null => {
            },
            Bool(_b) => {
            },
            Number(_n) => {
            },
            String(s) => {
                let re = Regex::new(r#"\$\{[A-Za-z0-9_]+\}"#).unwrap();
                if re.is_match(s) {
                    let f = format!("{res}{s}");

                    if !found.contains(&f) {
                        found.push(f);
                    }
                }
            },
            Array(a) => {
                for v in a {
                    finder(v, res.clone(), found);
                }
            },
            Object(o) => {
                for (k, v) in o.iter() {
                    finder(v, res.clone() + k + ":", found);
                }
            },
        };

        found.clone()
    }

    finder(&v, String::new(), &mut vec![])
}

peg::parser!( grammar llmfunc() for str {
    pub rule func(func: &str, all_args: &str, mand_args: &str) -> String
        = "\n"* fc:func_comment()+ ac:arg_comment()+ "fn"? _ f:ident() _ "(" a:arg_ident() ** comma() ")" _ "\n"* _ {
            //format!("{:?}", ( fc, ac, f, a ))
            let cnt = ac.iter().enumerate()
                .filter(|(i, arg)| {
                    //*arg.0 == *a[*i] || "*".to_string() + &arg.0 == a[*i]
                    arg.starts_with(a[*i]) || format!("*{arg}").starts_with(a[*i])
                }).count();
            if ac.len() == a.len() && a.len() == cnt {
                let ma: Vec<&str> = a.iter()
                    .filter(|a| !a.starts_with("*"))
                    .map(|a| &a[..])
                    .collect();
                let a: Vec<&str> = a.iter()
                    .map(|a| if a.starts_with("*") { &a[1..] } else { a })
                    .collect();
                let mut h: HashMap<&str, String> = HashMap::new();

                h.insert("func", f.to_string());
                h.insert("func_desc", fc[0].to_string());
                h.insert("arg", a.join("|"));
                h.insert("arg_desc", ac.join("|"));
                h.insert("marg", ma.join("|"));
                h.insert("all_args", all_args.to_string());
                h.insert("mand_args", mand_args.to_string());
                h.insert("func_call", func.to_string());

                Template::new("${func_call}").render(&h)
                //format!("{:?}", ( fc, ac, f, a.clone(), ma))
            } else {
                "Error: Argument names do not match".to_string()
            }
        }

    rule _ = [' ']*

    rule comma() = "," " "* 

    rule ident() -> &'input str
        = s:$(['a'..='z'|'0'..='9'|'_']+) { s }

    rule arg_ident() -> &'input str
        = s:$("*"? ident()) { s }

    rule func_comment() -> &'input str
        = _ "/"*<2,3> _ s:comment() _ "\n"+ { s }

    rule arg_comment() -> &'input str
        = _ "/"*<2,3> _ a:$(ident() ":" comment()) _ "\n"+ { a }
    //rule arg_comment() -> (&'input str, &'input str)
        //= "//" _ a:ident() ":" _ s:comment() _ "\n"+ { (a, s) }

    rule comment() -> &'input str
        = s:$(['a'..='z'|'A'..='Z'|'0'..='9'|'_'|' ']+) { s }
});