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}