ai_bot/
ai_bot.rs

1//! An AI chatbot. Replies to every message with an LLM completion, showing a
2//! "typing…" indicator while it thinks and threading its answer to your message.
3//! Keeps a short rolling history per conversation, so it follows context.
4//!
5//! Works in DMs and Community channels alike — the SDK hides the difference, so
6//! the entire "make it an AI" part is just the `ask_llm` call below.
7//!
8//! Point it at any OpenAI-compatible chat-completions endpoint:
9//! ```sh
10//! OPENAI_API_KEY=sk-...  \
11//! VECTOR_NSEC=nsec1...   \
12//! cargo run -p vector-sdk --example ai_bot
13//! # optional: OPENAI_BASE_URL (default https://api.openai.com/v1), OPENAI_MODEL (default gpt-4o-mini)
14//! ```
15
16use std::collections::HashMap;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19
20use serde_json::{json, Value};
21use vector_sdk::vector_core::net::build_http_client;
22use vector_sdk::VectorBot;
23
24const SYSTEM_PROMPT: &str =
25    "You are a friendly, concise assistant chatting inside Vector, a private messenger. Keep replies short.";
26const HISTORY_LIMIT: usize = 16; // recent user/assistant messages kept per conversation
27
28#[tokio::main]
29async fn main() -> vector_sdk::Result<()> {
30    let nsec = std::env::var("VECTOR_NSEC").expect("set VECTOR_NSEC to your bot's nsec");
31    std::env::var("OPENAI_API_KEY").expect("set OPENAI_API_KEY for the LLM endpoint");
32
33    let bot = VectorBot::builder().nsec(nsec).build().await?;
34    println!("AI bot online as {}", bot.npub());
35
36    // Per-conversation rolling history, shared across handler invocations.
37    let history: Arc<Mutex<HashMap<String, Vec<Value>>>> = Arc::new(Mutex::new(HashMap::new()));
38
39    bot.on_message(move |_bot, msg| {
40        let history = history.clone();
41        async move {
42            if msg.is_mine() || msg.text().trim().is_empty() {
43                return;
44            }
45
46            // Let the user see we're working while the model generates.
47            let _ = msg.channel().typing().await;
48
49            // Prompt = system + recent history for this conversation + the new turn.
50            let key = msg.chat_id.clone();
51            let mut messages = vec![json!({ "role": "system", "content": SYSTEM_PROMPT })];
52            if let Some(prior) = history.lock().unwrap().get(&key) {
53                messages.extend(prior.iter().cloned());
54            }
55            messages.push(json!({ "role": "user", "content": msg.text() }));
56
57            match ask_llm(&messages).await {
58                Ok(answer) => {
59                    let _ = msg.reply(&answer).await;
60                    // Remember both sides, trimmed to the last HISTORY_LIMIT messages.
61                    let mut store = history.lock().unwrap();
62                    let convo = store.entry(key).or_default();
63                    convo.push(json!({ "role": "user", "content": msg.text() }));
64                    convo.push(json!({ "role": "assistant", "content": answer }));
65                    let overflow = convo.len().saturating_sub(HISTORY_LIMIT);
66                    convo.drain(0..overflow);
67                }
68                Err(e) => {
69                    let _ = msg.reply(&format!("(LLM error: {e})")).await;
70                }
71            }
72        }
73    })
74    .await?;
75
76    Ok(())
77}
78
79/// Call an OpenAI-compatible `/chat/completions` endpoint and return the reply text.
80async fn ask_llm(messages: &[Value]) -> Result<String, String> {
81    let base =
82        std::env::var("OPENAI_BASE_URL").unwrap_or_else(|_| "https://api.openai.com/v1".into());
83    let model = std::env::var("OPENAI_MODEL").unwrap_or_else(|_| "gpt-4o-mini".into());
84    let key = std::env::var("OPENAI_API_KEY").map_err(|_| "OPENAI_API_KEY not set".to_string())?;
85
86    // The SDK hands you a hardened HTTP client — same one Vector uses everywhere — so even your
87    // LLM calls inherit the Tor failsafe and SSRF guards. No need to add `reqwest` yourself.
88    let body: Value = build_http_client(Duration::from_secs(60))?
89        .post(format!("{base}/chat/completions"))
90        .bearer_auth(key)
91        .json(&json!({ "model": model, "messages": messages }))
92        .send()
93        .await
94        .map_err(|e| e.to_string())?
95        .error_for_status()
96        .map_err(|e| e.to_string())?
97        .json()
98        .await
99        .map_err(|e| e.to_string())?;
100
101    body["choices"][0]["message"]["content"]
102        .as_str()
103        .map(|s| s.trim().to_string())
104        .ok_or_else(|| "no completion in response".to_string())
105}
ai_bot/ai_bot.rs

ai_bot/
ai_bot.rs
