1use serde_json::{json, Value};
2use std::env;
3use std::fmt::Debug;
4use std::future::Future;
5use std::sync::OnceLock;
6use thiserror::Error;
7
8#[derive(Error, Debug)]
9pub enum VibesortError {
10 #[error("Environment variable CEREBRAS_API_KEY not found")]
11 MissingApiKey,
12 #[error("HTTP request failed: {0}")]
13 RequestError(#[from] reqwest::Error),
14 #[error("Failed to parse response: {0}")]
15 ParseError(String),
16 #[error("Invalid response format")]
17 InvalidResponse,
18}
19
20static HTTP_CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
21static API_KEY: OnceLock<String> = OnceLock::new();
22
23fn get_http_client() -> &'static reqwest::Client {
24 HTTP_CLIENT.get_or_init(|| reqwest::Client::new())
25}
26
27fn get_api_key() -> Result<&'static str, VibesortError> {
28 if let Some(key) = API_KEY.get() {
29 return Ok(key.as_str());
30 }
31
32 let key = env::var("CEREBRAS_API_KEY").map_err(|_| VibesortError::MissingApiKey)?;
33 let _ = API_KEY.set(key);
34
35 Ok(API_KEY.get().unwrap().as_str())
36}
37
38pub async fn vibesort<T: std::fmt::Display + std::str::FromStr + Clone>(
47 vec: Vec<T>,
48) -> Result<Vec<T>, VibesortError>
49where
50 <T as std::str::FromStr>::Err: std::fmt::Debug,
51{
52 let api_key = get_api_key()?;
53 let client = get_http_client();
54
55 let mut input_str = String::new();
56 for (i, item) in vec.iter().enumerate() {
57 if i > 0 {
58 input_str.push_str(", ");
59 }
60 input_str.push_str(&item.to_string());
61 }
62
63 let payload = json!({
64 "model": "llama-4-maverick-17b-128e-instruct",
65 "stream": false,
66 "max_tokens": 128,
67 "temperature": 0.1,
68 "messages": [
69 {
70 "role": "system",
71 "content": "Return only sorted array in format [1, 2, 3]. No text."
72 },
73 {
74 "role": "user",
75 "content": format!("[{}]", input_str)
76 }
77 ]
78 });
79
80 let response = client
81 .post("https://api.cerebras.ai/v1/chat/completions")
82 .header("Content-Type", "application/json")
83 .header("Authorization", format!("Bearer {}", api_key))
84 .json(&payload)
85 .send()
86 .await?
87 .json::<Value>()
88 .await?;
89
90 let content = response["choices"][0]["message"]["content"]
91 .as_str()
92 .ok_or(VibesortError::InvalidResponse)?
93 .trim();
94
95 let clean_content = content.trim_start_matches('[').trim_end_matches(']');
96
97 let result: Result<Vec<T>, _> = clean_content
98 .split(", ")
99 .map(|s| {
100 s.trim()
101 .parse()
102 .map_err(|_| VibesortError::ParseError(s.to_string()))
103 })
104 .collect();
105
106 result
107}
108
109pub trait VibeSort<T: std::fmt::Display + std::str::FromStr + Clone>
110where
111 <T as std::str::FromStr>::Err: std::fmt::Debug,
112{
113 fn vibesort(self) -> impl Future<Output = Result<Vec<T>, VibesortError>>;
122}
123
124impl<T: std::fmt::Display + std::str::FromStr + Clone> VibeSort<T> for Vec<T>
125where
126 <T as std::str::FromStr>::Err: std::fmt::Debug,
127{
128 async fn vibesort(self) -> Result<Vec<T>, VibesortError> {
129 vibesort(self).await
130 }
131}