Skip to main content

vibesort/
lib.rs

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
38/// Sorts Vec<T> blazingly fast with enterprise AI powered precision
39///
40/// ```rust
41/// let my_vec = vec![69, 41, 67, 9 + 10, 21, 420, 4];
42///
43/// // [4, 19, 21, 41, 67, 69, 420] (hopefully)
44/// let sorted_vec = vibesort(my_vec).await;
45/// ```
46pub 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    /// Sorts Vec<T> blazingly fast with enterprise AI powered precision
114    ///
115    /// ```rust
116    /// let my_vec = vec![69, 41, 67, 9 + 10, 21, 420, 4];
117    ///
118    /// // [4, 19, 21, 41, 67, 69, 420] (hopefully)
119    /// let sorted_vec = my_vec.vibe_sort().await;
120    /// ```
121    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}