Skip to main content

streaming/
streaming.rs

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