rig-model-meta 0.1.0

Provider-agnostic model metadata (context window, capabilities, quantization) and probe traits for Rig agents.
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

//! Integration coverage for the bundled static catalog.

#![allow(clippy::unwrap_used, clippy::panic, clippy::indexing_slicing)]

use rig_model_meta::{Capability, ModelMetaProbe, StaticProbe};

#[tokio::test]
async fn openai_catalog_has_gpt_family() {
    let probe = StaticProbe::builtin();
    for model in [
        "gpt-4o",
        "gpt-4o-mini",
        "gpt-4-turbo",
        "gpt-3.5-turbo",
        "o1",
        "o1-mini",
    ] {
        let desc = probe
            .describe(model)
            .await
            .unwrap()
            .unwrap_or_else(|| panic!("missing {model} from openai catalog"));
        assert!(
            desc.context_window.is_some(),
            "{model} missing context window"
        );
        assert!(
            desc.has_capability(Capability::Completion)
                || desc.has_capability(Capability::Embedding)
        );
    }
}

#[tokio::test]
async fn anthropic_catalog_has_claude_family() {
    let probe = StaticProbe::builtin();
    for model in [
        "claude-3-5-sonnet-latest",
        "claude-3-5-sonnet-20241022",
        "claude-3-5-haiku-latest",
        "claude-3-opus-latest",
        "claude-3-haiku-20240307",
    ] {
        let desc = probe
            .describe(model)
            .await
            .unwrap()
            .unwrap_or_else(|| panic!("missing {model} from anthropic catalog"));
        assert_eq!(desc.context_window, Some(200_000));
        assert!(desc.has_capability(Capability::Vision));
    }
}

#[tokio::test]
async fn extend_from_custom_json() {
    let mut probe = StaticProbe::default();
    let json = r#"[{"model":"custom-7b","context_window":32768,"capabilities":["completion"]}]"#;
    probe.extend_from_json_str("custom-provider", json).unwrap();
    let desc = probe.describe("custom-7b").await.unwrap().unwrap();
    assert_eq!(desc.context_window, Some(32_768));
    assert_eq!(desc.provider.as_str(), "custom-provider");
}