LangChainからAmazon Bedrock Knowledge Baseでインデックスを生成したPineconeインデックスにアクセスする

以前BerdockのKnowledge Baseを使ってEmbeddingしたデータのインデックスなどを作りました。その際にPineconeをベクターストアに使ったのですが、このインデックスをLangChainから利用したらどうなるか気になったので試してみました。

OpenAI APIとPineconeで検索だけ動かしてみる

ベクターストア周りはLangChainのPineconeインテグレーションをそのまま利用します。シンプルにOpenAI APIを利用する形でまず組んでみました。

    const pinecone = new Pinecone({
        apiKey: env.pinecone.API_KEY,
        environment: env.pinecone.ENV
    });
    
    const pineconeIndex = pinecone.Index(env.pinecone.PINECONE_INDEX);

    const vectorStore = await PineconeStore.fromExistingIndex(
        new OpenAIEmbeddings({
            openAIApiKey: env.openai.apiKey
        }),
        { pineconeIndex }
      );
      const results = await vectorStore.similaritySearch("subscription", 1);
      console.log(results);

結果がこちら。Knowledge Base経由でEmbddingしていたS3のデータを取れているのがなんとなく伺えます。

[
  {
    "metadata": {
      "metadata": "{\"source\":\"s3://DUMMY_BUCKET_NAME/app-flow/stripe-product-pricing/stripe-price/c7f15f4e-906f-4217-9a42-9675341a9773/-1259199976-2023-12-11T14:11:13\"}",
      "question1": "{\"id\":\"price_1KJALKL6R1kGwUF47gErXS96\",\"object\":\"price\",\"active\":true,\"billing_scheme\":\"per_unit\",\"created\":1642484094,\"currency\":\"jpy\",\"custom_unit_amount\":null,\"livemode\":false,\"lookup_key\":\"membership\",\"metadata\":{},\"nickname\":\"monthly membership\",\"product\":\"prod_Kz8cjSg54JYbjd\",\"recurring\":{\"aggregate_usage\":null,\"interval\":\"month\",\"interval_count\":1,\"trial_period_days\":null,\"usage_type\":\"licensed\"},\"tax_behavior\":\"unspecified\",\"tiers_mode\":null,\"transform_quantity\":null,\"type\":\"recurring\",\"unit_amount\":1000,\"unit_amount_decimal\":\"1000\"} {\"id\":\"price_1KJALKL6R1kGwUF4NejvgEUz\",\"object\":\"price\",\"active\":true,\"billing_scheme\":\"per_unit\",\"created\":1642484094,\"currency\":\"jpy\",\"custom_unit_amount\":null,\"livemode\":false,\"lookup_key\":null,\"metadata\":{},\"nickname\":\"simple"
    }
  }
]

ChainでRAGのようなものを作ってみる

データは取れそうなので、RAGを目指してChainを組んでいきます。

Bedrockのモデルを利用できるようにSDKを追加する

LangChainが内部的に利用するランタイムをインストールしましょう。

npm i '@aws-sdk/client-bedrock-runtime'

LLMのモデルで実装してみる

まずはLLMs側でワンショットのChainを作ってみます。Embeddingで生成するベクトルデータを揃えた方がよいかもと思ったので、ここでEmbeddingに利用するモデルをBedrockに切り替えています。

jsonApp.get('load', async c => {

    const pinecone = new Pinecone({
        apiKey: env.pinecone.API_KEY,
        environment: env.pinecone.ENV
    });
    
    const pineconeIndex = pinecone.Index(env.pinecone.PINECONE_INDEX);
    const embeddings = new BedrockEmbeddings({
      region: env.aws.region,
      credentials: {
        accessKeyId: env.aws.credentials.accessKeyId,
        secretAccessKey: env.aws.credentials.secretAccessKey,
      },
      model: "amazon.titan-embed-text-v1", // Default value
    });

    const vectorStore = await PineconeStore.fromExistingIndex(
        embeddings,
        { pineconeIndex }
      );
      const model = new OpenAI({
        openAIApiKey: env.openai.apiKey
      });
      const chain = VectorDBQAChain.fromLLM(model, vectorStore, {
        k: 1,
        returnSourceDocuments: true,
      });
      const response = await chain.call({ query: "Stripeでのサブスクリプションの作り方" });
      console.log(response);
      return c.json(response)
})

生成された文章がこちら。returnSourceDocumentsを含めているからか、参照したデータまで見れるのがよいですね。ただし回答の精度については、「I don’t know」が混ざっているあたり、要調整に見えます。

