Advanced Start

use axllm::{ax, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("question:string -> answer:string")?;
    let output = program.forward(&mut client, json!({"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."}))?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{agent_with_options, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;

fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| {
            axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")
        })?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;

    // The handbook can be arbitrarily large. Listing it in `contextFields` keeps
    // it in the agent's runtime so it never inflates the model prompt -- the agent
    // reads it through code, not through tokens.
    let handbook = r#"
# Acme Cloud -- Support Handbook

## Billing
- Invoices are issued on the 1st of each month and are due net-15.
- Plan downgrades take effect at the END of the current billing cycle, not immediately.
- Refunds are issued to the original payment method within 5 business days.

## Access
- Seats can be added by any workspace Owner under Settings -> Members.
- SSO (SAML) is available on Enterprise; SCIM provisioning is Owner-only.

## Incidents
- Status and uptime are published at status.acme.example.
- Sev-1 incidents page the on-call within 5 minutes; updates post every 30 minutes.

## Data
- Exports are available in CSV and JSON from Settings -> Data.
- Deleted workspaces are recoverable for 30 days, then permanently purged.
"#;

    // `with_runtime` attaches the embedded JS engine so the agent loop can run.
    let mut assistant = agent_with_options(
        "question:string, handbook:string -> answer:string, citations:string[] \"Handbook sections the answer relies on\"",
        json!({"contextFields": ["handbook"], "runtime": {"language": "JavaScript"}}),
    )?
    .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;

    let output = assistant.forward_with_options(
        &mut client,
        json!({
            "question": "A customer downgraded their plan today. When does it take effect, and can they get a refund for the current cycle?",
            "handbook": handbook,
        }),
        json!({"max_actor_steps": 12}),
    )?;

    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::{ax, flow, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::env;


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let step = ax("documentText:string -> summaryText:string")?;
    let mut program = axllm::flow("examples.sequentialFlow").execute("step", step).returns(json!({"step": "step"}));
    let output = program.forward(&mut client, json!({"documentText": "Ax gives developers signatures, provider clients, agents, flows, tracing, and optimization."}))?;
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}

use axllm::{ax, AxAIClient, AxResult, OpenAICompatibleClient};
use serde_json::json;
use std::{env, fs};


fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let speech = client.speak(json!({"text": "Ax turns LLM prompts into typed programs.", "voice": "alloy", "format": "mp3"}))?;
    let audio_len = speech["audio"].as_str().unwrap_or("").len();
    println!("{}", serde_json::to_string_pretty(&json!({"format": speech["format"].clone(), "audioBytesBase64": audio_len}))?);
    Ok(())
}

use std::{env, sync::{Arc, Mutex}};

use axllm::{
    AxAIClient, AxBalancer, AxBalancerAdaptiveStrategy, AxBalancerOptions,
    AxInMemoryBalancerStatsStore, AxResult, OpenAICompatibleClient,
};
use serde_json::json;

fn main() -> AxResult<()> {
    let key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".into());
    let clients: Vec<Box<dyn AxAIClient>> = vec![
        Box::new(OpenAICompatibleClient::new(&key, &model)),
        Box::new(OpenAICompatibleClient::new(&key, &model)),
    ];

    let store = Arc::new(AxInMemoryBalancerStatsStore::new());
    let route_keys = ["openai-primary".to_string(), "openai-backup".to_string()];
    let events = Arc::new(Mutex::new(Vec::new()));
    let event_sink = events.clone();
    let strategy = AxBalancerAdaptiveStrategy::new(6_000.0, 0.02)
        .with_expected_tokens(1_200, 300)
        .with_namespace("support-summary-v1")
        .with_store(store)
        .with_route_key(Arc::new(move |_service, index| route_keys[index].clone()))
        .with_slice(Arc::new(|context| if context["options"]["stream"] == true { "streaming".into() } else { "interactive".into() }))
        .on_routing_event(Arc::new(move |event| event_sink.lock().unwrap().push(event["type"].clone())));
    let mut balancer = AxBalancer::from_clients(clients, AxBalancerOptions { strategy: Some(strategy), ..AxBalancerOptions::default() })?;
    let response = balancer.chat(json!({"model": model, "chat_prompt": [{"role": "user", "content": "Summarize why shared routing state matters."}]}))?;
    println!("{}", serde_json::to_string_pretty(&response)?);
    println!("{:?}", events.lock().unwrap());
    Ok(())
}

