[LangChain.jsでいろんなRAGを作る]Cloudflare Workers AIで作ったRAGに翻訳機能を追加してみた

CloudflareでRAGを構築するメリットの１つは、「Cloudflareだけで必要なリソースを揃えることができる」ことです。LLMの呼び出しにはWorkers AIが利用できますし、Vectorizeを使ってRAGに利用するベクターストアを用意できます。

ただ、利用するモデルの関係かプロンプトの関係か、RAGで生成した回答文章が英語になることが時々ありました。そこで今回は、「テキストを翻訳するステップ（Chain）」を明示的に追加して、指定した言語で回答が生成されるようにChainを構築する方法を試してみました。

質問に対する回答を生成する実装

まずはシンプルなテキスト生成処理を作りましょう。Cloudflare Workers AIを利用して、質問に回答するChainを用意します。

  const question = "AWS Lambdaの使い方を教えて"
  const chatCloudflare = new ChatCloudflareWorkersAI({
    model: "@hf/thebloke/neural-chat-7b-v3-1-awq", // Default value
    cloudflareAccountId: c.env.CLOUDFLARE_ACCOUNT_ID,
    cloudflareApiToken: c.env.CLOUDFLARE_API_TOKEN,
    cache: false,
  });

  const chatGenerateAnswerPrompt = ChatPromptTemplate.fromMessages([
    [
      "system",
      `Answer the question:`,
    ],
    ["human", "{question}"],
  ])

  const generateAnswerChain = RunnableSequence.from([
    chatGenerateAnswerPrompt,
    chatCloudflare,
    new StringOutputParser()
  ])

  const result = await generateAnswerChain.invoke({question})
  console.log(result)

この実装でも、質問に対する回答は得られます。ただしここで得られる回答は、利用するLLMが学習した結果に依存します。RAGとして構築するには、このChainに対して検索結果を渡してやる必要があります。

検索のステップを追加する

上のChainのプロンプトに渡してやる情報を検索するステップを追加しましょう。今回はインデックス作成のステップを省くため、任意のWordPressサイトに対して検索を行う形で代用します。下のサンプルコードは、{ここをWordPressサイトのドメインに変える}と書かれている部分を、所有されているWordPressサイトのドメインに変更して実行しましょう。

  // REST APIの検索クエリを生成するChain
  const chatExtractSearchKeywordPrompt = ChatPromptTemplate.fromMessages([
      ['system', `Given a question, extract the essential keywords for a search query. 
      Ensure only keywords are produced, with no sentences or explanations. 
  
      - Input should be processed to identify the core subjects and actions. 
      - Output strictly as keywords, separated by spaces, capturing the essence of the query.
  
      Example inputs and expected outputs:
      - Input: "How to use Hono on AWS Lambda?" => Expected Output: "Hono AWS Lambda"
      - Input: "How to use AWS?" => Expected Output: "AWS"
  
      Reminder: Do not generate any additional text. Focus solely on extracting and outputting keywords.`],
      ['human', '{question}']
  ])
  const extractSearchKeywordChain = RunnableSequence.from([
    chatExtractSearchKeywordPrompt,
    chatCloudflare,
    new StringOutputParser()
  ])


  // WP APIに検索リクエストを投げるChain
  const searchPostChain = RunnableSequence.from([
    {
      keyword: extractSearchKeywordChain
    },
    new RunnableLambda({
      func: async (input) => {
        const result = await fetch(`https://{ここをWordPressサイトのドメインに変える}/wp-json/wp/v2/posts?search=${encodeURIComponent(input.keyword)}`)
        const posts = await result.json()
        const excerpt = (posts as any)[0].excerpt.rendered
        return excerpt
      }
    }),
    new StringOutputParser(),
  ])

最後に初めに作った回答生成のChainを更新して、検索結果を利用するようにしましょう。

  const chatGenerateAnswerPrompt = ChatPromptTemplate.fromMessages([
    [
      "system",
      `Answer the question based on only the following context:
    
  {context}`,
    ],
    ["human", "{question}"],
  ])
  const generateAnswerChain = RunnableSequence.from([
    {
      context: searchPostChain,
      question: new RunnablePassthrough()
    },
    chatGenerateAnswerPrompt,
    chatCloudflare,
    new StringOutputParser()
  ])

