Skip to main content

agentzero_core/common/
local_providers.rs

1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2pub enum LocalProviderType {
3    ChatCompletion,
4    Transcription,
5}
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct LocalProviderMeta {
9    pub id: &'static str,
10    pub default_port: u16,
11    pub default_base_url: &'static str,
12    pub models_endpoint: Option<&'static str>,
13    pub supports_pull: bool,
14    pub provider_type: LocalProviderType,
15}
16
17const LOCAL_PROVIDERS: &[LocalProviderMeta] = &[
18    LocalProviderMeta {
19        id: "ollama",
20        default_port: 11434,
21        default_base_url: "http://localhost:11434",
22        models_endpoint: Some("/api/tags"),
23        supports_pull: true,
24        provider_type: LocalProviderType::ChatCompletion,
25    },
26    LocalProviderMeta {
27        id: "llamacpp",
28        default_port: 8080,
29        default_base_url: "http://localhost:8080",
30        models_endpoint: Some("/v1/models"),
31        supports_pull: false,
32        provider_type: LocalProviderType::ChatCompletion,
33    },
34    LocalProviderMeta {
35        id: "lmstudio",
36        default_port: 1234,
37        default_base_url: "http://localhost:1234",
38        models_endpoint: Some("/v1/models"),
39        supports_pull: false,
40        provider_type: LocalProviderType::ChatCompletion,
41    },
42    LocalProviderMeta {
43        id: "vllm",
44        default_port: 8000,
45        default_base_url: "http://localhost:8000",
46        models_endpoint: Some("/v1/models"),
47        supports_pull: false,
48        provider_type: LocalProviderType::ChatCompletion,
49    },
50    LocalProviderMeta {
51        id: "sglang",
52        default_port: 30000,
53        default_base_url: "http://localhost:30000",
54        models_endpoint: Some("/v1/models"),
55        supports_pull: false,
56        provider_type: LocalProviderType::ChatCompletion,
57    },
58    LocalProviderMeta {
59        id: "osaurus",
60        default_port: 8080,
61        default_base_url: "http://localhost:8080",
62        models_endpoint: Some("/v1/models"),
63        supports_pull: false,
64        provider_type: LocalProviderType::ChatCompletion,
65    },
66    LocalProviderMeta {
67        id: "whispercpp",
68        default_port: 8080,
69        default_base_url: "http://localhost:8080",
70        models_endpoint: None,
71        supports_pull: false,
72        provider_type: LocalProviderType::Transcription,
73    },
74];
75
76pub fn is_local_provider(id: &str) -> bool {
77    let needle = id.trim().to_ascii_lowercase();
78    LOCAL_PROVIDERS.iter().any(|p| p.id == needle)
79}
80
81pub fn local_provider_meta(id: &str) -> Option<&'static LocalProviderMeta> {
82    let needle = id.trim().to_ascii_lowercase();
83    LOCAL_PROVIDERS.iter().find(|p| p.id == needle)
84}
85
86pub fn all_local_providers() -> &'static [LocalProviderMeta] {
87    LOCAL_PROVIDERS
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn is_local_provider_recognises_ollama() {
96        assert!(is_local_provider("ollama"));
97        assert!(is_local_provider("Ollama"));
98        assert!(is_local_provider("  ollama  "));
99    }
100
101    #[test]
102    fn is_local_provider_rejects_cloud_providers() {
103        assert!(!is_local_provider("openrouter"));
104        assert!(!is_local_provider("openai"));
105        assert!(!is_local_provider("anthropic"));
106    }
107
108    #[test]
109    fn local_provider_meta_returns_correct_defaults() {
110        let meta = local_provider_meta("ollama").expect("ollama should be found");
111        assert_eq!(meta.default_port, 11434);
112        assert_eq!(meta.default_base_url, "http://localhost:11434");
113        assert!(meta.supports_pull);
114    }
115
116    #[test]
117    fn local_provider_meta_returns_none_for_unknown() {
118        assert!(local_provider_meta("openai").is_none());
119        assert!(local_provider_meta("").is_none());
120    }
121
122    #[test]
123    fn all_local_providers_contains_expected_count() {
124        assert_eq!(all_local_providers().len(), 7);
125    }
126
127    #[test]
128    fn whispercpp_is_transcription_type() {
129        let meta = local_provider_meta("whispercpp").expect("whispercpp should be found");
130        assert_eq!(meta.provider_type, LocalProviderType::Transcription);
131        assert!(meta.models_endpoint.is_none());
132    }
133}