gems 0.1.4

💎 A cli, tui, and sdk for interacting with the Gemini API (WIP)
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use anyhow::Result;
use gems::embed::{BatchEmbeddingBuilder, EmbeddingBuilder};
use gems::messages::{Content, Message};
use gems::models::Model;
use gems::traits::CTrait;
use gems::Client;

#[test]
fn test_embed_build_required() {
    let message = Message::User {
        content: Content::Text("Test input".into()),
        name: None,
    };

    let embed = EmbeddingBuilder::default()
        .model(Model::Embedding)
        .input(message.clone())
        .build()
        .unwrap();

    assert_eq!(embed.model, Model::Embedding);
    assert_eq!(embed.input, message);
}

#[test]
fn test_batch_embed_build() {
    let inputs = vec![
        Message::User {
            content: Content::Text("First".into()),
            name: None,
        },
        Message::User {
            content: Content::Text("Second".into()),
            name: None,
        },
    ];

    let batch = BatchEmbeddingBuilder::default()
        .model(Model::Embedding)
        .input(inputs.clone())
        .build()
        .unwrap();

    assert_eq!(batch.model, Model::Embedding);
    assert_eq!(batch.input, inputs);
}

#[tokio::test]
async fn test_embed_create() -> Result<()> {
    let gemini_client = Client::builder()
        .model(&Model::Embedding.to_string())
        .build()?;

    gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default());

    let input = Message::User {
        content: Content::Text("What is Rust programming?".into()),
        name: None,
    };

    let params = EmbeddingBuilder::default()
        .model(Model::Embedding)
        .input(input)
        .build()?;

    let response = gemini_client.embeddings().create(params).await?;

    let embeddings = response.embedding.clone().unwrap_or_default().values;

    assert!(
        !embeddings.is_empty(),
        "Expected embedding vector in response."
    );

    for &value in &embeddings {
        assert!(
            value.is_finite(),
            "Embedding value should be finite (not NaN or infinite)"
        );
        assert!(
            (-1.0..=1.0).contains(&value),
            "Embedding value out of expected range: {}",
            value
        );
    }

    assert_eq!(embeddings.len(), 768, "Unexpected embedding vector size");

    Ok(())
}

#[tokio::test]
async fn test_embed_batch() -> Result<()> {
    let gemini_client = Client::builder()
        .model(&Model::Embedding.to_string())
        .build()?;

    gemini_client.set_api_key(std::env::var("GEMINI_API_KEY").unwrap_or_default());

    let inputs = vec![
        Message::User {
            content: Content::Text("First embedding".into()),
            name: None,
        },
        Message::User {
            content: Content::Text("Second embedding".into()),
            name: None,
        },
    ];

    let params = BatchEmbeddingBuilder::default()
        .model(Model::Embedding)
        .input(inputs)
        .build()?;

    let response = gemini_client.embeddings().batch(params).await?;

    assert_eq!(
        response.embeddings.unwrap_or_default().len(),
        2,
        "Expected two embeddings in batch response."
    );

    Ok(())
}