tea-model 0.1.0

Provider-neutral model port for tea-rs
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
93
94
95
96
97
98
99
100
101
102

use std::str::FromStr;

use futures_util::{StreamExt, stream};
use tea_model::{
    BoxModelStream, ModelCancellation, ModelCapabilities, ModelCompletion, ModelDisplayName,
    ModelEvent, ModelProvider, ModelRequest, ModelResponseInfo, ModelSpec, ProviderId, Utf8Delta,
};
use tea_protocol::{
    CanonicalMessage, ContentBlock, MessageId, ModelId, ProtocolTimestamp, StopReason, TokenCount,
};

#[derive(Debug)]
struct InMemoryProvider {
    provider_id: ProviderId,
    models: Vec<ModelSpec>,
}

impl ModelProvider for InMemoryProvider {
    fn provider_id(&self) -> &ProviderId {
        &self.provider_id
    }

    fn models(&self) -> &[ModelSpec] {
        &self.models
    }

    fn stream(&self, _request: ModelRequest, _cancellation: ModelCancellation) -> BoxModelStream {
        Box::pin(stream::iter([
            ModelEvent::Started(ModelResponseInfo::new()),
            ModelEvent::TextDelta(Utf8Delta::new("hello").unwrap()),
            ModelEvent::Completed(ModelCompletion::new(StopReason::Completed).unwrap()),
        ]))
    }
}

fn provider() -> InMemoryProvider {
    InMemoryProvider {
        provider_id: ProviderId::from_str("memory").unwrap(),
        models: vec![
            ModelSpec::new(
                ModelId::from_str("memory/test").unwrap(),
                ProviderId::from_str("memory").unwrap(),
                ModelDisplayName::from_str("Memory Test").unwrap(),
                TokenCount::new(8_000).unwrap(),
                TokenCount::new(2_000).unwrap(),
                ModelCapabilities::text(),
            )
            .unwrap(),
        ],
    }
}

fn request() -> ModelRequest {
    ModelRequest::new(
        ModelId::from_str("memory/test").unwrap(),
        vec![
            CanonicalMessage::user(
                MessageId::from_str("0195a0b1-5e3d-73de-b461-0aa7aa000004").unwrap(),
                vec![ContentBlock::text("hello").unwrap()],
                ProtocolTimestamp::from_str("2026-07-23T09:30:12.123Z").unwrap(),
            )
            .unwrap(),
        ],
    )
    .unwrap()
}

#[tokio::test(flavor = "current_thread")]
async fn provider_is_object_safe_and_returns_normalized_stream() {
    let provider: Box<dyn ModelProvider> = Box::new(provider());
    assert_eq!(provider.provider_id().as_str(), "memory");
    assert_eq!(provider.models().len(), 1);
    assert_eq!(
        provider
            .model(&ModelId::from_str("memory/test").unwrap())
            .unwrap()
            .display_name()
            .as_str(),
        "Memory Test"
    );
    assert!(
        provider
            .model(&ModelId::from_str("memory/missing").unwrap())
            .is_none()
    );

    let events = provider
        .stream(request(), ModelCancellation::new())
        .collect::<Vec<_>>()
        .await;
    assert_eq!(events.len(), 3);
    assert!(matches!(events.first(), Some(ModelEvent::Started(_))));
    assert!(matches!(events.last(), Some(ModelEvent::Completed(_))));
}

#[test]
fn provider_and_stream_contracts_are_send_sync_as_required() {
    fn assert_provider<T: ModelProvider>() {}
    fn assert_stream<T: tea_model::ModelStream>() {}
    assert_provider::<InMemoryProvider>();
    assert_stream::<futures_util::stream::Iter<std::array::IntoIter<ModelEvent, 0>>>();
}