🥊

RAGを OpenAI と Elasticsearch を用いて、ローカルでシュッと試してみる

2023/10/04に公開

Elasticsearch

はじめに

RAG（Retrieval Augmented Generation）とは、外部のデータを LLM に参照させて、解答を生成させることです。

今回は、OpenAI, Elasticsearch, Next.js, LangChain.js を使って、ミニマムに RAG アプリを構築してみます。RAG の具体的な実装イメージを掴むことが目的です。

全体のコードは以下のリポジトリにアップロードされています。

実際に試す

とりあえず、Next.js プロジェクトを作ります。

npx create-next-app@latest

Elasticsearch を用意します。今回はローカルで試すだけなので、パスワードは無効にしています。

docker run -p 9200:9200 -it -m 2GB -e xpack.security.enabled=false -e discovery.type=single-node docker.elastic.co/elasticsearch/elasticsearch:8.10.2

今回、LLM に参考にしてもらうデータを用意します。

ChatGPT に「適当な物語を作って」と依頼して、作成しました。

data/example_1.txt

タイトル: 「星の子供と魔法の庭」

昔、遠い星に住む小さな星の精霊たちがいました。彼らは星々を点灯し、夜空を美しく輝かせる責任を担っていました。ある日、星の精霊たちは特別な星を発見しました。その星には魔法の庭が広がっていました。

この魔法の庭には、不思議な花や植物が咲き乱れ、夜空の輝きをさらに美しくしました。星の精霊たちはこの庭を大切にし、星々を点灯する仕事を手伝ってくれる小さな子供の星を見つけました。彼女の名前はルナで、星の精霊たちは彼女を大切な家族のように思っていました。

しかし、魔法の庭には一つだけ問題がありました。庭の奥深くに眠る魔法の樹が、庭全体を守る大切な存在でしたが、その樹は元気を失い、葉っぱがしぼんでしまいました。星の精霊たちは悲しんでいましたが、ルナは決意しました。

ルナは夜空に輝く星々から魔法の力を借り、魔法の樹を元気づけるために庭の中心に行きました。彼女は一生懸命に魔法の歌を歌い、星々の光を樹に送りました。すると、樹は再び元気を取り戻し、新しい葉っぱが生え始めました。

魔法の樹が回復すると、星の精霊たちは感謝の気持ちでルナを包み込みました。彼女の勇気と愛情が魔法の庭を救い、星々の輝きをより美しくしました。以降、ルナは星の精霊たちと共に、夜空を美しく輝かせる仕事を全うし、彼女の冒険譚は星々の中で語り継がれました。

この童話の物語から、勇気と愛情がどんな困難も乗り越えられることを学びます。また、小さな存在でも大きな影響を持つことができることを示しています。

data/example_2.txt

タイトル: 「森の精霊と魔法の鳥」

昔々、深い森の中に美しい魔法の鳥が住んでいました。この鳥は七つの色の羽毛を持ち、その歌声は風に乗って遠くの村まで聞こえました。村人たちは鳥の歌声を聞くと、幸せな夢を見ることができました。

しかし、森の精霊たちは鳥が幸せな歌を歌い続けることを心配していました。鳥の歌声が人々に幸せをもたらす一方で、森の生物たちには不幸をもたらしていたのです。鳥の歌声が美しくても、それは森の生態系にとって破壊的でした。

ある日、小さな森の精霊であるエルフィンは、鳥に会いに行くことを決意しました。彼は鳥に説明し、森の生物たちの幸せも大切だと伝えました。最初は鳥は困惑しましたが、エルフィンの言葉を聞いて理解しました。

鳥は新しい歌を作り、それは森の生物たちに幸せをもたらすものでした。鳥の新しい歌声は森中に広がり、動物たちは喜びに満ちた日々を送ることができました。同時に、村の人々も新しい歌声に感動し、幸せな夢を見ることができました。

この童話から、人々と自然の調和が大切であり、協力と優しさが問題を解決する手助けになることを学びます。

data/example_3.txt

タイトル: 「小さな魔法使いの大きな夢」

昔々、小さな村に住む一人の少年がいました。その少年の名前はリオと言い、彼は小さな魔法使いでした。村の人々はリオを笑い者にし、彼の魔法の力を信じていませんでした。しかし、リオは大きな夢を抱いていました。

ある日、村に恐ろしいドラゴンが現れ、村を襲いました。村人たちは逃げ惑い、絶望の中にいました。しかし、リオは決意しました。彼は小さな体に宿る魔法の力を信じ、村を救うことを決心しました。

リオは村の中心に立ち、勇気を振り絞りました。彼の小さな手には魔法の杖があり、その杖を振るうと鮮やかな光が放たれました。リオは魔法の力を使ってドラゴンに立ち向かい、驚くべき戦いが始まりました。

戦いは激しいものでしたが、リオの決意と勇気が彼を導きました。最終的に、リオはドラゴンを打ち破り、村を救いました。村人たちは彼の勇敢さを称賛し、魔法の力を信じるようになりました。

この童話から、小さな者でも大きな夢を持ち、信じることができれば、どんな困難も克服できることを学びます。また、信念と勇気が力に変わることを示しています。

Elasticsearch に物語のデータを追加します。アプリからファイルをアップロードして...という仕組みを作っても良いのですが、今回はスクリプトを用意しました。

scripts/loadDocs.ts