{
  "text": " Stripeのサブスクリプションを作るには、まずStripeのウェブサイトにアクセスし、アカウントを作成します。その後、ダッシュボードから「サブスクリプション」を選択し、必要な情報を入力してプランを作成します。作成したプランをウェブサイトに組み込んで、顧客がサブスクリプションを購入できるようにします。\n\nI don't know.",
  "sourceDocuments": [
    {
      "metadata": {
        "metadata": "{\"source\":\"s3://DUMMY_BUCKET_NAME/app-flow/stripe-product-pricing/3300a769-6250-4833-80c5-d93af6ce2d16/-1264520871-2023-12-01T06:56:56\"}",
        "question1": "{\"id\":\"prod_KugwwcPLrFjsBW\",\"object\":\"product\",\"active\":true,\"attributes\":[],\"created\":1641458563,\"default_price\":\"price_1KErYWL6R1kGwUF40RAg2a4Y\",\"description\":\"(created by Stripe CLI)\",\"images\":[],\"livemode\":false,\"metadata\":{\"industry\":\"pvAeqLVP\"},\"name\":\"SubscriptionTestProduct\",\"package_dimensions\":null,\"shippable\":null,\"statement_descriptor\":null,\"tax_code\":null,\"type\":\"service\",\"unit_label\":null,\"updated\":1652121707,\"url\":null} {\"id\":\"prod_KmSIFB9LJUXdU3\",\"object\":\"product\",\"active\":true,\"attributes\":[],\"created\":1639558988,\"default_price\":\"price_1K6tOCL6R1kGwUF4NC25BLww\",\"description\":null,\"images\":[],\"livemode\":false,\"metadata\":{},\"name\":\"たべものパッケージ\",\"package_dimensions\":null,\"shippable\":null,\"statement_descriptor\":null,\"tax_code\":\"txcd_40060003\",\"type\":\"service\",\"unit_label\":null,\"updated\":1696839595,\"url\":null} {\""
      }
    }
  ]
}

BedrockChatモデルを使ってみた場合

別のモデルも試してみましょう。今度はLLMsではなくChat Modelで、Bedrockのみを使ってみます。

    const pinecone = new Pinecone({
        apiKey: env.pinecone.API_KEY,
        environment: env.pinecone.ENV
    });
    
    const pineconeIndex = pinecone.Index(env.pinecone.PINECONE_INDEX);
    const embeddings = new BedrockEmbeddings({
      region: env.aws.region,
      credentials: {
        accessKeyId: env.aws.credentials.accessKeyId,
        secretAccessKey: env.aws.credentials.secretAccessKey,
      },
      model: "amazon.titan-embed-text-v1", // Default value
    });

    const vectorStore = await PineconeStore.fromExistingIndex(
        embeddings,
        { pineconeIndex }
    );
    const model = new BedrockChat({
        model: "anthropic.claude-v2",
        region: env.aws.region,
        credentials: {
          accessKeyId: env.aws.credentials.accessKeyId,
          secretAccessKey: env.aws.credentials.secretAccessKey,
        },
    });
    const chain = VectorDBQAChain.fromLLM(model, vectorStore, {
        k: 5,
        returnSourceDocuments: true,
    });
    const response = await chain.invoke({
        query:"Stripeでのサブスクリプションの作り方"
    });

こちらの回答は次のようになりました。文章が途切れていたり、Embeddingしたデータを利用した回答ができていないなど、調整がかなり必要そうな気配を感じます。

{
  "text": " すみません、Stripeでのサブスクリプションの作り方について詳しく知りません。Stripeは決済サービスのプラットフォ",
  "sourceDocuments": [
    {
      "metadata": {
        "metadata": "{\"source\":\"s3://DUMMY_BUCKET_NAME/app-flow/stripe-product-pricing/3300a769-6250-4833-80c5-d93af6ce2d16/-1264520871-2023-12-01T06:56:56\"}",
        "question1": "{\"id\":\"prod_KugwwcPLrFjsBW\",\"object\":\"product\",\"active\":true,\"attributes\":[],\"created\":1641458563,\"default_price\":\"price_1KErYWL6R1kGwUF40RAg2a4Y\",\"description\":\"(created by Stripe CLI)\",\"images\":[],\"livemode\":false,\"metadata\":{\"industry\":\"pvAeqLVP\"},\"name\":\"SubscriptionTestProduct\",\"package_dimensions\":null,\"shippable\":null,\"statement_descriptor\":null,\"tax_code\":null,\"type\":\"service\",\"unit_label\":null,\"updated\":1652121707,\"url\":null} {\"id\":\"prod_KmSIFB9LJUXdU3\",\"object\":\"product\",\"active\":true,\"attributes\":[],\"created\":1639558988,\"default_price\":\"price_1K6tOCL6R1kGwUF4NC25BLww\",\"description\":null,\"images\":[],\"livemode\":false,\"metadata\":{},\"name\":\"たべものパッケージ\",\"package_dimensions\":null,\"shippable\":null,\"statement_descriptor\":null,\"tax_code\":\"txcd_40060003\",\"type\":\"service\",\"unit_label\":null,\"updated\":1696839595,\"url\":null} {\""
      }
    },
    {

感想

同じようなプロンプトをいれたのですが、期待した回答は得られませんでした。LangChainのChainに相当する部分がBedrockのAgentやKnowledge Base内部で実装されていると考えると、LangChainで同じようなChainに作れるのではという予測をしています。が、そこまでするくらいならLangChainからBedrockのAgentを呼び出すようなChainを作ってみるのも一つかもしれません。

ヒントになるのは、RetrievalQAChainな気がしますので、日を改めてこちらと連携したものも作ってみようと思います。