use axllm::{ax, AxResult, OpenAICompatibleClient, OptimizerEngine};
use serde_json::{json, Value};
use std::env;

struct ExampleOptimizer;
impl OptimizerEngine for ExampleOptimizer {
    fn optimize(&mut self, _request: Value, _evaluator: &mut dyn FnMut(Value) -> AxResult<Value>) -> AxResult<Value> {
        Ok(json!({"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "axgen"}}))
    }
}

fn openai_client() -> AxResult<OpenAICompatibleClient> {
    let api_key = env::var("OPENAI_API_KEY").or_else(|_| env::var("OPENAI_APIKEY")).map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example."))?;
    let model = env::var("AX_OPENAI_MODEL").unwrap_or_else(|_| "gpt-5.4-mini".to_string());
    Ok(OpenAICompatibleClient::new(api_key, model).with_model_config(json!({"temperature": 0})))
}

fn main() -> AxResult<()> {
    let mut client = openai_client()?;
    let mut program = ax("emailText:string -> priority:class \"high, normal, low\", rationale:string")?;
    let baseline = program.forward(&mut client, json!({"emailText": "Production checkout is failing for enterprise customers."}))?;
    let mut optimizer = ExampleOptimizer;
    let artifact = optimizer.optimize(json!({"candidate": "priority"}), &mut |_candidate| Ok(json!({"score": 1.0})))?;
    println!("{}", serde_json::to_string_pretty(&json!({"baseline": baseline, "artifact": artifact}))?);
    Ok(())
}

use axllm::{
    ax, AxExecutionContext, AxMCPClient, AxMCPStreamableHTTPTransport, AxResult,
    OpenAICompatibleClient,
};
use serde_json::json;
use std::{
    env,
    sync::{Arc, Mutex},
};

fn main() -> AxResult<()> {
    let key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY."))?;
    let endpoint = env::var("MCP_URL").map_err(|_| axllm::AxError::runtime("Set MCP_URL."))?;
    let transport = AxMCPStreamableHTTPTransport::new(endpoint, json!({}))?;
    let mcp = Arc::new(Mutex::new(AxMCPClient::new(
        Box::new(transport),
        json!({"namespace":"inventory"}),
    )));
    let context = AxExecutionContext::new(vec![mcp.clone()], vec![])?;
    let catalog = mcp.lock().unwrap().inspect_catalog(false)?;
    println!(
        "MCP catalog: {} tools, {} resources, {} templates",
        catalog.tools.len(),
        catalog.resources.len(),
        catalog.resource_templates.len()
    );
    let mut program = ax("request:string -> answer:string")?.with_execution_context(context)?;
    let mut llm = OpenAICompatibleClient::new(key, "gpt-5.4-mini");
    println!(
        "{}",
        program.forward(&mut llm, json!({"request":"Reindex inventory."}))?
    );
    mcp.lock().unwrap().close()?;
    Ok(())
}

use axllm::runtime::quickjs::QuickJsCodeRuntime;
use axllm::{
    agent_with_options, AxEventRoute, AxEventRuntime, AxEventTarget, AxMCPClient, AxMCPEventSource,
    AxMCPResourceSubscriptionPolicy, AxMCPStreamableHTTPTransport, AxResult,
    OpenAICompatibleClient,
};
use serde_json::{json, Value};
use std::{
    env,
    sync::{Arc, Condvar, Mutex},
    time::{Duration, Instant},
};

fn main() -> AxResult<()> {
    let key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY."))?;
    let endpoint = env::var("AX_MCP_ENDPOINT").map_err(|_| {
        axllm::AxError::runtime("Set AX_MCP_ENDPOINT to a Streamable HTTP MCP server.")
    })?;
    let local = endpoint.starts_with("http://127.0.0.1");
    let transport = AxMCPStreamableHTTPTransport::new(
        endpoint,
        json!({"ssrfProtection":{"requireHttps":!local,"allowLocalhost":local,"allowPrivateNetworks":local}}),
    )?;
    let client = Arc::new(Mutex::new(AxMCPClient::new(
        Box::new(transport),
        json!({"namespace":"inventory"}),
    )));
    let mut llm = OpenAICompatibleClient::new(key, "gpt-5.4-mini");
    let mut agent = agent_with_options(
        "uri:string -> summary:string",
        json!({"runtime":{"language":"JavaScript"}}),
    )?
    .with_runtime(Box::new(QuickJsCodeRuntime::new()))?;
    let completed = Arc::new((Mutex::new(false), Condvar::new()));
    let completed_target = completed.clone();
    let mut target = AxEventTarget::new("inventory-agent", move |input, _| {
        let output = agent.forward(&mut llm, input)?;
        println!("{output}");
        let (lock, changed) = &*completed_target;
        *lock.lock().unwrap() = true;
        changed.notify_all();
        Ok(output)
    });
    target.retry_safety = "idempotent".into();
    target.map_input = Some(Arc::new(|event, _| Ok(json!({"uri":event.data["uri"]}))));
    let mut runtime = AxEventRuntime::new(
        vec![AxEventRoute {
            id: "resource-wake".into(),
            action: "wake".into(),
            r#match: json!({"types":["mcp.resource.updated"]}),
            target_id: Some("inventory-agent".into()),
            require_authenticated: true,
            ordering: "strict".into(),
            debounce_ms: 0,
            instance_key: None,
        }],
        json!({}),
    )?;
    runtime.register_target(target);
    runtime.start()?;
    let runtime = Arc::new(Mutex::new(runtime));
    let mut source = AxMCPEventSource::with_policy(
        client.clone(),
        runtime.clone(),
        "inventory",
        "tenant:demo",
        "authenticated",
        AxMCPResourceSubscriptionPolicy::All,
    );
    source.start()?;
    let deadline = Instant::now() + Duration::from_secs(60);
    loop {
        source.poll();
        if *completed.0.lock().unwrap() {
            break;
        }
        if Instant::now() >= deadline {
            return Err(axllm::AxError::runtime(
                "Timed out waiting for an MCP resource notification",
            ));
        }
        std::thread::sleep(Duration::from_millis(10));
    }
    source.close()?;
    client.lock().unwrap().close()?;
    runtime.lock().unwrap().close()?;
    Ok(())
}

use axllm::{
    ax, AxEventEnvelope, AxEventRoute, AxEventRuntime, AxEventTarget, AxMCPClient,
    AxMCPEventSource, AxMCPStreamableHTTPTransport, AxResult, OpenAICompatibleClient,
};
use serde_json::{json, Map};
use std::{
    env,
    sync::{Arc, Condvar, Mutex},
    time::{Duration, Instant},
};

fn main() -> AxResult<()> {
    let key = env::var("OPENAI_API_KEY")
        .or_else(|_| env::var("OPENAI_APIKEY"))
        .map_err(|_| axllm::AxError::runtime("Set OPENAI_API_KEY."))?;
    let endpoint = env::var("AX_MCP_ENDPOINT").map_err(|_| {
        axllm::AxError::runtime("Set AX_MCP_ENDPOINT to a Streamable HTTP MCP server.")
    })?;
    let local = endpoint.starts_with("http://127.0.0.1");
    let transport = AxMCPStreamableHTTPTransport::new(
        endpoint,
        json!({"ssrfProtection":{"requireHttps":!local,"allowLocalhost":local,"allowPrivateNetworks":local}}),
    )?;
    let client = Arc::new(Mutex::new(AxMCPClient::new(
        Box::new(transport),
        json!({"namespace":"inventory"}),
    )));
    client.lock().unwrap().init()?;
    let task = client
        .lock()
        .unwrap()
        .call_tool("start_reindex", json!({"scope":"all"}))?;
    let task_id = task["task"]["taskId"].as_str().unwrap().to_string();
    let mut llm = OpenAICompatibleClient::new(key, "gpt-5.4-mini");
    let mut flow = axllm::flow("reindex-flow")
        .execute("status", ax("taskId:string -> status:string")?)
        .returns(json!({"status":"status"}));
    let completed = Arc::new((Mutex::new(0usize), Condvar::new()));
    let completed_target = completed.clone();
    let mut target = AxEventTarget::new("reindex-flow", move |input, _| {
        let output = flow.forward(&mut llm, input)?;
        println!("{output}");
        let (lock, changed) = &*completed_target;
        *lock.lock().unwrap() += 1;
        changed.notify_all();
        Ok(output)
    });
    target.retry_safety = "idempotent".into();
    target.wait_for =
        vec![json!({"kind":"mcp.task","value":"taskKey","metadata":{"taskId":task_id}})];
    target.map_input = Some(Arc::new(|event, continuation| {
        Ok(
            json!({"taskId":continuation.map(|value|value.metadata["taskId"].clone()).unwrap_or_else(||event.data["taskId"].clone())}),
        )
    }));
    let routes = vec![
        AxEventRoute {
            id: "task-start".into(),
            action: "wake".into(),
            r#match: json!({"types":["app.task.started"]}),
            target_id: Some("reindex-flow".into()),
            require_authenticated: false,
            ordering: "strict".into(),
            debounce_ms: 0,
            instance_key: None,
        },
        AxEventRoute {
            id: "task-progress".into(),
            action: "observe".into(),
            r#match: json!({"types":["mcp.progress"]}),
            target_id: None,
            require_authenticated: false,
            ordering: "strict".into(),
            debounce_ms: 0,
            instance_key: None,
        },
        AxEventRoute {
            id: "task-resume".into(),
            action: "resume".into(),
            r#match: json!({"types":["mcp.task.status"]}),
            target_id: Some("reindex-flow".into()),
            require_authenticated: false,
            ordering: "strict".into(),
            debounce_ms: 0,
            instance_key: None,
        },
    ];
    let mut runtime_value = AxEventRuntime::new(routes, json!({}))?;
    runtime_value.register_target(target);
    runtime_value.start()?;
    runtime_value.publish(
        AxEventEnvelope {
            specversion: "1.0".into(),
            id: "task-start".into(),
            source: "app://tasks".into(),
            r#type: "app.task.started".into(),
            subject: Some(task_id.clone()),
            data: json!({"taskId":task_id,"taskKey":format!("inventory:{task_id}")}),
            extensions: Map::new(),
            correlation: vec![],
        },
        "tenant:demo",
        "authenticated",
    )?;
    let runtime = Arc::new(Mutex::new(runtime_value));
    let mut source = AxMCPEventSource::new(
        client.clone(),
        runtime.clone(),
        "inventory",
        "tenant:demo",
        "authenticated",
        vec![],
    );
    source.start()?;
    println!("Waiting for terminal MCP task notification {task_id}");
    let deadline = Instant::now() + Duration::from_secs(60);
    loop {
        source.poll();
        if *completed.0.lock().unwrap() >= 2 {
            break;
        }
        if Instant::now() >= deadline {
            return Err(axllm::AxError::runtime(
                "Timed out waiting for the MCP task continuation",
            ));
        }
        std::thread::sleep(Duration::from_millis(10));
    }
    source.close()?;
    client.lock().unwrap().close()?;
    runtime.lock().unwrap().close()?;
    Ok(())
}