import { TextLoader } from "langchain/document_loaders/fs/text";
import { CharacterTextSplitter } from "langchain/text_splitter";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { ElasticVectorSearch } from "langchain/vectorstores/elasticsearch";
import { Client } from "@elastic/elasticsearch";

const embeddings = new OpenAIEmbeddings();
const vectorStore = new ElasticVectorSearch(embeddings, {
  client: new Client({ node: "http://localhost:9200" }),
  indexName: "sample_index",
});

async function loadDocument(textFilePath: string) {
  // データをテキストファイルから読み込む
  const loader = new TextLoader(textFilePath);
  const docs = await loader.load();

  // データを分割する
  const splitter = new CharacterTextSplitter({
    chunkSize: 7,
    chunkOverlap: 3,
  });
  const splitDocs = await splitter.splitDocuments(docs);

  console.log("Split documents", splitDocs);

  // データをElasticsearchに投入する
  const ids = await vectorStore.addDocuments(
    splitDocs.map((d) => ({
      pageContent: d.pageContent,
      metadata: { source: d.metadata.source },
    }))
  );

  console.log("Added documents", ids);
  console.info("Loaded documents", textFilePath);
}

const main = async () => {
  await loadDocument("./data/example_1.txt");
  await loadDocument("./data/example_2.txt");
  await loadDocument("./data/example_3.txt");
};

main();

OPENAI の API キーだけ設定しておきましょう。

export OPENAI_API_KEY=<Your API Key>

スクリプトを実行します。

npx ts-node ./scripts/loadDocs.ts

データが投入できました。

アプリケーションのコードを作成します。

src/components/QueryForm.tsx

"use client";

import { useState, useCallback, ChangeEvent, FormEvent } from "react";

export function QueryForm() {
  const [query, setQuery] = useState("");
  const [answer, setAnswer] = useState("");
  const [reference, setReference] = useState("");
  const [loading, setLoading] = useState(false);

  const onChange = useCallback((e: ChangeEvent<HTMLInputElement>) => {
    setQuery(e.target.value);
  }, []);

  const onSubmit = useCallback(
    async (e: FormEvent<HTMLFormElement>) => {
      e.preventDefault();
      setLoading(true);

      try {
        const res = await fetch(`/api/question?query=${query}`);

        const data = await res.json();

        if (!res.ok) {
          console.error("Failed to fetch answer: ", data);
          return;
        }

        setAnswer(data.answer);
        setReference(data.ref);
      } catch (err) {
        setAnswer("");
        setReference("");
      } finally {
        setLoading(false);
      }
    },
    [query]
  );

  return (
    <>
      <form onSubmit={onSubmit}>
        <div className="py-1 space-x-4">
          <input
            className="outline"
            type="text"
            value={query}
            onChange={onChange}
          />
          <button
            className="outline"
            type="submit"
            disabled={loading || query.length <= 0}
          >
            送信
          </button>
        </div>
      </form>
      <div className="flex flex-col items-center">
        <div className="py-1"> 回答: </div>
        <div className="py-1">{loading ? "loading..." : answer}</div>
        <div className="py-1"> 参考にしたデータ: </div>
        <div className="py-1">{loading ? "loading..." : reference}</div>
      </div>
    </>
  );
}

src/app/page.tsx

import { QueryForm } from "@/components/QueryForm";

export default function Home() {
  return (
    <main className="flex min-h-screen flex-col items-center p-24">
      <div className="py-1 text-xl">RAG sample</div>
      <QueryForm />
    </main>
  );
}

解答を返す API エンドポイント GET api/question も作ります。

src/app/api/question/route.ts

import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { ElasticVectorSearch } from "langchain/vectorstores/elasticsearch";
import { OpenAI } from "langchain/llms/openai";
import { PromptTemplate } from "langchain/prompts";
import { Client } from "@elastic/elasticsearch";
import { type NextRequest } from "next/server";

const embeddings = new OpenAIEmbeddings();
const vectorStore = new ElasticVectorSearch(embeddings, {
  client: new Client({ node: "http://localhost:9200" }),
  indexName: "sample_index",
});

const llm = new OpenAI();

export async function GET(request: NextRequest) {
  const query = request.nextUrl.searchParams.get("query");

  if (!query) {
    return Response.json({ error: "No query provided" }, { status: 400 });
  }

  // queryに最も類似するデータをElasticsearchから取得する
  const results = await vectorStore.similaritySearch(query, 1);
  const refData = results[0].pageContent;

  // 取得したデータを参考にして、queryに対する解答を生成する
  const prompt = PromptTemplate.fromTemplate(
    `
    Please answer the following question:
    Question: {query}
    If necessary, please refer to the following information.
    Information: {refData}
    Answer in the same language as the question.`
  );
  const formattedPrompt = await prompt.format({
    query,
    refData,
  });
  const answer = await llm.predict(formattedPrompt);

  return Response.json({
    answer,
    ref: refData,
  });
}

Next.js を起動します。

npm run dev

物語のデータを参照できるなら回答できる質問をしてみます。

追加した物語のデータを参照して、回答してくれました 🎉

ちなみに ChatGPT に聞くと、今回投入した物語のデータを参照できないので、"知らない"と返ってきます。~~でも、ちゃんと推測はしてる~~

おわりに

RAG は、とりあえず作る分には、さっと作れることがわかりました。特にプロンプトなどを色々いじってイメージを掴むには、ローカルで動かせるのはメリットだと感じます。

実運用するなら、考慮しないといけない点は山ほどあるので、今後も勉強していきます。

はじめに

実際に試す

おわりに

Discussion