ここまでで質問に対して、モデルの外で取得した情報を利用した回答文章生成を行うフローが完成しました。この処理を実行して、テキストが意図した言語以外で生成されることがある場合、次のステップを追加して翻訳を行うようにしましょう。

テキストを翻訳するステップ（Chain）を追加する

Cloudflare Workers AIを利用している場合、テキストの翻訳に利用できるモデルが用意されています。このモデルは2024/03時点でLangChain.jsから直接利用ができませんので、RunnableLambdaとWorkers AI SDKを利用して実装しましょう。

  const japaneseTranslationChain = RunnableSequence.from([
    new RunnableLambda({
      func: async (input) => {
        console.log(input.answer)
        const { translated_text: translatedQuestion } = await ai.run('@cf/meta/m2m100-1.2b', {
            text: input.answer,
            source_lang: "english", // defaults to english
            target_lang: "japanese"
          }
        );
        return translatedQuestion
      }
    }),
    new StringOutputParser()
  ])

前のステップで実行するChain(generateAnswerChain)の結果を翻訳させたいので、配列の先頭にオブジェクトを追加します。

  const japaneseTranslationChain = RunnableSequence.from([
    {
      answer: generateAnswerChain
    },
    new RunnableLambda({
      func: async (input) => {
        console.log(input.answer)
        const { translated_text: translatedQuestion } = await ai.run('@cf/meta/m2m100-1.2b', {
            text: input.answer,
            source_lang: "english", // defaults to english
            target_lang: "japanese"
          }
        );
        return translatedQuestion
      }
    }),
    new StringOutputParser()
  ])

このオブジェクトを追加することで、generateAnswerChainの実行結果を、input.answerとしてWorkers AI SDKのリクエストに送信できます。

ここまでのChainをつなぎ合わせることで、自然言語で入力された質問に対して、WordPressサイトの記事情報を検索し、その結果を元に回答文章を日本語で生成するようになりました。

Appendix: 質問文（Embeddingや検索クエリに渡す文章）も英語にする

記事で紹介した方法は、回答文章を翻訳しました。今回紹介した実装方法は、「入力された質問文を翻訳する」処理にも利用できます。例えば次のChainは、日本語で入力されたテキストを英語に翻訳します。

  const ai = new Ai(c.env.AI);
  const englishTranslationChain = RunnableSequence.from([
    {
      question: input => input.question,
    },
    new RunnableLambda({
      func: async (input) => {
        const { translated_text: translatedQuestion } = await ai.run('@cf/meta/m2m100-1.2b', {
            text: input.question,
            source_lang: "japanese", // defaults to english
            target_lang: "english"
          }
        );
        return translatedQuestion
      }
    }),
    new StringOutputParser()
  ])

このChainを「検索クエリを生成するChain」の前に実行することで、検索クエリを英語のみに強制することができます。

  const chatExtractSearchKeywordPrompt = ChatPromptTemplate.fromMessages([
    ['system', `Given a question, extract the essential keywords for a search query. 
    Ensure only keywords are produced, with no sentences or explanations. 

    - Input should be processed to identify the core subjects and actions. 
    - Output strictly as keywords, separated by spaces, capturing the essence of the query.

    Example inputs and expected outputs:
    - Input: "How to use Hono on AWS Lambda?" => Expected Output: "Hono AWS Lambda"
    - Input: "How to use AWS?" => Expected Output: "AWS"
    - Input: "Tell me about Amazon S3?" => Expected Output: "Amazon S3"
    - Input: "AWS Lambdaの使い方を教えて"=> Expected Output: "AWS Lambda"
    - Input: "Learn how to use AWS Lambda" => Expected Output: "AWS Lambda"

    Reminder: Do not generate any additional text. Focus solely on extracting and outputting keywords.`],
    ['human', '{question}']
  ]);
  const extractSearchKeywordChain = RunnableSequence.from([
    {
      question: englishTranslationChain
    },
    chatExtractSearchKeywordPrompt,
    chatCloudflare,
    new StringOutputParser()
  ]);

多言語サイトなどでRAGを提供したい場合には、このような前処理を追加することで検索の精度を高めることができるかもしれません。

まとめ

LLMは入力された言語やそのモデルが得意とする言語で回答文章を生成するように見えます。そのため、時より意図しない言語で回答文章が生成されることがあります。そのような場合に備えて、今回のCloudflare Workers AIを利用した例のような、翻訳を行うステップを準備しておくことも検討したほうがよいかもしれません。