Rustで学術論文からテキストを抽出する #18 応用編 - arXiv論文収集システム構築⑥ OpenAI API

Summary

• OpenAIのAPIラッパーを実装しました
  • プログラムの構成はSemantic ScholarのAPIラッパーとほとんど同じ
  • 生成AIが含まれるライブラリでは出力が安定しないため単体テストによる品質管理が困難

crate	GitHub
rsrpp	rsrpp
rsrpp-cli	rsrpp
arxiv-tools	rs-arxiv-tools
ss-tools	rs-ss-tools
keyword-tools	keywords
openai-tools	rs-openai-tools

前回までのあらすじ

前回：Rustで学術論文からテキストを抽出する #17

前回までで，①arXivからのスクレイピング，②Semantic Scholarからの論文情報取得，③キーワード抽出，の実装が完了しました．
④はrsrppで対応するので，残りは⑤と⑥です．

今回実装すること

arXiv論文収集システムでは，集めた論文の情報を要約するために生成AIを用います．ChatGPTが登場したばかりの2022年末時点ならいざ知らず，今は高性能なLLMが割とそこら辺に転がっている時代になってしまいました．
細かい差異はあるかもしれませんが，現時点ではそこまで高度な要約を求めているわけではないので，とりあえずOpenAIのAPIを使うことにします．
今後情報抽出の精度改善を実施するときになったら必要に応じて利用するサービスを変更しようと思います．

というわけで，今回はまずOpenAIのAPIを手軽に叩けるようにAPIラッパーを実装していきます．

完成系は以下のようなイメージ．

let mut openai = OpenAI::new();
let messages = vec![
    Message::new("user".to_string(), "Hi there!".to_string())
];

openai
    .model_id("gpt-4o-mini")
    .messages(messages)
    .temperature(1.0);

let response: Response = openai.chat().unwrap();
println!("{}", &response.choices[0].message.content);
// Hello! How can I assist you today?

このシリーズ，後半はAPIのラッパーを作ってばかりな気がする．

OpenAI APIラッパーの実装

OpenAIのAPIにはたくさんの種類がありますが，今回欲しいのはText GenerationのAPIです．OpenAIではChat Completionと呼ばれています．

APIの使い方はこちらの公式ページを参照．

Rustの公式クレートはまだない？と思うので，これまでと同じく自分が使いやすいように実装してしまおうと思います．

利用するのはcurlによるこんなタイプのコマンド．

curl https://api.openai.com/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $OPENAI_API_KEY" \
  -d '{
    "model": "gpt-4o-2024-08-06",
    "messages": [
      {
        "role": "system",
        "content": "You extract email addresses into JSON data."
      },
      {
        "role": "user",
        "content": "Feeling stuck? Send a message to [e-mail]."
      }
    ],
    "response_format": {
      "type": "json_schema",
      "json_schema": {
        "name": "email_schema",
        "schema": {
            "type": "object",
            "properties": {
                "email": {
                    "description": "The email address that appears in the input",
                    "type": "string"
                }
            },
            "additionalProperties": false
        }
      }
    }
  }'

Bodyの実装

重要な部分はbodyに設定されているJSONです．ここで使用するモデルやプロンプト，レスポンスのフォーマットなどを指定できます．

Chat Completionで指定できるパラメータの一覧がこちらのページで見つかります．

今回の実装では以下のパラメータを使えるようにしておこうと思います．
実際に使うのはmodelとmessages以外ではtemperatureくらいな気もしますが．

パラメータ名	型	必須	デフォルト値	説明
messages	array	Required	-	会話を構成するメッセージのリスト。モデルによってサポートされるメッセージの種類（テキスト、画像、音声など）が異なります。
model	string	Required	-	使用するモデルのID。Chat APIで動作するモデルについては互換性表を参照してください。
store	boolean or null	Optional	false	このチャット完了リクエストの出力を保存して、モデル蒸留や評価製品で使用するかどうか。
frequency_penalty	number or null	Optional	0	-2.0から2.0の範囲で指定可能。正の値は、これまでのテキスト内で既存頻度に基づいて新しいトークンをペナルティ化し、同じ行をそのまま繰り返す可能性を減少させます。
logit_bias	map	Optional	null	特定のトークンが補完に現れる可能性を修正します。トークンIDとバイアス値（-100から100）をマッピングするJSONオブジェクトを受け入れます。
logprobs	boolean or null	Optional	false	出力トークンのログ確率を返すかどうか。trueの場合、各出力トークンのログ確率がメッセージ内容に含まれます。
max_completion_tokens	integer or null	Optional	-	補完で生成できるトークン数の上限（可視出力トークンと推論トークンを含む）。
n	integer or null	Optional	1	各入力メッセージに対して生成するチャット補完選択肢の数。生成されたすべての選択肢で発生したトークン数に基づいて課金されます。
modalities	array or null	Optional	["text"]	このリクエストでモデルが生成する出力タイプ（例: テキスト、音声）。
presence_penalty	number or null	Optional	0	-2.0から2.0の範囲で指定可能。正の値はこれまでのテキスト内に現れるかどうかに基づいて新しいトークンをペナルティ化し、新しい話題について話す可能性を増加させます。
response_format	object	Optional	-	モデルが出力すべき形式を指定します（例: JSONモードや構造化出力）。
temperature	number or null	Optional	1	サンプリング温度（0から2）。高い値はランダム性を増加させ、低い値はより焦点が絞られた決定論的な出力になります。

