LangChain.jsでCloudflare Workers AIの翻訳モデルを利用する

LangChain.jsでCloudflare Workers AIを利用すると、簡単にチャットモデルやテキスト生成モデルを利用できます。複雑なプロンプトなどでなければ、model部分だけ後からOpenAIやClaude・Amazon Bedrockなどに差し替えることができるのもLangChain.jsのいいところだと思っています。

import { ChatCloudflareWorkersAI } from "@langchain/cloudflare";

const model = new ChatCloudflareWorkersAI({
  model: "@cf/meta/llama-2-7b-chat-int8", // Default value
  cloudflareAccountId: process.env.CLOUDFLARE_ACCOUNT_ID,
  cloudflareApiToken: process.env.CLOUDFLARE_API_TOKEN,
  // Pass a custom base URL to use Cloudflare AI Gateway
  // baseUrl: `https://gateway.ai.cloudflare.com/v1/{YOUR_ACCOUNT_ID}/{GATEWAY_NAME}/workers-ai/`,
});

const response = await model.invoke([
  ["system", "You are a helpful assistant that translates English to German."],
  ["human", `Translate "I love programming".`],
]);

console.log(response);

翻訳など、テキスト生成 / Embedding以外のモデルが動かない(2024/3時点)

ただ少し気になっているのが、Workers AIが公開している利用可能モデル全てがLangChain.jsで利用できない様子なことです。llama-2系mistral-7b系などの[Text Generation]モデルやText Embeddingsのbge-base-en系などは問題なく利用できます。ただしTranslationモデルについては、ChatCloudflareWorkersAIなどのmodelプロパティに指定しても意図した動作になりませんでした。

これは推測ですが、APIに対して入力するパラメーターがモデルの種類によって異なることが原因かなと思っています。Text Generationモデル系では、このような形でAPIリクエストを送信します。

    const messages = [
      { role: "system", content: "You are a friendly assistant" },
      {
        role: "user",
        content: "What is the origin of the phrase Hello, World",
      },
    ];

    const stream = await ai.run(
      "@cf/meta/llama-2-7b-chat-fp16", //  "@cf/mistral/mistral-7b-instruct-v0.1"なども同様
      {
        messages,
        stream: true,
      }
    );

これが翻訳のモデル（meta/m2m100-1.2b）とかですと、全く異なる内容を送信します。

    const response = await ai.run(
      "@cf/meta/m2m100-1.2b",
      {
        text: "I'll have an order of the moule frites",
        source_lang: "english", // defaults to english
        target_lang: "french",
      }
    );

LangChain.js側で、モデルの用途ごとにリクエスト内容を書き換えるとなると、モデル名ベースのif分岐が必要になりそうですので、ChatCloudflareWorkersAIやCloudflareWorkersAIクラスがテキスト生成系以外をサポートするのは難しそうに見えます。LangChain.js側でサポートしてもらうとすれば、CloudflareWorkesTranslationAIのような用途別のクラスを作るくらいでしょうか・・

RunnableLambdaを使って、Workers AI SDKを直接実行する

2024/03時点でテキスト生成系以外のモデルをLangChain.jsで利用したい場合、RunnableLambdaを利用するのが現実的かなと思います。RunnableLambdaを利用することで、動的な処理をChainにすることができますので、「テキストを生成するChain」と組み合わせて「テキスト生成 -> 翻訳」のようなフローが作れます。

実装方法としては、RunnableLambdaの中でWorkers AI SDKのモデル呼び出しを実行するだけです。

  const ai = new Ai(c.env.AI);

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

まとめ

Cloudflare Workers AIは、テキスト生成とEmbedding以外にもさまざまな種類・用途のモデルを提供しています。その一方でAPIリクエストの内容（body）が変わる関係上、LangChain.jsからストレートに利用できるモデルがすこし限定されています(2024/03時点)。

一方でLangChain.js側には、動的な処理を自前で実装してChainにする仕組みも用意されています。そのため、今回のようなLangChain.jsが提供する機能やAPIだけで達成できない実装については、ある程度自前の実装をChain化して持ち込むことも検討する必要がありそうです。

RunnableLambdaをうまく使って、Cloudflare Workers AIを使いこなしていきましょう。