open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-stack/examples/seeded_cascade_with_adapters.rs

//! End-to-end composition example: walks a query through the
//! full free-seeded cascade using the adapter crates shipped in
//! [f168037] / [b8367d0] / [7328d8c] / [03e7571] / [03c063c] /
//! [6594a54].
//!
//! Runs entirely offline — every adapter accepts a caller-injected
//! transport, and this example wires deterministic mock transports
//! so the binary compiles + runs with no API keys, no network
//! traffic, and no model-weight downloads. The point is to
//! demonstrate the COMPOSITION shape, not to call real services.
//!
//! For a production wire-up, swap the mock transports for ones
//! that wrap `reqwest::blocking` (sync transports) or `reqwest`
//! (async transports) and forward HTTP requests / responses
//! verbatim — the adapter shapes don't change.
//!
//! Walk:
//!   1. L3.5 fresh retrieval — `ArxivProvider.search(query, 5)`
//!   2. L3.5 fetch — `JinaReader.read(top_hit)`
//!   3. L2 embedding — `OpenAiEmbedBackend.embed_batch([query, …docs])`
//!   4. L2.5 rerank — `CohereRerankBackend.rerank(query, docs, top=2)`
//!   5. L3 LLM — `OllamaBackend.complete(prompt with top doc + query)`
//!   6. Receipt — print joule + token accounting
//!
//! Run: `cargo run -p jouleclaw-stack --example seeded_cascade_with_adapters`

use async_trait::async_trait;
use jouleclaw_embed_api::{
    HttpTransport as EmbedHttpTransport, OpenAiEmbedBackend,
};
use jouleclaw_fresh::providers::{
    ArxivProvider, BraveProvider, HttpMethod, HttpTransportExt, JinaReader,
};
use jouleclaw_fresh::{FreshError, Transport, TransportResponse};
use jouleclaw_llm::ollama::{HttpTransport as OllamaHttpTransport, OllamaBackend};
use jouleclaw_llm::{LlmBackend, LlmRequest};
use jouleclaw_rerank_api::{
    CohereRerankBackend, HttpTransport as RerankHttpTransport,
};
use std::sync::Arc;

// ─────────────────────────────────────────────────────────────────────
// Deterministic mock transports — one per adapter trait shape
// ─────────────────────────────────────────────────────────────────────

/// Mock for `jouleclaw_fresh::Transport` (GET-only, no headers).
/// Returns canned content based on URL substring matches.
struct MockGetTransport;

#[async_trait]
impl Transport for MockGetTransport {
    async fn get(&self, url: &str) -> Result<TransportResponse, FreshError> {
        let body = if url.contains("export.arxiv.org") {
            // arXiv Atom feed with 3 entries.
            arxiv_atom_feed()
        } else if url.contains("r.jina.ai") {
            // Jina Reader returns clean markdown.
            "# Tokio task scheduler\n\nTokio uses a work-stealing scheduler with M:N \
             green-thread mapping. Each worker thread has a local run queue …".into()
        } else {
            "{}".into()
        };
        Ok(TransportResponse {
            body: body.into_bytes(),
            content_type: Some("application/atom+xml".into()),
        })
    }
}

/// Mock for `HttpTransportExt` (paid search providers).
struct MockExtTransport;

#[async_trait]
impl HttpTransportExt for MockExtTransport {
    async fn request(
        &self,
        _method: HttpMethod,
        _url: &str,
        _headers: &[(&str, &str)],
        _body: Option<&[u8]>,
    ) -> Result<TransportResponse, FreshError> {
        let body = serde_json::json!({
            "web": {
                "results": [
                    {"title": "Brave hit 1", "url": "https://example.com/1", "description": "Tokio's runtime …"},
                    {"title": "Brave hit 2", "url": "https://example.com/2", "description": "Work-stealing scheduler …"}
                ]
            }
        })
        .to_string();
        Ok(TransportResponse {
            body: body.into_bytes(),
            content_type: Some("application/json".into()),
        })
    }
}

/// Mock for `jouleclaw_embed_api::HttpTransport`. Returns
/// deterministic 4-d vectors keyed by input position so the
/// example can show "different inputs → different vectors."
struct MockEmbedTransport;

