Rust Examples

use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use std::env;

#[derive(Debug, Serialize)]
struct QueryRequest {
    query: String,
    store_id: String,
}

#[derive(Debug, Deserialize)]
struct QueryResponse {
    answer: String,
    citations: Vec<Citation>,
    metadata: Metadata,
}

#[derive(Debug, Deserialize)]
struct Citation {
    document: String,
    page: i32,
    excerpt: String,
    relevance: f64,
}

#[derive(Debug, Deserialize)]
struct Metadata {
    query_id: String,
    tokens_used: i32,
    response_time_ms: i32,
}

async fn query(question: &str, store_id: &str) -> Result<QueryResponse, Box<dyn std::error::Error>> {
    let api_key = env::var("BRAINUS_API_KEY")?;

    let client = reqwest::Client::new();

    let mut headers = HeaderMap::new();
    headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
    headers.insert("X-API-Key", HeaderValue::from_str(&api_key)?);

    let request_body = QueryRequest {
        query: question.to_string(),
        store_id: store_id.to_string(),
    };

    let response = client
        .post("https://api.brainus.lk/api/v1/query")
        .headers(headers)
        .json(&request_body)
        .send()
        .await?;

    if !response.status().is_success() {
        let error_text = response.text().await?;
        return Err(format!("API error: {}", error_text).into());
    }

    let result: QueryResponse = response.json().await?;
    Ok(result)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let result = query("What is photosynthesis?", "default").await?;

    println!("Answer: {}", result.answer);
    println!("Citations: {}", result.citations.len());

    Ok(())
}

use thiserror::Error;

#[derive(Error, Debug)]
pub enum BrainUsError {
    #[error("Authentication failed: {0}")]
    AuthError(String),

    #[error("Rate limit exceeded, retry after {retry_after}s")]
    RateLimitError { retry_after: u64 },

    #[error("Validation error: {0}")]
    ValidationError(String),

    #[error("Server error: {0}")]
    ServerError(String),

    #[error("Network error: {0}")]
    NetworkError(#[from] reqwest::Error),
}

pub struct BrainUsClient {
    api_key: String,
    base_url: String,
    client: reqwest::Client,
}

impl BrainUsClient {
    pub fn new(api_key: String) -> Self {
        Self {
            api_key,
            base_url: "https://api.brainus.lk".to_string(),
            client: reqwest::Client::new(),
        }
    }

    pub async fn query(&self, question: &str, store_id: &str) -> Result<QueryResponse, BrainUsError> {
        let url = format!("{}/api/v1/query", self.base_url);

        let request_body = QueryRequest {
            query: question.to_string(),
            store_id: store_id.to_string(),
        };

        let response = self.client
            .post(&url)
            .header("Content-Type", "application/json")
            .header("X-API-Key", &self.api_key)
            .json(&request_body)
            .send()
            .await?;

        match response.status().as_u16() {
            200 => Ok(response.json().await?),
            401 => Err(BrainUsError::AuthError("Invalid API key".to_string())),
            429 => {
                let retry_after = response
                    .headers()
                    .get("Retry-After")
                    .and_then(|h| h.to_str().ok())
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(60);
                Err(BrainUsError::RateLimitError { retry_after })
            }
            400 => {
                let error_text = response.text().await?;
                Err(BrainUsError::ValidationError(error_text))
            }
            500..=599 => {
                let error_text = response.text().await?;
                Err(BrainUsError::ServerError(error_text))
            }
            _ => {
                let error_text = response.text().await?;
                Err(BrainUsError::ServerError(format!("Unexpected error: {}", error_text)))
            }
        }
    }
}

use axum::{
    extract::Json,
    http::StatusCode,
    response::IntoResponse,
    routing::post,
    Router,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
use std::collections::HashMap;

#[derive(Clone)]
struct AppState {
    client: Arc<BrainUsClient>,
    cache: Arc<RwLock<HashMap<String, CachedResponse>>>,
}

#[derive(Clone)]
struct CachedResponse {
    data: QueryResponse,
    timestamp: std::time::Instant,
}

#[derive(Deserialize)]
struct QueryPayload {
    query: String,
    store_id: Option<String>,
}

#[derive(Serialize)]
struct ApiResponse {
    data: QueryResponse,
    cached: bool,
}

async fn query_handler(
    Json(payload): Json<QueryPayload>,
    state: axum::extract::State<AppState>,
) -> Result<Json<ApiResponse>, (StatusCode, String)> {
    let store_id = payload.store_id.unwrap_or_else(|| "default".to_string());
    let cache_key = format!("{}-{}", payload.query, store_id);

    // Check cache
    {
        let cache = state.cache.read().await;
        if let Some(cached) = cache.get(&cache_key) {
            if cached.timestamp.elapsed().as_secs() < 3600 {
                return Ok(Json(ApiResponse {
                    data: cached.data.clone(),
                    cached: true,
                }));
            }
        }
    }

    // Query API
    let result = state
        .client
        .query(&payload.query, &store_id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    // Update cache
    {
        let mut cache = state.cache.write().await;
        cache.insert(
            cache_key,
            CachedResponse {
                data: result.clone(),
                timestamp: std::time::Instant::now(),
            },
        );
    }

    Ok(Json(ApiResponse {
        data: result,
        cached: false,
    }))
}

#[tokio::main]
async fn main() {
    let api_key = std::env::var("BRAINUS_API_KEY").expect("BRAINUS_API_KEY must be set");

    let state = AppState {
        client: Arc::new(BrainUsClient::new(api_key)),
        cache: Arc::new(RwLock::new(HashMap::new())),
    };

    let app = Router::new()
        .route("/api/query", post(query_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();

    println!("Server running on http://127.0.0.1:3000");

    axum::serve(listener, app).await.unwrap();
}

use tokio::task::JoinSet;

async fn query_multiple(
    client: &BrainUsClient,
    queries: Vec<&str>,
    store_id: &str,
) -> Vec<Result<QueryResponse, BrainUsError>> {
    let mut set = JoinSet::new();

    for query in queries {
        let client = client.clone();
        let query = query.to_string();
        let store_id = store_id.to_string();

        set.spawn(async move {
            client.query(&query, &store_id).await
        });
    }

    let mut results = Vec::new();
    while let Some(result) = set.join_next().await {
        results.push(result.unwrap());
    }

    results
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let api_key = std::env::var("BRAINUS_API_KEY")?;
    let client = BrainUsClient::new(api_key);

    let queries = vec![
        "What is photosynthesis?",
        "Explain the water cycle",
        "What causes earthquakes?",
    ];

    let results = query_multiple(&client, queries, "default").await;

    for (i, result) in results.iter().enumerate() {
        match result {
            Ok(response) => println!("Query {}: {}", i + 1, &response.answer[..100]),
            Err(e) => println!("Query {}: Error - {}", i + 1, e),
        }
    }

    Ok(())
}