Add support for Amazon Bedrock Knowledge Base document filtering using metadata

[nick-elpedes]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

When using Bedrock to retrieve data using the AWS SDK, you are allowed to pass filters to only retrieve certain documents in your knowledge base that are tagged with metadata. This functionality should be added to Langchain, since it supports the bedrock:Retrieve function. This would allow for greater output control when using langchain to query an Amazon Knowledge Base, since you can specify what documents are available for retrieval using filters.

Filtering Using Metadata
Example API Request

Motivation

Currently, langchain supports retrieving metadata from the returned documents (see #5535 ). However, langchain does not currently support the filter aspect of the retrieval function from Bedrock. It is frustrating to have to use AWS SDK for half of a chain and use langchain for the other half. It would be nice if all the parts of the retrieval function were available to be used in langchain.

Proposal (If applicable)

Change the Amazon Knowledge Base query functionality to support filtering by metadata
Maybe something like this?

import {
  RetrieveCommand,
  BedrockAgentRuntimeClient,
  BedrockAgentRuntimeClientConfig,
  RetrievalFilter
} from "@aws-sdk/client-bedrock-agent-runtime";

import { BaseRetriever } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";

/**
 * Interface for the arguments required to initialize an
 * AmazonKnowledgeBaseRetriever instance.
 */
export interface AmazonKnowledgeBaseRetrieverArgs {
  knowledgeBaseId: string;
  topK: number;
  region: string;
  clientOptions?: BedrockAgentRuntimeClientConfig;
}

/**
 * Class for interacting with Amazon Bedrock Knowledge Bases, a RAG workflow oriented service
 * provided by AWS. Extends the BaseRetriever class.
 * @example
 * ```typescript
 * const retriever = new AmazonKnowledgeBaseRetriever({
 *   topK: 10,
 *   knowledgeBaseId: "YOUR_KNOWLEDGE_BASE_ID",
 *   region: "us-east-2",
 *   clientOptions: {
 *     credentials: {
 *       accessKeyId: "YOUR_ACCESS_KEY_ID",
 *       secretAccessKey: "YOUR_SECRET_ACCESS_KEY",
 *     },
 *   },
 * });
 *
 * const docs = await retriever.getRelevantDocuments("How are clouds formed?");
 * ```
 */
export class AmazonKnowledgeBaseRetriever extends BaseRetriever {
  static lc_name() {
    return "AmazonKnowledgeBaseRetriever";
  }

  lc_namespace = ["langchain", "retrievers", "amazon_bedrock_knowledge_base"];

  knowledgeBaseId: string;

  topK: number;

  bedrockAgentRuntimeClient: BedrockAgentRuntimeClient;

  constructor({
    knowledgeBaseId,
    topK = 10,
    clientOptions,
    region,
  }: AmazonKnowledgeBaseRetrieverArgs) {
    super();

    this.topK = topK;
    this.bedrockAgentRuntimeClient = new BedrockAgentRuntimeClient({
      region,
      ...clientOptions,
    });
    this.knowledgeBaseId = knowledgeBaseId;
  }

  /**
   * Cleans the result text by replacing sequences of whitespace with a
   * single space and removing ellipses.
   * @param resText The result text to clean.
   * @returns The cleaned result text.
   */
  cleanResult(resText: string) {
    const res = resText.replace(/\s+/g, " ").replace(/\.\.\./g, "");
    return res;
  }

  async queryKnowledgeBase(query: string, topK: number, metadataFilter?: RetrievalFilter) {
    const retrieveCommand = new RetrieveCommand({
      knowledgeBaseId: this.knowledgeBaseId,
      retrievalQuery: {
        text: query,
      },
      retrievalConfiguration: {
        vectorSearchConfiguration: {
          filter: metadataFilter, // Add metadata filter here
          numberOfResults: topK,
        },
      },
    });

    const retrieveResponse = await this.bedrockAgentRuntimeClient.send(
      retrieveCommand
    );

    return (
      retrieveResponse.retrievalResults?.map((result) => ({
        pageContent: this.cleanResult(result.content?.text || ""),
        metadata: {
          source: result.location?.s3Location?.uri,
          score: result.score,
          ...result.metadata,
        },
      })) ?? ([] as Array<Document>)
    );
  }

  async _getRelevantDocuments(query: string, metadataFilter?: RetrievalFilter): Promise<Document[]> {
    const docs = await this.queryKnowledgeBase(query, this.topK, metadataFilter);
    return docs;
  }
}

[nick-elpedes]

Edited initial post to include a possible proposal of implementation in langchain js

[shahbhavik01]

Any updates on this? This is a crucial functionality for something we are trying to build.

[RPKosec]

If it helps, it was implemented in libs/langchain-aws/src/retrievers/bedrock.ts and released 3 days ago in version 0.0.6 of the langchain/aws. But I agree, it would make more sense do extend the langchain/community implementation instead of having two.

[shahbhavik01]

The Langchain AWS documentation is really confusing. It would really help to have that streamlined. For example, in the langchain-aws documentation, the link for "See Usage Example" takes you to a documentation that's for the community implementation! It's also missing the "API Reference"

https://python.langchain.com/v0.2/docs/integrations/platforms/aws/#amazon-bedrock-knowledge-bases