rigspark_core/
coverage.rs

1use crate::{
2    catalog::{Catalog, require},
3    sizing::ValidationError,
4};
5use regex::Regex;
6use serde::Serialize;
7use std::collections::BTreeSet;
8
9pub const INVENTORY_URL: &str =
10    "https://raw.githubusercontent.com/ollama/ollama/main/integration/reg_library_test.go";
11pub const MAX_INVENTORY_BYTES: usize = 1024 * 1024;
12
13#[derive(Serialize)]
14#[serde(rename_all = "camelCase")]
15pub struct Coverage {
16    pub upstream_count: usize,
17    pub covered_count: usize,
18    pub missing: Vec<String>,
19}
20
21fn regex(pattern: &str) -> Result<Regex, ValidationError> {
22    Regex::new(pattern).map_err(|error| ValidationError(error.to_string()))
23}
24
25pub fn parse_inventory(source: &str) -> Result<Vec<String>, ValidationError> {
26    require(
27        source.len() <= MAX_INVENTORY_BYTES,
28        "Ollama library inventory response is oversized",
29    )?;
30    let declaration = "var libraryModels = []string{";
31    let start = source
32        .find(declaration)
33        .ok_or_else(|| ValidationError("Ollama library inventory declaration is missing".into()))?
34        + declaration.len();
35    let end = source[start..]
36        .find("\n}")
37        .ok_or_else(|| ValidationError("Ollama library inventory terminator is missing".into()))?
38        + start;
39    let body = regex(r"(?s)/\*.*?\*/")?.replace_all(&source[start..end], "");
40    let body = regex(r"(?m)//.*$")?.replace_all(&body, "");
41    let entries = regex(r#"(?m)^\s*"([a-z0-9][a-z0-9._-]*)",\s*$"#)?;
42    let models: Vec<String> = entries
43        .captures_iter(&body)
44        .map(|capture| capture[1].into())
45        .collect();
46    require(
47        !models.is_empty(),
48        "Ollama library inventory contains no active models",
49    )?;
50    sorted(&models)
51}
52
53fn weight(byte: u8) -> u16 {
54    match byte {
55        b'_' => 0,
56        b'-' => 1,
57        b'.' => 2,
58        _ => u16::from(byte) + 3,
59    }
60}
61
62fn sorted(models: &[String]) -> Result<Vec<String>, ValidationError> {
63    require(models.len() <= 10000, "inventory exceeds 10000 models")?;
64    for model in models {
65        require(
66            !model.is_empty()
67                && model.len() <= 256
68                && model.as_bytes()[0].is_ascii_alphanumeric()
69                && model.bytes().all(|byte| {
70                    byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"._-".contains(&byte)
71                }),
72            "Ollama library inventory contains an invalid model name",
73        )?;
74    }
75    let mut models: Vec<_> = models
76        .iter()
77        .cloned()
78        .collect::<BTreeSet<_>>()
79        .into_iter()
80        .collect();
81    models.sort_by(|left, right| left.bytes().map(weight).cmp(right.bytes().map(weight)));
82    Ok(models)
83}
84
85fn official_repository(reference: &str) -> Option<&str> {
86    let reference = reference.trim();
87    let path = match reference.rsplit_once(':') {
88        Some((path, tag)) if !path.is_empty() && !tag.contains('/') => path,
89        _ => reference,
90    };
91    if let Some(path) = path.strip_prefix("library/") {
92        Some(path)
93    } else if path.contains('/') {
94        None
95    } else {
96        Some(path)
97    }
98}
99
100fn lineage(repository: &str) -> &str {
101    let length = repository
102        .bytes()
103        .take_while(u8::is_ascii_lowercase)
104        .count();
105    if length == 0 {
106        repository
107    } else {
108        &repository[..length]
109    }
110}
111
112pub fn monitored(catalog: &Catalog, upstream: &[String]) -> Result<Vec<String>, ValidationError> {
113    let vendors: BTreeSet<_> = catalog
114        .models
115        .iter()
116        .filter_map(|model| model.source.ollama.as_deref())
117        .filter_map(official_repository)
118        .map(lineage)
119        .collect();
120    Ok(sorted(upstream)?
121        .into_iter()
122        .filter(|name| vendors.contains(lineage(name)))
123        .collect())
124}
125
126pub fn evaluate(catalog: &Catalog, upstream: &[String]) -> Result<Coverage, ValidationError> {
127    let upstream = sorted(upstream)?;
128    let covered: BTreeSet<_> = catalog
129        .models
130        .iter()
131        .filter_map(|model| model.source.ollama.as_deref())
132        .filter_map(official_repository)
133        .collect();
134    let missing: Vec<_> = upstream
135        .iter()
136        .filter(|name| !covered.contains(name.as_str()))
137        .cloned()
138        .collect();
139    Ok(Coverage {
140        upstream_count: upstream.len(),
141        covered_count: upstream.len() - missing.len(),
142        missing,
143    })
144}
rigspark_core/coverage.rs

rigspark_core/
coverage.rs