impl EmbedHttpTransport for MockEmbedTransport {
    fn post_with_bearer(
        &self,
        _url: &str,
        _api_key: &str,
        body: &[u8],
    ) -> Result<Vec<u8>, String> {
        let req: serde_json::Value =
            serde_json::from_slice(body).map_err(|e| e.to_string())?;
        let inputs = req
            .get("input")
            .and_then(|i| i.as_array())
            .cloned()
            .unwrap_or_default();
        let mut data = Vec::with_capacity(inputs.len());
        for (i, _) in inputs.iter().enumerate() {
            let base = (i as f32) * 0.1;
            data.push(serde_json::json!({
                "object": "embedding",
                "index": i,
                "embedding": [base, base + 0.05, base + 0.10, base + 0.15],
            }));
        }
        let resp = serde_json::json!({
            "data": data,
            "usage": {"prompt_tokens": 12, "total_tokens": 12},
        });
        Ok(resp.to_string().into_bytes())
    }
}

/// Mock for `jouleclaw_rerank_api::HttpTransport`. Returns scores
/// that are monotonically decreasing by index so the example
/// shows a clear ranking.
struct MockRerankTransport;

impl RerankHttpTransport for MockRerankTransport {
    fn post_with_bearer(
        &self,
        _url: &str,
        _api_key: &str,
        body: &[u8],
    ) -> Result<Vec<u8>, String> {
        let req: serde_json::Value =
            serde_json::from_slice(body).map_err(|e| e.to_string())?;
        let docs = req
            .get("documents")
            .and_then(|d| d.as_array())
            .cloned()
            .unwrap_or_default();
        let mut results = Vec::with_capacity(docs.len());
        for (i, _) in docs.iter().enumerate() {
            results.push(serde_json::json!({
                "index": i,
                "relevance_score": 1.0 - (i as f64) * 0.2,
            }));
        }
        Ok(serde_json::json!({"results": results}).to_string().into_bytes())
    }
}

/// Mock for `jouleclaw_llm::ollama::HttpTransport`.
struct MockOllamaTransport;

impl OllamaHttpTransport for MockOllamaTransport {
    fn post_json(&self, _url: &str, _body: &[u8]) -> Result<Vec<u8>, String> {
        let body = serde_json::json!({
            "model": "qwen2.5:7b",
            "response": "Tokio's scheduler is work-stealing with M:N threading. Each worker keeps a local run queue; idle workers steal tasks from busy peers.",
            "done": true,
            "done_reason": "stop",
            "prompt_eval_count": 240,
            "eval_count": 38,
        });
        Ok(body.to_string().into_bytes())
    }
}

// ─────────────────────────────────────────────────────────────────────
// Tiny synthetic arXiv Atom feed (kept inline so the example
// stays self-contained)
// ─────────────────────────────────────────────────────────────────────

fn arxiv_atom_feed() -> String {
    r#"<?xml version="1.0" encoding="UTF-8"?>
<feed xmlns="http://www.w3.org/2005/Atom">
  <entry>
    <id>http://arxiv.org/abs/2024.tokio.scheduler</id>
    <title>Work-Stealing Task Scheduling in Async Runtimes</title>
    <summary>Survey of M:N green-thread schedulers with focus on Tokio's design.</summary>
  </entry>
  <entry>
    <id>http://arxiv.org/abs/2024.runtime.parking</id>
    <title>Parking Strategies for Idle Workers in Async Runtimes</title>
    <summary>Comparison of LIFO parking with work-stealing for tail latency.</summary>
  </entry>
  <entry>
    <id>http://arxiv.org/abs/2024.scheduler.fairness</id>
    <title>Cooperative Scheduling and Fairness Under Adversarial Workloads</title>
    <summary>Studies cooperative-yield mechanisms in pre-emptive vs cooperative schedulers.</summary>
  </entry>
</feed>"#
        .to_string()
}

