atomr-infer-core 0.8.0

Foundation crate for atomr-infer — typed errors, deployment value object, ModelRunner trait, batch primitives.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

//! Default runtime selection — `Deployment::infer_runtime()` (doc §3.2).
//!
//! Operators always override; this is the ergonomic path for the common
//! case where the model name carries the runtime affinity unambiguously
//! (`gpt-4o*` → openai, `claude-*` → anthropic, etc.).

use crate::runtime::RuntimeKind;

/// Map a model name to its default runtime backend.
pub fn infer_runtime(model: &str) -> RuntimeKind {
    let m = model.to_ascii_lowercase();

    // Remote provider model families take precedence. The patterns are
    // intentionally conservative — operators with custom local
    // fine-tunes named `gpt-*` should set `runtime` explicitly.
    if matches_any(
        &m,
        &[
            "gpt-4",
            "gpt-4o",
            "gpt-4-turbo",
            "gpt-3.5",
            "o1-",
            "o3-",
            "chatgpt-",
        ],
    ) {
        return RuntimeKind::OpenAi;
    }
    if m.starts_with("claude-") || m.starts_with("anthropic/") {
        return RuntimeKind::Anthropic;
    }
    if m.starts_with("gemini-") || m.starts_with("google/gemini") {
        return RuntimeKind::Gemini;
    }
    if m.contains("via-litellm") || m.starts_with("litellm/") {
        return RuntimeKind::LiteLlm;
    }

    // Local LLM families: prefer Rust-native runtimes where available.
    if m.contains("mistral") {
        return RuntimeKind::MistralRs;
    }

    // Default for local LLM-shaped names where no Rust-native backend
    // is well established.
    RuntimeKind::Vllm
}

fn matches_any(model: &str, prefixes: &[&str]) -> bool {
    prefixes.iter().any(|p| model.starts_with(p))
}

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

    #[test]
    fn openai_defaults() {
        assert_eq!(infer_runtime("gpt-4o"), RuntimeKind::OpenAi);
        assert_eq!(infer_runtime("gpt-4o-mini"), RuntimeKind::OpenAi);
        assert_eq!(infer_runtime("o1-preview"), RuntimeKind::OpenAi);
    }

    #[test]
    fn anthropic_defaults() {
        assert_eq!(infer_runtime("claude-sonnet-4"), RuntimeKind::Anthropic);
        assert_eq!(infer_runtime("anthropic/claude-3-haiku"), RuntimeKind::Anthropic);
    }

    #[test]
    fn gemini_defaults() {
        assert_eq!(infer_runtime("gemini-2.0-pro"), RuntimeKind::Gemini);
        assert_eq!(infer_runtime("google/gemini-1.5-flash"), RuntimeKind::Gemini);
    }

    #[test]
    fn local_fallthrough_for_unknown() {
        assert_eq!(
            infer_runtime("meta-llama/Llama-3.1-70B-Instruct"),
            RuntimeKind::Vllm
        );
    }

    #[test]
    fn mistral_picks_rust_native() {
        assert_eq!(
            infer_runtime("mistralai/Mistral-7B-Instruct-v0.3"),
            RuntimeKind::MistralRs
        );
    }
}