curlのbodyを構築するための構造体をそのまま実装します．
modelとmessages以外はパラメータが指定されないケースもあるため，OptionにしてNoneの場合にはJSONにシリアライズしないように設定しています．

#[derive(Debug, Deserialize, Serialize)]
pub struct ChatCompletionRequestBody {
    // ID of the model to use. (https://platform.openai.com/docs/models#model-endpoint-compatibility)
    pub model: String,
    // A list of messages comprising teh conversation so far.
    pub messages: Vec<Message>,
    // Whether or not to store the output of this chat completion request for user. false by default.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub store: Option<bool>,
    // -2.0 ~ 2.0. Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    // Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object that maps tokens to an associated bias value from 100 to 100.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logit_bias: Option<FxHashMap<String, i32>>,
    // Whether to return log probabilities of the output tokens or not.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub logprobs: Option<bool>,
    // 0 ~ 20. Specify the number of most likely tokens to return at each token position, each with an associated log probability.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_logprobs: Option<u8>,
    // An upper bound for the number of tokens that can be generated for a completion.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_completion_tokens: Option<u64>,
    // How many chat completion choices to generate for each input message. 1 by default.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u32>,
    // Output types that you would like the model to generate for this request. ["text"] for most models.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub modalities: Option<Vec<String>>,
    // -2.0 ~ 2.0. Positive values penalize new tokens based on whether they apper in the text so far, increasing the model's likelihood to talk about new topics.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    // 0 ~ 2. What sampling temperature to use. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    // An object specifying the format that the model must output. (https://platform.openai.com/docs/guides/structured-outputs)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
}

API本体の実装

続いてOpenAIのAPI本体を扱う構造体です．
APIキーとbodyを持つだけの非常にシンプルな構造体です．

pub struct OpenAI {
    api_key: String,
    pub completion_body: ChatCompletionRequestBody,
}

今回はSemantic Scholarの場合と異なりパラメータの数が多いため，引数に指定するとライブラリとして使う場合に引数が多すぎて面倒です．
なので，必要なパラメータだけを設定できるようにチェイン形式でパラメータを設定できる実装にします．
ちなみに，これもSemantic Scholarの場合とは異なり，OpenAIではAPIキーが必須になるので環境変数からキーを読み取って失敗した場合にはエラーを返すようにします．

impl OpenAI {
    pub fn new() -> Self {
        dotenv().ok();
        let api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY is not set.");
        return Self {
            api_key,
            completion_body: ChatCompletionRequestBody::default(),
        };
    }

    pub fn model_id(&mut self, model_id: &str) -> &mut Self {
        self.completion_body.model = model_id.to_string();
        return self;
    }

    pub fn messages(&mut self, messages: Vec<Message>) -> &mut Self {
        self.completion_body.messages = messages;
        return self;
    }

    pub fn store(&mut self, store: bool) -> &mut Self {
        self.completion_body.store = Option::from(store);
        return self;
    }

    pub fn frequency_penalty(&mut self, frequency_penalty: f32) -> &mut Self {
        self.completion_body.frequency_penalty = Option::from(frequency_penalty);
        return self;
    }

    pub fn logit_bias(&mut self, logit_bias: FxHashMap<String, i32>) -> &mut Self {
        self.completion_body.logit_bias = Option::from(logit_bias);
        return self;
    }

    pub fn logprobs(&mut self, logprobs: bool) -> &mut Self {
        self.completion_body.logprobs = Option::from(logprobs);
        return self;
    }

    pub fn top_logprobs(&mut self, top_logprobs: u8) -> &mut Self {
        self.completion_body.top_logprobs = Option::from(top_logprobs);
        return self;
    }

    pub fn max_completion_tokens(&mut self, max_completion_tokens: u64) -> &mut Self {
        self.completion_body.max_completion_tokens = Option::from(max_completion_tokens);
        return self;
    }

    pub fn n(&mut self, n: u32) -> &mut Self {
        self.completion_body.n = Option::from(n);
        return self;
    }

    pub fn modalities(&mut self, modalities: Vec<String>) -> &mut Self {
        self.completion_body.modalities = Option::from(modalities);
        return self;
    }