// ─────────────────────────────────────────────────────────────────────
// The walk
// ─────────────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let query = "How does Tokio's task scheduler work?";
    println!("Q: {query}\n");

    // Joule accounting — every step adds to the running total.
    // The numbers here are pinned-typical-cost stand-ins; a real
    // run reads them from the energy backends per cascade walk.
    let mut joules_uj: u64 = 0;

    // ── Step 1: L3.5 fresh search (arXiv) ───────────────────────
    let get_tx: Arc<dyn Transport> = Arc::new(MockGetTransport);
    let arxiv = ArxivProvider::new(get_tx.clone());
    use jouleclaw_fresh::SearchProvider;
    let hits = arxiv.search(query, 5).await?;
    joules_uj += 1_500; // mock typical cost for an arXiv search
    println!("L3.5 fresh (arxiv) — {} hits", hits.len());
    for (i, h) in hits.iter().take(3).enumerate() {
        println!("  [{i}] {} — {}", h.title, h.url);
    }

    // ── Step 2: L3.5 fetch (Jina Reader) ────────────────────────
    let reader = JinaReader::new(get_tx.clone());
    let top_url = hits
        .first()
        .map(|h| h.url.as_str())
        .unwrap_or("https://example.com/none");
    let claim = reader.read(top_url, 7).await?;
    joules_uj += 800;
    println!(
        "\nL3.5 fetch (jina reader) — fetched {} bytes from {}",
        claim.content.len(),
        claim.url
    );

    // ── Step 3: L2 embed (OpenAI Embeddings — mock transport) ───
    let embed_tx: Arc<dyn EmbedHttpTransport> = Arc::new(MockEmbedTransport);
    let embedder =
        OpenAiEmbedBackend::new(embed_tx, "sk-mock", "text-embedding-3-small");
    let docs: Vec<&str> = hits.iter().take(3).map(|h| h.title.as_str()).collect();
    let mut inputs: Vec<&str> = vec![query];
    inputs.extend(docs.iter().copied());
    let vectors = embedder.embed_batch(&inputs)?;
    joules_uj += 2_000;
    println!(
        "\nL2 embed (openai-shape) — {} vectors of dim {}",
        vectors.len(),
        vectors[0].len()
    );

    // ── Step 4: L2.5 rerank (Cohere — mock transport) ───────────
    let rerank_tx: Arc<dyn RerankHttpTransport> = Arc::new(MockRerankTransport);
    let reranker =
        CohereRerankBackend::new(rerank_tx, "co-mock", "rerank-english-v3.0");
    let scored = reranker.rerank(query, &docs, Some(2))?;
    joules_uj += 1_200;
    println!("\nL2.5 rerank (cohere v3) — top {}:", scored.len());
    for (i, (idx, score)) in scored.iter().enumerate() {
        println!("  [{i}] doc[{idx}] = `{}` (score {:.3})", docs[*idx], score);
    }

    // ── Step 5: L3 LLM (Ollama — mock transport) ────────────────
    let ollama_tx: Arc<dyn OllamaHttpTransport> = Arc::new(MockOllamaTransport);
    let llm = OllamaBackend::new(ollama_tx, "qwen2.5:7b");
    let top_doc_idx = scored.first().map(|(i, _)| *i).unwrap_or(0);
    let top_doc = docs[top_doc_idx];
    let prompt = format!(
        "Use this evidence: \"{}\". Answer the question: {}",
        top_doc, query
    );
    let response = llm.complete(&LlmRequest::from_prompt(prompt, 200))?;
    joules_uj += 6_500; // mock typical cost for a small local LLM call
    println!(
        "\nL3 LLM (ollama qwen2.5:7b) — {} prompt tokens, {} output tokens",
        response.input_tokens, response.output_tokens
    );

    // ── Step 6: receipt ─────────────────────────────────────────
    println!("\n──── ANSWER ────");
    println!("{}", response.text);
    println!("──── RECEIPT ───");
    println!("  walk: fresh → fetch → embed → rerank → llm");
    println!(
        "  joules:        {:.3} J (~{:.3} J/answer)",
        (joules_uj as f64) / 1_000_000.0,
        (joules_uj as f64) / 1_000_000.0
    );
    println!(
        "  tokens:        prompt={} output={}",
        response.input_tokens, response.output_tokens
    );
    println!(
        "  joules/token:  {:.3} mJ",
        (joules_uj as f64) / 1_000.0
            / (response.input_tokens + response.output_tokens).max(1) as f64
    );

    // ── (optional) paid-path example: identical walk shape with
    //    a real key would just swap the providers. Demonstrate
    //    that the Brave path also wires under the same trait.
    let ext_tx: Arc<dyn HttpTransportExt> = Arc::new(MockExtTransport);
    let brave = BraveProvider::new(ext_tx, "BSA-mock-key");
    let brave_hits = brave.search(query, 3).await?;
    println!(
        "\n(also wired — brave search returned {} hits using HttpTransportExt)",
        brave_hits.len()
    );

    Ok(())
}