代码—Rust 实现微信公众号 OpenAI 机器人

Cargo.toml 配置：

[package]
name = "robot"
version = "0.1.0"
authors = ["xwb"]
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
actix-web = "4.*"
actix = "0.*"
actix-rt = "2.*"
awc = { version = "3.*", features = ["openssl"] }
serde = {version = "1.*", features = ["derive"]}
serde_json = "1.*"
quick-xml = "0.*"
hex = "0.*"
sha1 = "0.*"
redis = "0.*"


[profile.release]
incremental = true

main.rs 代码：

use actix_web::{App, HttpMessage, HttpRequest, HttpResponse, HttpServer};
use actix_web::web::Query;

use awc::ClientBuilder;
use serde::Deserialize;
use serde_json;
use serde_json::value::Value;
use redis::{Commands, RedisResult};
use quick_xml::events::{Event, BytesCData, BytesEnd, BytesStart};
use quick_xml::reader::Reader;
use quick_xml::writer::Writer;
use sha1::{Sha1, Digest};
use hex;

use std::{env, str};
use std::time::SystemTime;
use std::collections::HashMap;
use std::io::Cursor;
use std::time::Duration;


const OPENAI_COMPLETION_URL: &str = "https://api.openai.com/v1/completions";
const OPENAI_MODEL: &str = "text-davinci-003";
const OPENAI_MAX_TOKENS: u32 = 1024;
const OPENAI_IMAGE_URL: &str = "https://api.openai.com/v1/images/generations";
const OPENAI_IMAGE_SIZE: &str = "512x512";

const RETRY_MSG: &str = "请给小木一点时间思考，然后再问小木一次";


fn get_env_var(k: &str, default: &str) -> String {
    match env::var(k) {
        Ok(v) => v,
        Err(_) => default.to_owned()
    }
}


fn sha1(v: &str) -> String {
    let mut sha = Sha1::new();
    sha.update(v.as_bytes());
    hex::encode(sha.finalize().as_slice())
}


fn parse_xml(body: &str) -> HashMap {
    let mut reader = Reader::from_str(body);
    let mut buf = Vec::new();
    reader.trim_text(true);

    let mut map: HashMap = HashMap::new();
    let mut name = String::from("");
    loop {
        match reader.read_event_into(&mut buf) {
            Ok(Event::Start(e)) => {
                name = str::from_utf8(e.name().as_ref()).unwrap().to_owned()
            }
            Ok(Event::Text(e)) => {
                let v = e.unescape().unwrap().into_owned();
                map.insert(name.clone(), v);
            }
            Ok(Event::CData(e)) => {
                let v = str::from_utf8(e.into_inner().into_owned().as_slice()).unwrap().to_owned();
                map.insert(name.clone(), v);
            }
            Ok(Event::Eof) => break,
            _ => (),
        }
    }
    map
}


async fn request_openai(content: &str) -> String {
    let openai_api_key = get_env_var("OPENAI_API_KEY", "");
    let client = ClientBuilder::new()
        .timeout(Duration::from_secs(100))
        .bearer_auth(openai_api_key)
        .finish();
    if content.starts_with("#") {
        let data = serde_json::json!({
            "prompt": content,
            "size": OPENAI_IMAGE_SIZE,
        });
        let mut rsp = client.post(OPENAI_IMAGE_URL)
            .send_json(&data)
            .await
            .unwrap();
        let body = rsp.body().await.unwrap();
        let rst: HashMap = serde_json::from_str(
            str::from_utf8(body.as_ref()).unwrap()
        ).unwrap();
        rst["data"][0]["url"].as_str().unwrap().to_owned()
    } else {
        let data = serde_json::json!({
            "model": OPENAI_MODEL,
            "prompt": content,
            "max_tokens": OPENAI_MAX_TOKENS,
        });
        let mut rsp = client.post(OPENAI_COMPLETION_URL)
            .send_json(&data)
            .await
            .unwrap();
        let body = rsp.body().await.unwrap();
        let rst: HashMap = serde_json::from_str(
            str::from_utf8(body.as_ref()).unwrap()
        ).unwrap();
        rst["choices"][0]["text"].as_str().unwrap().to_owned()
    }
}



#[derive(Deserialize)]
struct VerifyTokenInfo {
    signature: String,
    timestamp: String,
    nonce: String,
    echostr: String
}


#[actix_web::get("/")]
async fn verify_token(info: Query) -> HttpResponse {
    let mut l = std::vec![
        get_env_var("WX_TOKEN", ""),
        info.timestamp.clone(),
        info.nonce.clone()
    ];
    l.sort();

    let msg = l.join("");
    if info.signature.eq(sha1(msg.as_str()).as_str()) {
        HttpResponse::Ok().body(info.echostr.clone())
    } else {
        HttpResponse::Forbidden().body("")
    }
}


#[actix_web::post("/")]
async fn chat(req: HttpRequest, body: String) -> HttpResponse {
    if !req.content_type().to_lowercase().ends_with("xml") {
        return HttpResponse::Ok().body("");
    }

    let map = parse_xml(body.as_str());
    if !map.get("MsgType").unwrap().eq("text") {
        return HttpResponse::Ok().body("");
    }

    let content = map.get("Content").unwrap();
    let mut cache = redis::Client::open("redis://127.0.0.1/0")
        .unwrap()
        .get_connection()
        .unwrap();
    let k = sha1(content);
    let v: RedisResult = cache.get(&k);
    let content_rsp: String;
    match v {
        Ok(d) => {
            content_rsp = d;
        }
        Err(_) => {
            let _: () = cache.set_ex(&k, RETRY_MSG, 30).unwrap();
            content_rsp = request_openai(content).await;
            let _: () = cache.set_ex(&k, content_rsp.as_str(), 600).unwrap();
        }
    }

    let create_time = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs()
        .to_string();

    let mut writer = Writer::new(Cursor::new(Vec::new()));
    assert!(writer.write_event(Event::Start(BytesStart::new("xml"))).is_ok());
    for (k, v) in [
        ("FromUserName", map.get("ToUserName").unwrap().as_str()),
        ("ToUserName", map.get("FromUserName").unwrap().as_str()),
        ("CreateTime", create_time.as_str()),
        ("MsgType", "text"),
        ("Content", content_rsp.as_str())
    ] {
        assert!(writer.write_event(Event::Start(BytesStart::new(k))).is_ok());
        assert!(writer.write_event(Event::CData(BytesCData::new(v))).is_ok());
        assert!(writer.write_event(Event::End(BytesEnd::new(k))).is_ok());
    }
    assert!(writer.write_event(Event::End(BytesEnd::new("xml"))).is_ok());

    let body_rsp = str::from_utf8(writer.into_inner().into_inner().as_slice()).unwrap().to_owned();
    println!("{}", body_rsp);
    return HttpResponse::Ok().body(body_rsp);
}



#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let workers = get_env_var("WEB_WORKERS", "4").parse::().unwrap();
    let port = get_env_var("WEB_PORT", "8000").parse::().unwrap();
    HttpServer::new(|| {
        App::new()
            .service(verify_token)
            .service(chat)
    })
    .bind(("0.0.0.0", port))?
    .workers(workers)
    .run()
    .await
}

start.sh 脚本：

export WX_TOKEN=
export WX_APP_ID=
export WX_APP_SECRET=

export OPENAI_API_KEY=

export WEB_PORT=8000
export WEB_WORKERS=8


pkill -f robot
nohup  ./target/release/robot &

占用资源非常少，适合小配置机器。

测试