Skip to main content

kronos_tide/
ai.rs

1use std::process::Command;
2
3pub struct AIClient {
4    provider: String,
5    model: String,
6    temperature: f32,
7    base_url: String,
8}
9
10impl AIClient {
11    pub fn new(provider: &str, model: &str, temperature: f32) -> Self {
12        let base_url = match provider {
13            "ollama" => "http://127.0.0.1:11434".to_string(),
14            _ => "http://127.0.0.1:11434".to_string(),
15        };
16        Self {
17            provider: provider.to_string(),
18            model: model.to_string(),
19            temperature,
20            base_url,
21        }
22    }
23
24    pub fn ask(&self, prompt: &str, system: &str) -> Result<String, String> {
25        let payload = serde_json::json!({
26            "model": self.model,
27            "prompt": prompt,
28            "system": system,
29            "stream": false,
30            "options": {
31                "temperature": self.temperature,
32                "num_predict": 4096
33            }
34        });
35
36        let output = Command::new("curl")
37            .args([
38                "-s",
39                "-X",
40                "POST",
41                &format!("{}/api/generate", self.base_url),
42                "-H",
43                "Content-Type: application/json",
44                "-d",
45                &payload.to_string(),
46            ])
47            .output()
48            .map_err(|e| format!("curl failed: {}", e))?;
49
50        if !output.status.success() {
51            return Err(format!(
52                "HTTP error: {}",
53                String::from_utf8_lossy(&output.stderr)
54            ));
55        }
56
57        let response: serde_json::Value = serde_json::from_slice(&output.stdout)
58            .map_err(|e| format!("JSON parse error: {}", e))?;
59
60        response["response"]
61            .as_str()
62            .map(|s| s.to_string())
63            .ok_or_else(|| "no response field".to_string())
64    }
65
66    pub fn function_call(
67        &self,
68        prompt: &str,
69        functions: &[FunctionSpec],
70    ) -> Result<FunctionCall, String> {
71        let fn_json: Vec<serde_json::Value> = functions
72            .iter()
73            .map(|f| {
74                serde_json::json!({
75                    "name": f.name,
76                    "description": f.description,
77                    "parameters": f.parameters
78                })
79            })
80            .collect();
81
82        let system = format!(
83            "You are a function-calling assistant. Given the user query, respond with a JSON object \
84             containing 'name' (function name) and 'arguments' (dict of args). Available functions: {}",
85            serde_json::to_string(&fn_json).unwrap_or_default()
86        );
87
88        let response = self.ask(prompt, &system)?;
89
90        let parsed: serde_json::Value = serde_json::from_str(&response)
91            .map_err(|_| format!("Could not parse function call: {}", response))?;
92
93        Ok(FunctionCall {
94            name: parsed["name"].as_str().unwrap_or("unknown").to_string(),
95            arguments: parsed["arguments"].clone(),
96        })
97    }
98}
99
100#[derive(Debug, Clone)]
101pub struct FunctionSpec {
102    pub name: String,
103    pub description: String,
104    pub parameters: serde_json::Value,
105}
106
107#[derive(Debug, Clone)]
108pub struct FunctionCall {
109    pub name: String,
110    pub arguments: serde_json::Value,
111}