tiny-agent 0.3.0

//! tiny-agent 极简示例。
//!
//! 用内存存储 + 自带的 `HostSandbox` + OpenAI 兼容 provider 跑一轮 agent loop。
//! 这是最小可运行形态:存储 trait 用 `HashMap` 现写,生产中请替换为你的持久化后端。
//!
//! 运行:
//! ```bash
//! LLM_API_KEY=sk-xxx cargo run -- "用 bash 执行 echo hello,然后用一句话总结"
//! ```
//! 默认指向 DeepSeek,可用 `LLM_BASE_URL` / `LLM_MODEL` 覆盖。
//!
//! 注意:`HostSandbox` 直接在**本机**执行 bash 命令,示例仅供本地试跑。

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

use async_trait::async_trait;
use tiny_agent::{
    Agent, AgentRunTimeBuilder, StorageError,
    checkpoint::CheckpointStorage,
    providers::openai_compatible::OpenAiCompatibleProvider,
    sandbox::HostSandbox,
    tools::{ToolRegistry, register_default_tools},
    transcript::{Transcript, TranscriptStorage},
};
use tokio_util::sync::CancellationToken;

/// 最小内存存储:同时实现 checkpoint 与 transcript 两个 trait。
/// `Agent` 没有 `Clone`,但可序列化,所以 checkpoint 按 JSON 字符串存。
#[derive(Default)]
struct MemoryStore {
    checkpoints: Mutex<HashMap<String, String>>,
    transcripts: Mutex<HashMap<String, Transcript>>,
}

#[async_trait]
impl CheckpointStorage for MemoryStore {
    async fn save_checkpoint(&self, session_id: &str, agent: &Agent) -> Result<(), StorageError> {
        let json = serde_json::to_string(agent).map_err(|e| StorageError::serde(e.to_string()))?;
        self.checkpoints
            .lock()
            .unwrap()
            .insert(session_id.to_string(), json);
        Ok(())
    }

    async fn get_checkpoint(&self, session_id: &str) -> Result<Option<Agent>, StorageError> {
        let json = self.checkpoints.lock().unwrap().get(session_id).cloned();
        match json {
            Some(json) => {
                let agent =
                    serde_json::from_str(&json).map_err(|e| StorageError::serde(e.to_string()))?;
                Ok(Some(agent))
            }
            None => Ok(None),
        }
    }

    async fn delete_checkpoint(&self, session_id: &str) -> Result<(), StorageError> {
        self.checkpoints.lock().unwrap().remove(session_id);
        Ok(())
    }
}

#[async_trait]
impl TranscriptStorage for MemoryStore {
    async fn save_transcript(
        &self,
        session_id: &str,
        transcript: &Transcript,
    ) -> Result<(), StorageError> {
        self.transcripts
            .lock()
            .unwrap()
            .insert(session_id.to_string(), transcript.clone());
        Ok(())
    }

    async fn get_transcript(&self, session_id: &str) -> Result<Option<Transcript>, StorageError> {
        Ok(self.transcripts.lock().unwrap().get(session_id).cloned())
    }

    async fn remove_transcript(&self, session_id: &str) -> Result<(), StorageError> {
        self.transcripts.lock().unwrap().remove(session_id);
        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let api_key = std::env::var("LLM_API_KEY")
        .or_else(|_| std::env::var("DEEPSEEK_API_KEY"))
        .expect("请设置 LLM_API_KEY 环境变量");
    let base_url =
        std::env::var("LLM_BASE_URL").unwrap_or_else(|_| "https://api.deepseek.com/v1".to_string());
    let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "deepseek-chat".to_string());

    let provider = Arc::new(OpenAiCompatibleProvider::new(api_key, base_url));
    let store = Arc::new(MemoryStore::default());

    let mut tools = ToolRegistry::new();
    register_default_tools(&mut tools); // read_file / write_file / bash / askFollowupQuestion

    let runtime = AgentRunTimeBuilder::new()
        .provider(provider)
        .model(model)
        .checkpoint_storage(store.clone())
        .transcript_storage(store.clone())
        .sandbox(Arc::new(HostSandbox::new()))
        .tools(tools)
        .system_prompt("你是 tiny-agent 的最小示例助手,请简洁作答。")
        .build()?;

    let session_id = runtime.create_session().await?;
    let prompt = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "用 bash 执行 echo hello,然后用一句话总结。".to_string());

    let ready = runtime.submit(&session_id, prompt).await?;
    let outcome = runtime.run(ready, CancellationToken::new()).await?;

    match outcome {
        Agent::Success(state) => {
            println!("\n=== 完成 ===\n{}", state.final_resp.unwrap_or_default());
        }
        Agent::WaitingForUser(_, interaction) => {
            println!(
                "\n=== 等待用户 ===\n{}: {}",
                interaction.kind, interaction.payload
            );
        }
        Agent::Fail(_, failure) => {
            println!("\n=== 失败 ===\n{}: {}", failure.kind, failure.message);
        }
        other => println!("\n=== 状态 ===\n{other:?}"),
    }

    Ok(())
}