ys memos
Blog
About
Features
Open Source

RustでOpenAIのChatCompletionsをストリームで受けるmod 〜2.型定義〜

rust
openai
chat-completions
gpt

2023/08/31

本コードの動作確認は2023/05時点までしか行っていません。
当時動作していたコードをそのまま公開いたします。

現在のAPIでの動作確認はしていないこと、ご了承ください。

以前、OpenAIのChat Completionsを叩くアプリを作っていた。 しかし、いつしか熱が冷めて開発が止まっていたので、眠らせておくくらいならと思い、コードを公開する。

また、一つの記事にまとめたかったが、ボリュームが大きくなりすぎてしまったので、記事をいくつかに分割した。

ここでは、各種型定義を行う。

src/
└── chatter
    ├── error.rs
    ├── headers.rs
    ├── json.rs            # <-
    ├── message.rs         # <-
    ├── mod.rs
    ├── request_builder.rs
    ├── role.rs            # <-
    └── stream_data.rs     # <-

リクエスト向けの型

use serde::Serialize;

use super::message::Message;

#[derive(Serialize)]
pub enum Model {
    #[serde(rename = "gpt-3.5-turbo")]
    Gpt35Turbo,
}

#[derive(Serialize)]
pub(super) struct Json<'a> {
    model: Model,
    messages: &'a Vec<Message>,
    stream: Option<bool>,
}

impl<'a> Json<'a> {
    pub fn new(messages: &'a Vec<Message>) -> Self {
        Self {
            model: Model::Gpt35Turbo,
            messages,
            stream: Some(true),
        }
    }
}

Chat Completionsとのやり取りで用いられるメッセージの型

use serde::Serialize;

use super::Role;

#[derive(Serialize, Clone)]
pub(super) struct Message {
    role: Role,
    content: String,
}

impl Message {
    pub fn new(role: Role, content: String) -> Self {
        Self { role, content }
    }
}

pub(super) trait ToMessage {
    fn to_message(self, role: Role) -> Message;
}

impl ToMessage for String {
    fn to_message(self, role: Role) -> Message {
        Message::new(role, self)
    }
}

Chat Completionsのメッセージの持つ役割（ロール）の型

use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub(super) enum Role {
    // #[serde(rename = "system")]
    // System,
    #[serde(rename = "user")]
    User,
    #[serde(rename = "assistant")]
    Assistant,
}

Chat Completions APIから返されるストリームの型

use std::str::FromStr;

use serde::{Deserialize, Serialize};

use super::error::{ChatterError, ChatterResult};
use super::Role;

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct Delta {
    role: Option<Role>,
    pub content: Option<String>,
}

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct Choice {
    pub delta: Delta,
    index: i128,
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct StreamData {
    id: String,
    object: String, // "chat.completion.chunk"
    created: u32,
    model: String,
    pub(super) choices: Vec<Choice>,
}

impl FromStr for StreamData {
    type Err = ChatterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match serde_json::from_str::<Self>(s) {
            Ok(res) => Ok(res),
            Err(_e) => {
                // if s.ends_with("[DONE]") {
                //     println!("FINISHED");
                // }
                Err(ChatterError::StreamParsingError)
            }
        }
    }
}

pub(super) trait ToStreamData {
    fn to_stream_data(&self) -> ChatterResult<StreamData>;
}

impl ToStreamData for str {
    fn to_stream_data(&self) -> ChatterResult<StreamData> {
        match self.strip_prefix("data:") {
            Some(data_str) => StreamData::from_str(data_str),
            _ => Err(ChatterError::StreamDataConvertingError),
        }
    }
}

use serde::Serialize;

use super::message::Message;

生成のモデルを記載するenumを定義した。開発時にはgpt-3.5-turboで試していたので、これだけ含まれている。

必要に応じて追加・選択できるようなコード設計をした。

#[derive(Serialize)]
pub enum Model {
    #[serde(rename = "gpt-3.5-turbo")]
    Gpt35Turbo,
}

実際に渡すjsonの形式を定義。命名が混乱しそうなので、ChatCompletionsRequestなどと名付けるほうが適切かもしれない。

#[derive(Serialize)]
pub(super) struct Json<'a> {
    model: Model,
    messages: &'a Vec<Message>,
    stream: Option<bool>,
}

リクエスト型のコンストラクタ。

今回はmodelは定数のように扱うし、ストリームで受けるのが確実なので、messagesのみ受け取ってリクエストjsonを生成する。

impl<'a> Json<'a> {
    pub fn new(messages: &'a Vec<Message>) -> Self {
        Self {
            model: Model::Gpt35Turbo,
            messages,
            stream: Some(true),
        }
    }
}

use serde::Serialize;

use super::Role;

リクエストに持たせるメッセージ一覧の各メッセージ。

roleは事前に定義されたenumのどれか、contentは任意の文字列を持たせられ、contentが実際のチャット内容。

#[derive(Serialize, Clone)]
pub(super) struct Message {
    role: Role,
    content: String,
}

impl Message {
    pub fn new(role: Role, content: String) -> Self {
        Self { role, content }
    }
}

Messageの利用部分を簡素化させるためのトレイト。Roleを引数とすることで、気軽にMessage変換を行う目的。

pub(super) trait ToMessage {
    fn to_message(self, role: Role) -> Message;
}

StringへのToMessage実装。

impl ToMessage for String {
    fn to_message(self, role: Role) -> Message {
        Message::new(role, self)
    }
}

use serde::{Deserialize, Serialize};

チャットにおいてはユーザとアシスタントが存在し、それらをenumにした。実際にはシステムも存在するので定義したが、当時利用しなかったのでコメントアウトしておいた。

#[derive(Serialize, Deserialize, Debug, Clone)]
pub(super) enum Role {
    // #[serde(rename = "system")]
    // System,
    #[serde(rename = "user")]
    User,
    #[serde(rename = "assistant")]
    Assistant,
}

use std::str::FromStr;

use serde::{Deserialize, Serialize};

use super::error::{ChatterError, ChatterResult};
use super::Role;

生成された文章の差分。roleはOption<Role>ではあるものの、確実にRole::Assistantになるはず。

このcontentは返答メッセージ全文ではなく差分のみが含まれるため、注意が必要。

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct Delta {
    role: Option<Role>,
    pub content: Option<String>,
}

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct Choice {
    pub delta: Delta,
    index: i128,
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize, Serialize)]
pub(super) struct StreamData {
    id: String,
    object: String, // "chat.completion.chunk"
    created: u32,
    model: String,
    pub(super) choices: Vec<Choice>,
}

&strからStreamDataへの変換を楽にするための実装。

serde_jsonを用いた変換を、それ自体を意識せずに記述できる目的である。

impl FromStr for StreamData {
    type Err = ChatterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match serde_json::from_str::<Self>(s) {
            Ok(res) => Ok(res),
            Err(_e) => {
                // if s.ends_with("[DONE]") {
                //     println!("FINISHED");
                // }
                Err(ChatterError::StreamParsingError)
            }
        }
    }
}

StreamDataへの変換トレイト

pub(super) trait ToStreamData {
    fn to_stream_data(&self) -> ChatterResult<StreamData>;
}

strへのToStreamData実装。

"data:"プレフィックスを意識せずに文字列を変換できるための意図を持つ。

impl ToStreamData for str {
    fn to_stream_data(&self) -> ChatterResult<StreamData> {
        match self.strip_prefix("data:") {
            Some(data_str) => StreamData::from_str(data_str),
            _ => Err(ChatterError::StreamDataConvertingError),
        }
    }
}

関連タグを探す

rust
openai
chat-completions
gpt