llm_api_access/tests/
openai.rs

1
2#[cfg(test)]
3mod tests {
4    use crate::openai::{call_gpt, get_embedding};
5    use crate::structs::general::Message;
6
7    #[tokio::test]
8    async fn test_call_gpt() {
9        let user_message = Message {
10            role: "user".to_string(),
11            content: "Hello, can you tell me a joke?".to_string(),
12        };
13        let messages = vec![user_message];
14        let res = call_gpt(messages).await;
15        match res {
16            Ok(response) => assert!(!response.is_empty(), "Response should not be empty"),
17            Err(err) => panic!("Call to OpenAI API failed: {}", err),
18        }
19    }
20
21    #[tokio::test]
22    async fn test_call_gpt_multi_prompt() {
23        let system_message = Message {
24            role: "system".to_string(),
25            content: "You are a helpful coding assistant.".to_string(),
26        };
27        let user_message_1 = Message {
28            role: "user".to_string(),
29            content: "Hello, can you write a python function that reverses a string?".to_string(),
30        };
31        let mut messages = vec![system_message, user_message_1];
32        let res = call_gpt(messages.clone()).await;
33        match res {
34            Ok(response) => {
35                assert!(!response.is_empty(), "Response should not be empty");
36                let user_message_2 = Message {
37                    role: "user".to_string(),
38                    content: "Can you also provide an example of how to use that function?".to_string(),
39                };
40                messages.push(user_message_2);
41                let res = call_gpt(messages).await;
42                match res {
43                    Ok(response) => assert!(!response.is_empty(), "Response should not be empty"),
44                    Err(err) => panic!("Call to OpenAI API failed on second prompt: {}", err),
45                }
46            }
47            Err(err) => panic!("Call to OpenAI API failed on first prompt: {}", err),
48        }
49    }
50
51    #[tokio::test]
52    async fn test_get_embedding() {
53        let input_text = "This is a test sentence.";
54        let res = get_embedding(input_text.to_string(), None).await;
55        match res {
56            Ok(embedding) => {
57                assert!(!embedding.is_empty(), "Embedding should not be empty");
58                println!("Embedding vector length: {}", embedding.len());
59                // Basic sanity check on the embedding vector (length might change with models)
60                assert!(embedding.len() > 100);
61            }
62            Err(err) => panic!("Failed to get embedding: {}", err),
63        }
64    }
65
66    #[tokio::test]
67    async fn test_get_embedding_with_dimensions() {
68        let input_text = "This is another test.";
69        let dimensions: u32 = 64;
70        let res = get_embedding(input_text.to_string(), Some(dimensions)).await;
71        match res {
72            Ok(embedding) => {
73                assert!(!embedding.is_empty(), "Embedding should not be empty");
74                assert_eq!(embedding.len() as u32, dimensions, "Embedding dimension should match requested dimension");
75                println!("Embedding vector length: {}", embedding.len());
76            }
77            Err(err) => panic!("Failed to get embedding with dimensions: {}", err),
78        }
79    }
80}
llm_api_access/tests/openai.rs

llm_api_access/tests/
openai.rs
