♊

【Rust】Gemini APIを使ってみた

に公開
Rust
Gemini
tech

はじめに

ターミナル上でGeminiと会話するCLIを作りました。

クレート

[dependencies]
tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.11", features = ["json"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
inquire = "0.6"
dotenv = "0.15.0"

できたもの

ソースコード
main.rs
use dotenv::dotenv;
use inquire::Text;
use reqwest::Client;
use serde_json::json;
use std::env;

// エンドポイント
const GEMINI_API_URL: &str =
    "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent";

fn fetch_api_key() -> String {
    // .envファイルからAPIキーを取得
    match dotenv() {
        Ok(_) => {
            if let Ok(key) = env::var("GEMINI_API_KEY") {
                return key;
            }
        }
        _ => {}
    }

    // コマンドライン引数があればそれをAPIキーとして使用
    if env::args().len() > 1 {
        return env::args().collect::<Vec<String>>()[1].clone();
    }

    // .envファイルが無く、引数も指定されていなければAPIキーの入力を求める
    Text::new("Enter your Gemini API key:").prompt().unwrap()
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let api_key = fetch_api_key();
    let client = Client::new();

    // 会話の履歴
    let mut messages = vec![];

    println!("Start Gemini Chat ('q' to quit):");

    // qが入力されるまでループ
    loop {
        // ユーザ入力
        let user_input = Text::new(" You:").prompt().unwrap();

        if user_input.trim().eq_ignore_ascii_case("q") {
            break;
        }

        // 会話履歴にユーザ入力を追加
        messages.push(json!({"role": "user", "parts": vec![json!({"text": user_input})]}));

        // リクエストのJSONを作成
        let json = json!({"contents": messages});

        // リクエストを送信
        let response = client
            .post(&format!("{}?key={}", GEMINI_API_URL, api_key))
            .header(reqwest::header::CONTENT_TYPE, "application/json")
            .body(json.to_string())
            .send()
            .await?
            .json::<serde_json::Value>()
            .await?;

        if let Some(reply) = response["candidates"][0]["content"]["parts"][0]["text"].as_str() {
            // レスポンスを出力
            println!("Gemini:\t{}", reply.trim().replace("\n", "\n\t"));
            // 会話履歴にレスポンスを追加
            messages.push(json!({"role": "model", "parts": vec![json!({"text": reply})]}));
        } else if let Some(reply) = response["error"]["message"].as_str() {
            // エラーが帰ってきた場合はエラーを出力して終了
            println!("Gemini:\t{}", reply.trim());
            break;
        } else {
            // その他の場合も終了
            println!("Gemini: Unexpected response");
            println!("Response:\n{:#?}", response);
            break;
        }
    }

    println!("Chat Terminated.");
    Ok(())
}

解説・補足

APIキーの管理

とりあえず動かすだけであればAPIキーはべた書きでも良いかもしれませんが、GitHub等を用いる場合はPushされるとまずいので、環境変数や外部ファイルに格納すべきです。

今回の場合はdotenvを用いてmain.rsと同じディレクトリに.envというファイルを用意し、そこに

GEMINI_API_KEY=...

という風に保存しています。もちろん.gitignoreには.envを追加しておきます。

継続会話

文脈を保ったまま会話をするには履歴を保存する必要があります。今回の場合はmessagesにユーザの入力とその返答を追加していくことで実現しています。

JSONのParse

きっちりと作るのであれば、リクエスト・レスポンスに対応した構造体を用意した方が良いと思います。
ただ、それをするとコードがかなり長くなってしまうので、今回は省きました。

詳細な仕様についてはこちらを参照ください。

おわり

参考

Discussion