    pub fn presence_penalty(&mut self, presence_penalty: f32) -> &mut Self {
        self.completion_body.presence_penalty = Option::from(presence_penalty);
        return self;
    }

    pub fn temperature(&mut self, temperature: f32) -> &mut Self {
        self.completion_body.temperature = Option::from(temperature);
        return self;
    }

    pub fn response_format(&mut self, response_format: ResponseFormat) -> &mut Self {
        self.completion_body.response_format = Option::from(response_format);
        return self;
    }
}

curlを叩く部分はstd::process::Commandを用いて実装します．
bodyをChatCompletionRequestBodyに持たせているので，APIのコール自体はシンプルにまとまります．

impl OpenAI {
    ...
    pub fn chat(&mut self) -> Result<Response> {
        let body = serde_json::to_string(&self.completion_body)?;
        let url = "https://api.openai.com/v1/chat/completions";
        let cmd = Command::new("curl")
            .arg(url)
            .arg("-H")
            .arg("Content-Type: application/json")
            .arg("-H")
            .arg(format!("Authorization: Bearer {}", self.api_key))
            .arg("-d")
            .arg(body)
            .output()
            .expect("Failed to execute command");

        let content = String::from_utf8_lossy(&cmd.stdout).to_string();
    ...

レスポンスのパース

次は戻り値のJSONをパースします．

APIの戻り値はJSONなので，Semantic Scholarで実装したパターンと同じです．

返ってくるJSONは以下のような感じ．

{
  "id": "chatcmpl-123",
  "object": "chat.completion",
  "created": 1677652288,
  "model": "gpt-4o-mini",
  "system_fingerprint": "fp_44709d6fcb",
  "choices": [{
    "index": 0,
    "message": {
      "role": "assistant",
      "content": "\n\nHello there, how may I assist you today?",
    },
    "logprobs": null,
    "finish_reason": "stop"
  }],
  "usage": {
    "prompt_tokens": 9,
    "completion_tokens": 12,
    "total_tokens": 21,
    "completion_tokens_details": {
      "reasoning_tokens": 0,
      "accepted_prediction_tokens": 0,
      "rejected_prediction_tokens": 0
    }
  }
}

色々情報が詰まっていますが，重要なのはchoices.message.contentの部分です．
このJSONを受け取れるようにResponse構造体を実装します．

#[derive(Debug, Deserialize, Serialize)]
pub struct Choice {
    pub index: u32,
    pub message: Message,
    pub finish_reason: String,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Usage {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
    pub completion_tokens_details: FxHashMap<String, u64>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Response {
    pub id: String,
    pub object: String,
    pub created: u64,
    pub model: String,
    pub system_fingerprint: String,
    pub choices: Vec<Choice>,
    pub usage: Usage,
}

serdeによるJSONのパースと構造体へのマッピングが本当に便利ですね．病みつきになりそう．

ともあれ，これで実装完了です．

動作確認

簡単に動かして，返事がちゃんと返ってくるかどうか確認してみます．

    let mut openai = OpenAI::new();
    let messages = vec![Message::new(
        "user".to_string(),
        "トンネルを抜けると？".to_string(),
    )];

    openai
        .model_id("gpt-4o-mini")
        .messages(messages)
        .temperature(1.0);

    let response = openai.chat().unwrap();
    println!("{}", &response.choices[0].message.content);
    // トンネルを抜けると、そこは雪国だった」というのは、川端康成の小説『雪国』の有名な冒頭文ですね。このフレーズは、トンネルを抜けた先の風景が一変する様子を象徴的に表現しています。あなたがこのフレーズについて何か特定のことを知りたい場合や、関連するテーマについてお話ししたいことがあれば、お知らせください！

temperature=1.5バージョン

// トンネルを抜けると、風景が一変することがありますよね。外に出た瞬間の光や景色は、トンネルの中にいた時とは全く異なるものが広がります。この言葉は多くの文脈で使われ、人生の新しいステージや変化を象徴することもありますよね。具体的にどんなシチュエーションを想像していますか？

ちゃんと動いているようですね！

ちなみに，openai-toolsのテストはあまり厳密には実装していません．せいぜいレスポンスが返ってきているかどうかを確認している程度です．
理由は，APIの戻り値が毎回変動するため，特定のルールで単体テストを実装してしまうと，テストが通ったり通らなかったりばらつきが生じてしまうためです．
この手の生成AIを組み込んだライブラリの品質管理は課題の一つです．

なお，OpenAIのAPIにはレスポンスのJSONフォーマットを指定する機能があるのですが，これについてはまた別の回で紹介しようと思います．

次回

OpenAIのAPIラッパーが完成したので，次は論文のテキストを要約して欲しい情報を効率的に収集できるようにプロンプトを構築していきます．

次回：Rustで学術論文からテキストを抽出する #19