LangChain (Amazon Bedrock & Claude)で、文章から検索キーワードを抽出してみた

この記事は「LangChain Advent Calendar 2023」22日目の記事です。

LLMやLangChainを使ってデータを検索する仕組みを作る場合、一般的にはベクターストアなどを使ったRAGを目指しがちです。ですがコストや案件規模によってはもう少し簡素な仕組みで良い場合もあります。今回はGoogle検索などであるような、「文章で検索できる仕組み」を作ってみました。

検索処理の考え方

今回のサンプルでは、LLMは「入力されたテキストを検索キーワードに変換する」処理でのみ利用します。「Amazon Bedrockの使い方を教えて」と入力された場合、「Amazon Bedrock 使い方」を出力するようなイメージですね。検索処理そのものは、WordPressなどのCMSが持つ検索APIを利用します。これによって組み込みの作業量を大きく減らすことができます。

LangChain.jsでキーワード抽出処理を実装する

検索処理そのものはすでにあるものを利用するので、LLMは検索キーワードの抽出だけを行います。LangChain.jsで実装してみましょう。

Modelを用意する

まずは利用する言語モデルを選びます。今回はAmazon Bedrock経由でClaudeを利用します。

import { BedrockChat } from "langchain/chat_models/bedrock";

const model = new BedrockChat({
    model: "anthropic.claude-v2",
    region: "us-east-1",
});

PromptTemplateとOutputParserで、処理内容を指定する

続いてLLMに依頼するタスクをプロンプトとして定義しましょう。LangChainの場合、PromptTemplateで入力をテンプレート化したり、各LLMのAPIが返してくるレスポンスをパースする処理をOutputParserで追加したりできます。

今回は「入力した質問」と「そこから抽出したキーワード」の２つをオブジェクトで受け取るように作りました。

import { PromptTemplate } from 'langchain/prompts';
import { StructuredOutputParser } from "langchain/output_parsers";
    
const queryGenerationTemplate = `あなたは入力された文章から、検索キーワードを抽出する作業を行っています。入力されたメッセージから、キーワードをスペース区切りのテキストで出力してください。

{format_instructions}

input: {query}
`;
const outputParser = StructuredOutputParser.fromNamesAndDescriptions({
  input: "Text that user asked to me.",
  output: "Suggested search keyword",
})

const queryGenerationPromptTemplate = new PromptTemplate({
      template: queryGenerationTemplate,
      inputVariables: ["query"],
      partialVariables: {
          format_instructions: outputParser.getFormatInstructions()
      },
      outputParser: outputParser
});

Chain経由で実行

設定したプロンプトなどの実行は、Chain経由で行います。Chainを使わずに行う方法もありますが、Chainを通しておくことで前段または後続に処理を足すことなどもやりやすくなります。

import { LLMChain } from "langchain/chains";

const queryChain = new LLMChain({
  llm: model,
  prompt: queryGenerationPromptTemplate,
});
    
// Call the chain with a query.
const res = await queryChain.run( "momentoの使い方を教えて");
console.log(res)

今回はシンプルにするため、質問文をハードコードしました。

Chainを実行した結果をみる

実装したコードを実行してみましょう。問題なく動けば、このようなデータが取得できるはずです。

{
  "input": "momentoの使い方を教えて",
  "output": "momento 使い方"
}

あとはoutputをWordPressなりmicroCMS / Algoliaなりの検索APIにクエリとして送りましょう。

      const res = await queryChain.run( "momentoの使い方を教えて");
      const searchParams = new URLSearchParams("");
      searchParams.set("search", res.output)
      const response = await fetch(`https://YOUR_WORDPRESS_SITE_URL/wp-json/wp/v2/search?${searchParams.toString()}`)
      const posts = await response.json()
      return c.json(posts.map((post) => ({ title: post.title, id: post.id })))

Momentoに関する記事を取得することに成功しました。

[
  {
    "title": "MomentoをCloudflare  Workers / Pages functionsで使う時は、Web SDKを使おう",
    "id": 13613
  },
  {
    "title": "momentoのベクターストア機能（Vector Index）を触ってみた",
    "id": 13570
  }
]

RAGとは別の検索システム構築に、LLMを活用する

「普通にRAG作る方が楽じゃないか？」という思いは無きにしも非ずですが、あちらはあちらでEmbeddingの追加・更新・削除をどうCMSなどの元データと連携するかみたいな部分で工数が発生します。ですので、検索のためのAPIがすでにある場合、もしくはある程度高機能な場合には、このような「検索クエリ生成システム」としてLLMを使うのもありじゃないかなと思います。

コード全部

参考までに、今回のコードをざっとまとめたものです。

import { Hono } from 'hono'
import { BedrockChat } from "langchain/chat_models/bedrock";
import { LLMChain } from "langchain/chains";
import { PromptTemplate } from 'langchain/prompts';
import { StructuredOutputParser } from "langchain/output_parsers";

const bedrockApp = new Hono()

bedrockApp.get('/chat', async c => {
    const model = new BedrockChat({
        model: "anthropic.claude-v2",
        region: "us-east-1",
      });
    
      const reviewTemplate = `あなたは入力された文章から、検索キーワードを抽出する作業を行っています。入力されたメッセージから、キーワードをスペース区切りのテキストで出力してください。

{format_instructions}

input: {query}
`;
      const outputParser = StructuredOutputParser.fromNamesAndDescriptions({
        input: "Text that user asked to me.",
        output: "Suggested search keyword",
      })

      const queryGenerationPromptTemplate = new PromptTemplate({
            template: reviewTemplate,
            inputVariables: ["query"],
            partialVariables: {
                format_instructions: outputParser.getFormatInstructions()
            },
            outputParser: outputParser
      });

      const queryChain = new LLMChain({
        llm: model,
        prompt: queryGenerationPromptTemplate,
      });
    
      // Call the chain with a query.
      const res = await queryChain.run( "Amazon Bedrockの使い方を教えて");
      return c.json(res)

})