awaken 0.6.0

Rust AI agent backend: build capabilities once, tune behavior live, serve every client
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Multi-turn conversation using AgentRuntime with persistent thread storage.
//!
//! Run with BigModel GLM:
//!   LLM_BASE_URL=https://open.bigmodel.cn/api/paas/v4/ LLM_API_KEY=<key> LLM_MODEL=GLM-4.7-Flash cargo run --example multi_turn
//!
//! Also supports: OPENAI_API_KEY + OPENAI_MODEL, or ANTHROPIC_API_KEY, etc.

use async_trait::async_trait;
use awaken::contract::event::AgentEvent;
use awaken::contract::event_sink::EventSink;
use awaken::contract::message::Message;
use awaken::engine::GenaiExecutor;
use awaken::loop_runner::build_agent_env;
use awaken::registry::ResolvedAgent;
use awaken::stores::InMemoryStore;
use awaken::*;
use std::sync::Arc;

struct ConsoleSink;

#[async_trait]
impl EventSink for ConsoleSink {
    async fn emit(&self, event: AgentEvent) {
        if let AgentEvent::TextDelta { delta } = &event {
            print!("{delta}");
        }
    }
}

struct SimpleResolver {
    agent: ResolvedAgent,
}

impl AgentResolver for SimpleResolver {
    fn resolve(&self, _agent_id: &str) -> Result<ResolvedAgent, awaken::RuntimeError> {
        let mut agent = self.agent.clone();
        agent.env = build_agent_env(&[], &agent)?;
        Ok(agent)
    }
}

fn build_llm_executor() -> GenaiExecutor {
    if let (Ok(mut base_url), Ok(api_key)) =
        (std::env::var("LLM_BASE_URL"), std::env::var("LLM_API_KEY"))
    {
        use genai::adapter::AdapterKind;
        use genai::resolver::{AuthData, Endpoint};
        use genai::{ModelIden, ServiceTarget};

        if !base_url.ends_with('/') {
            base_url.push('/');
        }
        let adapter = match std::env::var("LLM_ADAPTER").as_deref() {
            Ok("anthropic") => AdapterKind::Anthropic,
            _ => AdapterKind::OpenAI,
        };
        let client = genai::Client::builder()
            .with_service_target_resolver_fn(move |st: ServiceTarget| {
                Ok(ServiceTarget {
                    endpoint: Endpoint::from_owned(base_url.clone()),
                    auth: AuthData::from_single(api_key.clone()),
                    model: ModelIden::new(adapter, st.model.model_name),
                })
            })
            .build();
        GenaiExecutor::with_client(client)
    } else {
        GenaiExecutor::new()
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let model = std::env::var("LLM_MODEL")
        .or_else(|_| std::env::var("OPENAI_MODEL"))
        .unwrap_or_else(|_| "gpt-4o-mini".into());
    let llm = Arc::new(build_llm_executor());

    let agent = ResolvedAgent::new(
        "chat",
        &model,
        "You are a helpful assistant. Be concise. Remember what the user tells you.",
        llm,
    );
    let resolver = Arc::new(SimpleResolver {
        agent: agent.clone(),
    });

    let store = Arc::new(InMemoryStore::new());
    let runtime = AgentRuntime::new(resolver)
        .with_commit_coordinator(awaken::stores::MemoryCommitCoordinator::wrap(store));

    let thread_id = "conversation-1";

    let turns = [
        "My name is Alice. Remember it.",
        "What is the capital of France?",
        "What is my name?",
    ];

    for (i, user_msg) in turns.iter().enumerate() {
        print!(
            "\n[Turn {}] User: {}\n[Turn {}] Assistant: ",
            i + 1,
            user_msg,
            i + 1
        );

        let request =
            RunActivation::new(thread_id, vec![Message::user(*user_msg)]).with_agent_id("chat");

        match runtime.run(request, Arc::new(ConsoleSink)).await {
            Ok(result) => {
                print!(" [{:?}]", result.termination);
            }
            Err(e) => {
                print!(" [ERROR: {e}]");
            }
        }
    }

    print!("\n\n=== Multi-turn test complete ===\n");
}