kora_models/
provider.rs

1//! Provider clients: OpenAI (API) and Ollama (localhost).
2//!
3//! HTTP is isolated behind [`Transport`] so request construction and response
4//! handling are testable without touching the network.
5
6use serde_json::{json, Value};
7
8use crate::base64;
9use crate::schema::{build_json_schema, system_prompt, user_prompt};
10use crate::validate::{parse_response, truncate};
11use crate::{AnalyzeRequest, FieldType, ModelConfig, ModelError, Provider, Step, ToolSpec};
12
13/// Long enough for a large local model to read an image, since a timeout
14/// that fires on ordinary work teaches people to raise it blindly. Override
15/// with `[models] timeout_secs`.
16pub const DEFAULT_TIMEOUT_SECS: u64 = 600;
17const OPENAI_BASE: &str = "https://api.openai.com/v1";
18const OLLAMA_BASE: &str = "http://localhost:11434";
19
20/// One HTTP POST: (url, headers, body) -> response body text.
21pub(crate) type Transport = dyn Fn(&str, &[(&str, String)], &Value) -> Result<String, ModelError>;
22
23/// `"openai:gpt-4o"` / `"local:llama3.1:8b"` -> config.
24///
25/// Everything after the first colon is the model name, so Ollama tags
26/// (`llama3.1:8b`) survive intact.
27pub fn parse_model_spec(spec: &str) -> Result<ModelConfig, ModelError> {
28    let (scheme, model) = spec.split_once(':').ok_or_else(|| {
29        ModelError::new(format!(
30            "model spec `{spec}` needs a provider prefix, e.g. `openai:gpt-4o` or `local:llama3.1:8b`"
31        ))
32    })?;
33    if model.trim().is_empty() {
34        return Err(ModelError::new(format!(
35            "model spec `{spec}` has no model name after `{scheme}:`"
36        )));
37    }
38    let provider = match scheme {
39        "openai" => Provider::OpenAI,
40        "local" | "ollama" => Provider::Ollama,
41        other => {
42            return Err(ModelError::new(format!(
43                "unknown model provider `{other}` (expected `openai` or `local`)"
44            )))
45        }
46    };
47    Ok(ModelConfig {
48        provider,
49        model: model.to_string(),
50        endpoint: None,
51        api_key: None,
52        max_output_tokens: 4096,
53        timeout_secs: DEFAULT_TIMEOUT_SECS,
54    })
55}
56
57pub(crate) fn step_with(
58    config: &ModelConfig,
59    req: &AnalyzeRequest,
60    transport: &Transport,
61) -> Result<Step, ModelError> {
62    match config.provider {
63        Provider::OpenAI => openai(config, req, transport),
64        Provider::Ollama => ollama(config, req, transport),
65    }
66}
67
68/// Tool declarations in the shape both providers accept.
69fn tools_json(tools: &[ToolSpec]) -> Value {
70    Value::Array(
71        tools
72            .iter()
73            .map(|tool| {
74                let mut properties = serde_json::Map::new();
75                let mut required = Vec::new();
76                for (name, ty) in &tool.params {
77                    properties.insert(name.clone(), param_schema(ty));
78                    required.push(Value::String(name.clone()));
79                }
80                json!({
81                    "type": "function",
82                    "function": {
83                        "name": tool.name,
84                        "description": tool.description,
85                        "parameters": {
86                            "type": "object",
87                            "properties": properties,
88                            "required": required,
89                        }
90                    }
91                })
92            })
93            .collect(),
94    )
95}
96
97fn param_schema(ty: &FieldType) -> Value {
98    match ty {
99        FieldType::Str => json!({"type": "string"}),
100        FieldType::Int => json!({"type": "integer"}),
101        FieldType::Float => json!({"type": "number"}),
102        FieldType::Bool => json!({"type": "boolean"}),
103        FieldType::ListOfStr => json!({"type": "array", "items": {"type": "string"}}),
104    }
105}
106
107/// Conversation messages: system, user, then any tool exchanges so far.
108///
109/// The two providers attach images differently — OpenAI splits the user
110/// message into typed content parts, Ollama keeps plain text and hangs a
111/// parallel `images` array off the message — so the provider decides the
112/// shape rather than the caller.
113fn messages(req: &AnalyzeRequest, provider: &Provider) -> Vec<Value> {
114    let text = user_prompt(&req.prompt, &req.data_json);
115    let user = if req.images.is_empty() {
116        json!({"role": "user", "content": text})
117    } else {
118        match provider {
119            Provider::OpenAI => {
120                let mut parts = vec![json!({"type": "text", "text": text})];
121                for image in &req.images {
122                    parts.push(json!({
123                        "type": "image_url",
124                        "image_url": {
125                            "url": format!(
126                                "data:{};base64,{}",
127                                image.mime,
128                                base64::encode(&image.bytes)
129                            )
130                        }
131                    }));
132                }
133                json!({"role": "user", "content": parts})
134            }
135            Provider::Ollama => {
136                let encoded: Vec<Value> = req
137                    .images
138                    .iter()
139                    .map(|i| Value::String(base64::encode(&i.bytes)))
140                    .collect();
141                json!({"role": "user", "content": text, "images": encoded})
142            }
143        }
144    };
145
146    let mut out = vec![
147        json!({"role": "system", "content": system_prompt(&req.schema)}),
148        user,
149    ];
150    for exchange in &req.tool_history {
151        out.push(json!({
152            "role": "assistant",
153            "content": format!("Calling {}({})", exchange.name, exchange.arguments_json),
154        }));
155        out.push(json!({
156            "role": "user",
157            "content": format!("Result of {}: {}", exchange.name, exchange.result_json),
158        }));
159    }
160    out
161}
162
163fn openai(
164    config: &ModelConfig,
165    req: &AnalyzeRequest,
166    transport: &Transport,
167) -> Result<Step, ModelError> {
168    let key = config
169        .api_key
170        .clone()
171        .or_else(|| std::env::var("OPENAI_API_KEY").ok())
172        .filter(|k| !k.trim().is_empty())
173        .ok_or_else(|| {
174            ModelError::new("OPENAI_API_KEY not set (export it, or set api_key in kora.toml)")
175        })?;
176
177    let mut body = json!({
178        "model": config.model,
179        "max_completion_tokens": config.max_output_tokens,
180        "messages": messages(req, &Provider::OpenAI),
181    });
182    if req.tools.is_empty() {
183        // Structured output and tool calling are mutually exclusive shapes:
184        // constrain the answer only once no tool can still be requested.
185        body["response_format"] = json!({
186            "type": "json_schema",
187            "json_schema": {
188                "name": sanitize_schema_name(&req.schema.type_name),
189                "strict": true,
190                "schema": build_json_schema(&req.schema),
191            }
192        });
193    } else {
194        body["tools"] = tools_json(&req.tools);
195    }
196
197    let headers = [
198        ("Authorization", format!("Bearer {key}")),
199        ("Content-Type", "application/json".to_string()),
200    ];
201    let url = format!("{OPENAI_BASE}/chat/completions");
202    let text = transport(&url, &headers, &body)?;
203    let response: Value = serde_json::from_str(&text).map_err(|e| {
204        ModelError::new(format!(
205            "OpenAI returned a non-JSON body ({e}): {}",
206            truncate(&text, 300)
207        ))
208    })?;
209
210    let tokens_in = response["usage"]["prompt_tokens"].as_u64().unwrap_or(0);
211    let tokens_out = response["usage"]["completion_tokens"].as_u64().unwrap_or(0);
212
213    let message = &response["choices"][0]["message"];
214    if let Some(call) = message["tool_calls"].get(0) {
215        let name = call["function"]["name"].as_str().unwrap_or_default();
216        let arguments_json = call["function"]["arguments"]
217            .as_str()
218            .unwrap_or("{}")
219            .to_string();
220        return Ok(Step::CallTool {
221            name: name.to_string(),
222            arguments_json,
223            tokens_in,
224            tokens_out,
225        });
226    }
227
228    let content = message["content"].as_str().ok_or_else(|| {
229        ModelError::new(format!(
230            "OpenAI response had no message content: {}",
231            truncate(&text, 300)
232        ))
233    })?;
234    parse_response(content, &req.schema, tokens_in, tokens_out).map(Step::Done)
235}
236
237fn ollama(
238    config: &ModelConfig,
239    req: &AnalyzeRequest,
240    transport: &Transport,
241) -> Result<Step, ModelError> {
242    let base = config.endpoint.as_deref().unwrap_or(OLLAMA_BASE);
243    let mut body = json!({
244        "model": config.model,
245        "stream": false,
246        "options": {"num_predict": config.max_output_tokens},
247        "messages": messages(req, &Provider::Ollama),
248    });
249    if req.tools.is_empty() {
250        // Ollama takes the JSON schema directly in `format`.
251        body["format"] = build_json_schema(&req.schema);
252    } else {
253        body["tools"] = tools_json(&req.tools);
254    }
255
256    let headers = [("Content-Type", "application/json".to_string())];
257    let url = format!("{}/api/chat", base.trim_end_matches('/'));
258    let text = transport(&url, &headers, &body)?;
259    let response: Value = serde_json::from_str(&text).map_err(|e| {
260        ModelError::new(format!(
261            "Ollama returned a non-JSON body ({e}): {}",
262            truncate(&text, 300)
263        ))
264    })?;
265
266    let tokens_in = response["prompt_eval_count"].as_u64().unwrap_or(0);
267    let tokens_out = response["eval_count"].as_u64().unwrap_or(0);
268
269    let message = &response["message"];
270    if let Some(call) = message["tool_calls"].get(0) {
271        let name = call["function"]["name"].as_str().unwrap_or_default();
272        // Ollama returns arguments as a JSON object, not a string.
273        let arguments_json = match &call["function"]["arguments"] {
274            Value::String(s) => s.clone(),
275            other => other.to_string(),
276        };
277        return Ok(Step::CallTool {
278            name: name.to_string(),
279            arguments_json,
280            tokens_in,
281            tokens_out,
282        });
283    }
284
285    let content = message["content"].as_str().ok_or_else(|| {
286        ModelError::new(format!(
287            "Ollama response had no message content: {}",
288            truncate(&text, 300)
289        ))
290    })?;
291    parse_response(content, &req.schema, tokens_in, tokens_out).map(Step::Done)
292}
293
294/// OpenAI requires schema names to match `^[a-zA-Z0-9_-]+$`.
295fn sanitize_schema_name(name: &str) -> String {
296    let cleaned: String = name
297        .chars()
298        .map(|c| {
299            if c.is_ascii_alphanumeric() || c == '_' || c == '-' {
300                c
301            } else {
302                '_'
303            }
304        })
305        .collect();
306    if cleaned.is_empty() {
307        "Result".to_string()
308    } else {
309        cleaned
310    }
311}
312
313/// The real network transport, carrying this config's timeout.
314pub(crate) fn transport_for(config: &ModelConfig) -> Box<Transport> {
315    // Zero is how "no timeout" sneaks back in, so it is clamped rather than
316    // honoured -- the same rule the `http` module applies.
317    let timeout = std::time::Duration::from_secs(config.timeout_secs.max(1));
318    Box::new(move |url: &str, headers: &[(&str, String)], body: &Value| {
319        send(url, headers, body, timeout)
320    })
321}
322
323fn send(
324    url: &str,
325    headers: &[(&str, String)],
326    body: &Value,
327    timeout: std::time::Duration,
328) -> Result<String, ModelError> {
329    let agent = ureq::AgentBuilder::new().timeout(timeout).build();
330    let mut request = agent.post(url);
331    for (name, value) in headers {
332        request = request.set(name, value);
333    }
334    match request.send_json(body.clone()) {
335        Ok(response) => response
336            .into_string()
337            .map_err(|e| ModelError::new(format!("could not read response body from {url}: {e}"))),
338        Err(ureq::Error::Status(code, response)) => {
339            let body = response.into_string().unwrap_or_default();
340            Err(ModelError::new(format!(
341                "{url} returned HTTP {code}: {}",
342                truncate(&body, 300)
343            )))
344        }
345        Err(e) => Err(ModelError::new(format!("request to {url} failed: {e}"))),
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352    use crate::{AnalyzeOutcome, FieldType, Schema, SchemaField};
353    use std::cell::RefCell;
354
355    fn schema() -> Schema {
356        Schema {
357            type_name: "Insight".into(),
358            fields: vec![
359                SchemaField {
360                    name: "summary".into(),
361                    field_type: FieldType::Str,
362                    description: None,
363                    pattern: None,
364                },
365                SchemaField {
366                    name: "count".into(),
367                    field_type: FieldType::Int,
368                    description: None,
369                    pattern: None,
370                },
371            ],
372        }
373    }
374
375    fn request() -> AnalyzeRequest {
376        AnalyzeRequest {
377            prompt: "find anomalies".into(),
378            data_json: "{\"rows\":2}".into(),
379            images: Vec::new(),
380            schema: schema(),
381            tools: Vec::new(),
382            tool_history: Vec::new(),
383        }
384    }
385
386    /// Records the outgoing request and replays a canned response body.
387    /// What the recording transport captured: (url, request body).
388    type Captured = std::rc::Rc<RefCell<Option<(String, Value)>>>;
389
390    /// A boxed transport plus the handle that observes what it was sent.
391    type Recorder = (Box<Transport>, Captured);
392
393    /// Build a transport that replays `reply` and remembers the request.
394    fn recording(reply: &'static str) -> Recorder {
395        let seen: Captured = std::rc::Rc::new(RefCell::new(None));
396        let sink = seen.clone();
397        let transport = Box::new(move |url: &str, _h: &[(&str, String)], body: &Value| {
398            *sink.borrow_mut() = Some((url.to_string(), body.clone()));
399            Ok(reply.to_string())
400        });
401        (transport, seen)
402    }
403
404    #[test]
405    fn spec_openai() {
406        let c = parse_model_spec("openai:gpt-4o").unwrap();
407        assert_eq!(c.provider, Provider::OpenAI);
408        assert_eq!(c.model, "gpt-4o");
409        assert_eq!(c.max_output_tokens, 4096);
410        assert_eq!(c.timeout_secs, DEFAULT_TIMEOUT_SECS);
411    }
412
413    #[test]
414    fn spec_local_keeps_tag() {
415        let c = parse_model_spec("local:llama3.1:8b").unwrap();
416        assert_eq!(c.provider, Provider::Ollama);
417        assert_eq!(c.model, "llama3.1:8b");
418    }
419
420    #[test]
421    fn spec_errors() {
422        assert!(parse_model_spec("gpt-4o")
423            .unwrap_err()
424            .message
425            .contains("prefix"));
426        assert!(parse_model_spec("openai:")
427            .unwrap_err()
428            .message
429            .contains("no model name"));
430        assert!(parse_model_spec("groq:x")
431            .unwrap_err()
432            .message
433            .contains("unknown model provider"));
434    }
435
436    #[test]
437    fn openai_request_shape_and_parse() {
438        let reply = r#"{
439            "choices":[{"message":{"content":"{\"summary\":\"ok\",\"count\":2,\"__uncertain__\":\"\"}"}}],
440            "usage":{"prompt_tokens":11,"completion_tokens":7}
441        }"#;
442        let (transport, seen) = recording(reply);
443        let mut config = parse_model_spec("openai:gpt-4o").unwrap();
444        config.api_key = Some("test-key".into());
445
446        let outcome = step_with(&config, &request(), &*transport).unwrap();
447        match outcome {
448            Step::Done(AnalyzeOutcome::Ok {
449                fields_json,
450                tokens_in,
451                tokens_out,
452            }) => {
453                assert_eq!(fields_json["summary"], "ok");
454                assert_eq!(tokens_in, 11);
455                assert_eq!(tokens_out, 7);
456            }
457            other => panic!("expected Ok, got {other:?}"),
458        }
459
460        let (url, body) = seen.borrow().clone().unwrap();
461        assert_eq!(url, "https://api.openai.com/v1/chat/completions");
462        assert_eq!(body["response_format"]["type"], "json_schema");
463        assert_eq!(body["response_format"]["json_schema"]["strict"], true);
464        assert_eq!(body["messages"][0]["role"], "system");
465        assert!(body["messages"][1]["content"]
466            .as_str()
467            .unwrap()
468            .contains("DATA:"));
469    }
470
471    #[test]
472    fn openai_missing_key_is_clear() {
473        // Ensure the env var cannot rescue the call.
474        std::env::remove_var("OPENAI_API_KEY");
475        let (transport, _seen) = recording("{}");
476        let config = parse_model_spec("openai:gpt-4o").unwrap();
477        let err = step_with(&config, &request(), &*transport).unwrap_err();
478        assert!(
479            err.message.contains("OPENAI_API_KEY not set"),
480            "{}",
481            err.message
482        );
483    }
484
485    #[test]
486    fn ollama_request_shape_and_uncertain() {
487        let reply = r#"{
488            "message":{"content":"{\"summary\":\"\",\"count\":0,\"__uncertain__\":\"no revenue column\"}"},
489            "prompt_eval_count":30,"eval_count":9
490        }"#;
491        let (transport, seen) = recording(reply);
492        let config = parse_model_spec("local:llama3.1:8b").unwrap();
493
494        match step_with(&config, &request(), &*transport).unwrap() {
495            Step::Done(AnalyzeOutcome::Uncertain {
496                reason,
497                tokens_in,
498                tokens_out,
499            }) => {
500                assert_eq!(reason, "no revenue column");
501                assert_eq!(tokens_in, 30);
502                assert_eq!(tokens_out, 9);
503            }
504            other => panic!("expected Uncertain, got {other:?}"),
505        }
506
507        let (url, body) = seen.borrow().clone().unwrap();
508        assert_eq!(url, "http://localhost:11434/api/chat");
509        assert_eq!(body["stream"], false);
510        assert_eq!(body["format"]["type"], "object");
511    }
512
513    #[test]
514    fn ollama_endpoint_override() {
515        let reply =
516            r#"{"message":{"content":"{\"summary\":\"a\",\"count\":1,\"__uncertain__\":\"\"}"}}"#;
517        let (transport, seen) = recording(reply);
518        let mut config = parse_model_spec("local:llama3.1:8b").unwrap();
519        config.endpoint = Some("http://box:11434/".into());
520
521        step_with(&config, &request(), &*transport).unwrap();
522        assert_eq!(
523            seen.borrow().clone().unwrap().0,
524            "http://box:11434/api/chat"
525        );
526    }
527
528    /// The same image must arrive in each provider's own shape: OpenAI wants
529    /// typed content parts with a data URL, Ollama wants bare base64 in a
530    /// sibling array. Getting either wrong is a silently text-only request.
531    #[test]
532    fn openai_attaches_images_as_content_parts() {
533        let reply = r#"{
534            "choices":[{"message":{"content":"{\"summary\":\"ok\",\"count\":1,\"__uncertain__\":\"\"}"}}],
535            "usage":{"prompt_tokens":1,"completion_tokens":1}
536        }"#;
537        let (transport, seen) = recording(reply);
538        let mut config = parse_model_spec("openai:gpt-4o").unwrap();
539        config.api_key = Some("test-key".into());
540        let mut req = request();
541        req.images = vec![crate::ImagePart {
542            mime: "image/png".into(),
543            bytes: b"foobar".to_vec(),
544        }];
545
546        step_with(&config, &req, &*transport).unwrap();
547        let (_, body) = seen.borrow().clone().unwrap();
548        let parts = &body["messages"][1]["content"];
549        assert_eq!(parts[0]["type"], "text");
550        assert_eq!(parts[1]["type"], "image_url");
551        assert_eq!(
552            parts[1]["image_url"]["url"],
553            "data:image/png;base64,Zm9vYmFy"
554        );
555    }
556
557    #[test]
558    fn ollama_attaches_images_beside_the_text() {
559        let reply =
560            r#"{"message":{"content":"{\"summary\":\"a\",\"count\":1,\"__uncertain__\":\"\"}"}}"#;
561        let (transport, seen) = recording(reply);
562        let config = parse_model_spec("local:llava:7b").unwrap();
563        let mut req = request();
564        req.images = vec![crate::ImagePart {
565            mime: "image/png".into(),
566            bytes: b"foobar".to_vec(),
567        }];
568
569        step_with(&config, &req, &*transport).unwrap();
570        let (_, body) = seen.borrow().clone().unwrap();
571        let message = &body["messages"][1];
572        assert!(message["content"].as_str().unwrap().contains("DATA:"));
573        assert_eq!(message["images"][0], "Zm9vYmFy");
574    }
575
576    /// A text-only call must keep the plain-string content shape: some
577    /// providers and local models reject the content-parts form outright.
578    #[test]
579    fn no_images_keeps_plain_string_content() {
580        let reply = r#"{
581            "choices":[{"message":{"content":"{\"summary\":\"ok\",\"count\":1,\"__uncertain__\":\"\"}"}}]
582        }"#;
583        let (transport, seen) = recording(reply);
584        let mut config = parse_model_spec("openai:gpt-4o").unwrap();
585        config.api_key = Some("test-key".into());
586
587        step_with(&config, &request(), &*transport).unwrap();
588        let (_, body) = seen.borrow().clone().unwrap();
589        assert!(body["messages"][1]["content"].is_string());
590    }
591
592    #[test]
593    fn schema_name_sanitized() {
594        assert_eq!(sanitize_schema_name("Insight"), "Insight");
595        assert_eq!(sanitize_schema_name("my type!"), "my_type_");
596        assert_eq!(sanitize_schema_name(""), "Result");
597    }
598}
kora_models/provider.rs

kora_models/
provider.rs
