axllm 24.0.8

Generated Ax runtime library
Documentation
use axllm::{AxAIClient, AxError, AxResult, AxTransport, OpenAICompatibleClient};
use serde_json::{json, Value};
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Clone)]
struct Script(Arc<Mutex<State>>);
struct State {
    responses: VecDeque<Value>,
    requests: Vec<Value>,
}
impl Script {
    fn new(responses: Vec<Value>) -> Self {
        Self(Arc::new(Mutex::new(State {
            responses: responses.into(),
            requests: vec![],
        })))
    }
    fn methods(&self) -> Vec<String> {
        self.0
            .lock()
            .unwrap()
            .requests
            .iter()
            .map(|value| value["method"].as_str().unwrap().to_string())
            .collect()
    }
}
impl AxTransport for Script {
    fn send(&mut self, request: Value) -> AxResult<Value> {
        let mut state = self.0.lock().unwrap();
        state.requests.push(request);
        state
            .responses
            .pop_front()
            .ok_or_else(|| AxError::runtime("script exhausted"))
    }
}
fn now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis() as u64
}
fn success(text: &str) -> Value {
    json!({"status":200,"json":{"candidates":[{"content":{"parts":[{"text":text}]},"finishReason":"STOP"}]}})
}
fn cache(name: &str, seconds: u64) -> Value {
    json!({"status":200,"json":{"name":name,"expireTime":now()+seconds*1000}})
}
fn failure(status: u16, message: &str) -> Value {
    json!({"status":status,"json":{"error":{"message":message}}})
}
fn service(script: Script) -> OpenAICompatibleClient {
    OpenAICompatibleClient::new("gemini-key", "gemini-3.5-flash")
        .with_profile("google-gemini")
        .with_options(
            json!({"contextCache":{"minTokens":0,"ttlSeconds":3600,"refreshWindowSeconds":300}}),
        )
        .with_transport(script)
}
fn main() -> AxResult<()> {
    let request = json!({"chat_prompt":[{"role":"system","content":"stable context"},{"role":"user","content":"answer briefly"}]});
    let recovery = Script::new(vec![
        cache("cachedContents/cache-1", 3600),
        failure(400, "cachedContent is invalid"),
        success("uncached recovery"),
    ]);
    service(recovery.clone()).chat(request.clone())?;
    assert_eq!(recovery.methods(), vec!["POST", "POST", "POST"]);
    let refresh = Script::new(vec![
        cache("cachedContents/old", 1),
        success("old"),
        failure(500, "refresh failed"),
        cache("cachedContents/new", 3600),
        success("recreated"),
    ]);
    let mut refresh_client = service(refresh.clone());
    refresh_client.chat(request.clone())?;
    refresh_client.chat(request.clone())?;
    assert_eq!(
        refresh.methods(),
        vec!["POST", "POST", "PATCH", "POST", "POST"]
    );
    let fallback = Script::new(vec![
        cache("cachedContents/old", 1),
        success("old"),
        failure(500, "refresh failed"),
        failure(500, "recreate failed"),
        success("uncached fallback"),
    ]);
    let mut fallback_client = service(fallback.clone());
    fallback_client.chat(request.clone())?;
    fallback_client.chat(request)?;
    assert_eq!(
        fallback.methods(),
        vec!["POST", "POST", "PATCH", "POST", "POST"]
    );
    println!("rust-context-cache-recovery-ok");
    Ok(())
}