LangChainのRemoteLangChainRetrieverでREST APIを使う

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

LangChainを利用して外部のデータソースを元に回答を作る場合、一般的にはEmbeddingやベクターストアを利用したRAGを作ります。ですがLangChainのドキュメントを見ると、REST APIに検索リクエストを投げる形の実装もできそうでしたので、こちらを試してみました。

Docsのサンプルコード

ベクターストアを使わない場合、RemoteLangChainRetrieverを利用します。ドキュメントのサンプルコードはこのようなものでした。

import { OpenAI } from "langchain/llms/openai";
import { RetrievalQAChain } from "langchain/chains";
import { RemoteLangChainRetriever } from "langchain/retrievers/remote";

export const run = async () => {
  // Initialize the LLM to use to answer the question.
  const model = new OpenAI({});

  // Initialize the remote retriever.
  const retriever = new RemoteLangChainRetriever({
    url: "http://0.0.0.0:8080/retrieve", // Replace with your own URL.
    auth: { bearer: "foo" }, // Replace with your own auth.
    inputKey: "message",
    responseKey: "response",
  });

  // Create a chain that uses the OpenAI LLM and remote retriever.
  const chain = RetrievalQAChain.fromLLM(model, retriever);

  // Call the chain with a query.
  const res = await chain.call({
    query: "What did the president say about Justice Breyer?",
  });
  console.log({ res });
};

任意のAPIを呼び出して、その結果をRetrievalQAChainに渡していることがわかります。ちなみにこの時呼び出すAPIはPOSTで呼び出せる必要がある様子です。

認証のないAPIを呼び出す場合

もし認証が必要ない場合は、auth: falseにするとTypeScriptの型エラーも回避できます。

const retriever = new RemoteLangChainRetriever({
  auth: false,
  ...

RemoteLangChainRetrieverが呼び出すAPIを作ってみる

呼び出される側のAPIを作ってみましょう。Honoでざっくりと組み立てたものがこちらです。

app.post('/retrieve', async c => {
    const { message } = await c.req.json()
    console.log({message})

    return c.json({
        response: [{
           page_content:"He is awesome.",
           metadata: {
            id: 123
           }
        }]
    })
})

RemoteLangChainRetrieverを実行してみたところ、このようなログが表示されていました。inputKeyのキーを持つJSONをBodyにつけて投げてくように見えます。

{ message: 'What did the president say about Justice Breyer?' }

RemoteLangChainRetrieverのverboseをtrueにすると、作ったAPIレスポンスも確認できました。

[21ms] Exiting Retriever run with output: {
  "documents": [
    {
      "pageContent": "He is awesome.",
      "metadata": {
        "id": 123
      }
    }
  ]
}

「入力された文章をPOSTでAPIに送って、Documentの配列を受け取るAPI」を作る必要があるように見えます。

リクエストとレスポンスのハンドルをカスタマイズする

もしAPI側でリクエスト・レスポンスのフォーマットをカスタマイズしたい場合は、コンストラクタで設定します。

      const retriever = new RemoteLangChainRetriever({
        url: "http://0.0.0.0:3000/remote",
        auth: false,
        inputKey: "query",
        pageContentKey: "text",
        responseKey: "data",
        verbose: true
      });

この場合、APIの各パラメータが次のように変わります。

app.post('/remote', async c => {
    const { query } = await c.req.json()
    console.log({query})

    return c.json({
        data: [{
           text:"He is awesome.",
           metadata: {
            id: 123
           }
        }]
    })
})

このRetriever向けのAPIを作る場合にはあまり必要がないかもしれませんが、共通化などを目指す必要が出た時に出番がありそうです。

RemoteLangChainRetrieverを使う際のポイント

今の所、このAPIを使ってデータ（Document）の取得をしたい場合、次のような要件が発生しそうです。

• POSTのAPIを作る
• JSONリクエストを受け付ける
• リクエストとレスポンスのキーは、コードからデフォルトを見て設定する
• 面倒なら明示的に指定しろ

生成の結果などについては、もう少し調べてみてから改めて紹介します。