Skip to main content

streaming/
streaming.rs

1use antigravity_sdk_rust::agent::{Agent, AgentConfig};
2use antigravity_sdk_rust::policy;
3use antigravity_sdk_rust::types::{GeminiConfig, StreamChunk};
4use futures_util::StreamExt;
5use tracing_subscriber::EnvFilter;
6
7#[tokio::main]
8async fn main() -> Result<(), anyhow::Error> {
9    // Initialize tracing subscriber
10    tracing_subscriber::fmt()
11        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()))
12        .init();
13
14    // Load environment variables from .env file if present
15    dotenvy::dotenv().ok();
16
17    let mut config = AgentConfig::default();
18
19    if let Ok(harness_path) = std::env::var("ANTIGRAVITY_HARNESS_PATH") {
20        config.binary_path = Some(harness_path);
21    }
22
23    let mut gemini_config = GeminiConfig::default();
24    if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
25        gemini_config.api_key = Some(api_key);
26    }
27    gemini_config.models.default.name = "gemini-3.5-flash".to_string();
28    config.gemini_config = gemini_config;
29
30    config.policies = Some(vec![policy::allow_all()]);
31
32    let mut agent = Agent::new(config);
33    println!("Starting agent...");
34    agent.start().await?;
35
36    let prompt = "Solve this riddle: I speak without a mouth and hear without ears. I have no body, but I come alive with wind. What am I? Explain your reasoning.";
37    println!("\n  User: {}\n", prompt);
38
39    let conversation = agent.conversation()?;
40    let mut stream = conversation.chat(prompt).await?;
41
42    println!("  Agent (Streaming response):");
43    println!("  -------------------------------------------------------");
44
45    while let Some(chunk_res) = stream.next().await {
46        match chunk_res? {
47            StreamChunk::Thought { text, .. } => {
48                // Print thought chunks (e.g. in grey or wrapped in brackets if desired, or directly)
49                print!("{}", text);
50                std::io::Write::flush(&mut std::io::stdout())?;
51            }
52            StreamChunk::Text { text, .. } => {
53                // Print response text chunks
54                print!("{}", text);
55                std::io::Write::flush(&mut std::io::stdout())?;
56            }
57            StreamChunk::ToolCall(call) => {
58                println!(
59                    "\n[Tool Call Requested: {} with args: {}]",
60                    call.name, call.args
61                );
62            }
63        }
64    }
65    println!("\n  -------------------------------------------------------\n");
66
67    agent.stop().await?;
68    Ok(())
69}