hehe 0.0.1

A modular AI Agent framework in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

//! Streaming chat example
//!
//! Run with: cargo run --example streaming

use futures::StreamExt;
use hehe_agent::{Agent, AgentEvent};
use hehe_llm::OpenAiProvider;
use std::sync::Arc;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let api_key = std::env::var("OPENAI_API_KEY")
        .expect("OPENAI_API_KEY environment variable required");

    let llm = Arc::new(OpenAiProvider::new(api_key));

    let agent = Agent::builder()
        .system_prompt("You are a helpful assistant.")
        .model("gpt-4o")
        .llm(llm)
        .build()?;

    let session = agent.create_session();

    println!("Streaming response:\n");

    let mut stream = agent.chat_stream(&session, "Write a haiku about Rust programming");

    while let Some(event) = stream.next().await {
        match event {
            AgentEvent::TextDelta { delta } => {
                print!("{}", delta);
            }
            AgentEvent::TextComplete { text } => {
                println!("\n\n[Complete: {} chars]", text.len());
            }
            AgentEvent::ToolUseStart { name, .. } => {
                println!("\n[Using tool: {}]", name);
            }
            AgentEvent::ToolUseEnd { output, is_error, .. } => {
                if is_error {
                    println!("[Tool error: {}]", output);
                } else {
                    println!("[Tool result: {}]", output);
                }
            }
            AgentEvent::Error { message } => {
                eprintln!("\n[Error: {}]", message);
            }
            AgentEvent::MessageEnd { .. } => {
                println!("\n[Done]");
                break;
            }
            _ => {}
        }
    }

    Ok(())
}