🦀

RustでOllamaにAPIを打つ

2024/08/31に公開

はじめに

対象読者:Rust、Go、C、C++に触れたことがなく、知識もない人
最近Rustの学習を始めました。
普段はPythonでプログラムを書いています。
背景として、Pythonが重いと感じることがあり、軽量な言語も習得したいと考えたためです。
RustはKernelにも採用される軽い言語です。
まだまだ学習の途中ですが、途中経過のアウトプットとして、Ollama APIを利用するプログラムをRustで実装しました。
このプログラムを通じて、Rustについて説明します。

環境構築

Rustには、rustupというpyenvのようなツールがあり、様々なtoolchainを使えます。
Rustにはリリースバージョンの進行を示すchannelsというものがあり、主に3種類あります。

  • stable: 安定版。
  • beta: 次の安定版の候補としてテストされているバージョン。
  • nightly: 最新バージョン。

toolchainは、Rustコンパイラ(rustc)やその他の関連ツール(cargoなど)のセットのことです。
channelsごとにあるようです。

とりあえずrustupを試したいので環境構築。

FROM debian:bookworm-slim

WORKDIR /app

RUN apt-get update && apt-get install -y --no-install-recommends \
    curl \
    build-essential \
    ca-certificates \
    pkg-config \
    libssl-dev \
    && rm -rf /var/lib/apt/lists/*

RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y

ENV CARGO_HOME=/root/.cargo
ENV RUSTUP_HOME=/root/.rustup
ENV PATH=$CARGO_HOME/bin:$PATH

# 動くか確認
RUN rustc --version

起動

docker build -t rustup .
docker run -it --network host -v $(pwd):/app rustup /bin/bash

インストール済みのtoolchainを確認したらstableと出てきました。

$ rustup toolchain list
stable-x86_64-unknown-linux-gnu (default)

toolchainに入っているcargoでプロジェクトを新規作成します。

cargo new the-project

the-projectディレクトリが作成されます。

|-- the-project
    |-- Cargo.toml
    `-- src
        `-- main.rs

src配下にmain.rsができたので、コードを書きます。
依存関係等はCargo.tomlに記載します。
OllamaのModelやEndpointは、environmentディレクトリを作成して.envに外出しします。

コード

main.rs

use reqwest::blocking::{Client};
use serde_json::Value;
use std::env;
use dotenv::from_filename;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // .envファイルのパスを指定して読み込む
    from_filename("src/environment/.env").ok();

    // 環境変数からモデルとURLを取得
    let model = env::var("MODEL")?;
    let url = env::var("URL")?;

    let client = Client::new();
    let data = serde_json::json!({
        "model": model,
        "prompt": "Hello."
    });

    let mut response_text = String::new();

    let response = client.post(&url).json(&data).send()?;

    // ストリームされたデータを1行ずつ処理
    for line in response.text()?.lines() {
        if !line.trim().is_empty() {
            let partial_response: Value = serde_json::from_str(line)?;
            if let Some(text) = partial_response.get("response").and_then(|v| v.as_str()) {
                response_text.push_str(text);
            }
        }
    }

    println!("{}", response_text);
    Ok(())
}

.env
ローカルのOllamaには以下のようなendpointにmodelとpromptを指定してAPI投げられます。

MODEL=llama3.1:8b-instruct-q5_K_M
URL=http://127.0.0.1:11434/api/generate

Cargo.toml

[package]
name = "the-project"
version = "0.1.0"
edition = "2021"

[dependencies]
reqwest = { version = "0.11", features = ["blocking", "json"] }
serde_json = "1.0"
dotenv = "0.15"

実行

cargo build

でコンパイルができます。

cargo run

でコンパイル&実行できます。

$ cargo run
   Compiling the-project v0.1.0 (/app/the-project)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 22.94s
     Running `target/debug/the-project`
It's nice to meet you. Is there something I can help you with, or would you like to chat?

コード解説

use reqwest::blocking::{Client};
use serde_json::Value;
use std::env;
use dotenv::from_filename;

Pythonとは違い、Rustではuseでライブラリなどをimportします。
reqwestdotenvといったライブラリは、Crates(クレート)という概念の中の一部です。
クレートとは、Rustコンパイラが一度にコンパイルするコードの単位です。2種類あります。

  • バイナリクレート: 実行可能なプログラムを生成するクレートです。main 関数を持ちます。
  • ライブラリクレート: 共有される機能を定義し、他のプロジェクトで再利用できるライブラリを提供するクレートです。main 関数を持たず、直接実行されることはありません。

reqwestdotenvといったライブラリは、ライブラリクレートに該当します。
reqwestの次のblockingModules(モジュール)という、クレートを論理的に整理できる仕組みです。
例えば、reqwestのsrcの中にはblockingというディレクトリがあって、client.rsの中にClient構造体があります。

関数

fn main() -> Result<(), Box<dyn std::error::Error>> {
    ...
}

-> Result<(), Box<dyn std::error::Error>> は戻り値を示しています。
Result は、Rustの標準ライブラリに含まれる列挙型で、OkErr の2つを持ちます。
関数が正常に終了する場合は Ok(()) を返し、何かエラーが発生した場合は Err(Box<dyn std::error::Error>) を返します。

変数

letはimmutableで、let mutをつけることでmutableになります。

所有権

Rustには所有権という重要な概念があります。この説明は全体の一部になりますが、所有権の基本的なルールを簡単に紹介します。
まずRustには以下のルールがあります。

  1. 各値には所有者がいる。: Rustの各値(例えば、Stringなど)は、どこかに所有者(変数)が存在します。
  2. 一つの値には一人の所有者しかいない。 : 値は一度に一つの所有者だけを持つことができます。所有権が別の変数に移った場合、元の変数はその値へのアクセス権を失います。
  3. 所有者がスコープを抜けると、値はドロップされる。: 変数がスコープを抜けると、その変数が所有していた値はメモリから自動的に解放されます(ドロップされます)。
    具体例として、次のコードを見てみましょう:
let s1 = String::from("hello"); // s1がString::from("hello")を所持 
let s2 = s1; // s1の所有権がs2に移るため、s1は無効になる

※オブジェクトには、ヒープ、スタックがあって
ヒープは値そのものと各文字のindexがあり
スタックにはその値のヒープポインタ、長さ、容量が入っています。

1行目で、s1が"hello"という文字列の所有者になります。ここで、"hello"はヒープ領域に保存され、s1はスタック上にそのヒープへのポインタ、長さ、容量の情報を持ちます。
2行目で、s1の所有権がs2に移ります。この際、s1のスタック上のデータがs2にコピーされますが、s1自体は無効化され、もうアクセスすることはできなくなります。これにより、メモリが二重に解放されることを防ぎます。

では、関数に変数を渡す場合はどうでしょうか?所有権が移ると、元の変数は無効になりますが、関数にその変数を借用させることができます。これを実現するのが&です。&を使うことで、関数は値そのものではなく、その値への参照(ポインタ)を受け取り、所有権を奪わずに値を利用することができます。

    let response = client.post(&url).json(&data).send()?;

&urlと&dataは、それぞれurlとdataへの参照を関数に渡していることを意味します。この場合、urlやdataの所有権はそのまま保たれ、関数はこれらの値を一時的に「借用」するだけです。この仕組みを「借用」と呼びます。

ヘッドウォータース

Discussion