Research assistant

use std::sync::Arc;
use async_trait::async_trait;
use cognis::prelude::*;
use cognis_core::schemars::{self, JsonSchema};
use cognis_llm::tools::{Tool, ToolInput, ToolOutput};
use serde::Deserialize;
use serde_json::json;

#[derive(Debug, Deserialize, JsonSchema)]
struct SearchArgs {
    /// The query to search for.
    query: String,
    /// Maximum number of results.
    #[serde(default)]
    max_results: Option<u32>,
}

struct SearchTool;

#[async_trait]
impl Tool for SearchTool {
    fn name(&self) -> &str { "search" }
    fn description(&self) -> &str {
        "Search the web. Returns a list of results with title, url, snippet."
    }
    fn args_schema(&self) -> Option<serde_json::Value> {
        Some(serde_json::to_value(schemars::schema_for!(SearchArgs)).unwrap())
    }
    async fn _run(&self, input: ToolInput) -> Result<ToolOutput> {
        let args: SearchArgs = serde_json::from_value(input.into_json())?;
        // Plug your real provider here.
        let results = vec![
            json!({"title": "Example", "url": "https://example.com", "snippet": format!("Result for: {}", args.query)}),
        ];
        Ok(ToolOutput::Content(json!({"results": results})))
    }
}

use cognis::prelude::*;
use cognis::{AgentBuilder, MultiAgentOrchestrator, Sequential};
use cognis_llm::Client;

#[tokio::main]
async fn main() -> Result<()> {
    let new_client = || Client::from_env();

    let planner = AgentBuilder::new()
        .with_llm(new_client()?)
        .with_system_prompt(
            "You are a research planner. Given a question, return 3 to 5 \
             numbered subquestions whose answers, taken together, would \
             produce a thorough report. Output ONLY the numbered list.",
        )
        .with_max_iterations(2)
        .build()?;

    let researcher = AgentBuilder::new()
        .with_llm(new_client()?)
        .with_tool(Arc::new(SearchTool))
        .with_system_prompt(
            "You are a researcher. You receive a numbered plan. For each \
             subquestion, call the `search` tool, then summarize the most \
             relevant results in one short paragraph. Cite URLs inline.",
        )
        .with_max_iterations(8)
        .build()?;

    let writer = AgentBuilder::new()
        .with_llm(new_client()?)
        .with_system_prompt(
            "You are a writer. You receive a research dossier with cited \
             findings. Produce a one-page report: a 2-sentence summary, \
             then a section per subtopic, then a list of sources. Keep \
             citation links inline like (https://example.com).",
        )
        .with_max_iterations(2)
        .build()?;

    let orch = MultiAgentOrchestrator::new(Sequential)
        .add("planner", planner)
        .add("researcher", researcher)
        .add("writer", writer);

    let resp = orch.run("What changed between Rust 2021 and Rust 2024?").await?;
    println!("{}", resp.content);
    Ok(())
}

COGNIS_PROVIDER=openai COGNIS_OPENAI_API_KEY=sk-... cargo run