HNSWLib

// process.env.LANGSMITH_API_KEY = "<YOUR API KEY HERE>";
// process.env.LANGSMITH_TRACING = "true";

import IntegrationInstallTooltip from "@mdx_components/integration_install_tooltip.mdx";
<IntegrationInstallTooltip></IntegrationInstallTooltip>

<Npm2Yarn>
  @langchain/community langchain @langchain/openai @langchain/core
</Npm2Yarn>

import { OpenAIEmbeddings } from "@langchain/openai";
import { HNSWLib } from "@langchain/community/vectorstores/hnswlib";
import { Document } from "@langchain/core/documents";
import type { AttributeInfo } from "langchain/chains/query_constructor";

/**
 * First, we create a bunch of documents. You can load your own documents here instead.
 * Each document has a pageContent and a metadata field. Make sure your metadata matches the AttributeInfo below.
 */
const docs = [
  new Document({
    pageContent:
      "A bunch of scientists bring back dinosaurs and mayhem breaks loose",
    metadata: { year: 1993, rating: 7.7, genre: "science fiction" },
  }),
  new Document({
    pageContent:
      "Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
    metadata: { year: 2010, director: "Christopher Nolan", rating: 8.2 },
  }),
  new Document({
    pageContent:
      "A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea",
    metadata: { year: 2006, director: "Satoshi Kon", rating: 8.6 },
  }),
  new Document({
    pageContent:
      "A bunch of normal-sized women are supremely wholesome and some men pine after them",
    metadata: { year: 2019, director: "Greta Gerwig", rating: 8.3 },
  }),
  new Document({
    pageContent: "Toys come alive and have a blast doing so",
    metadata: { year: 1995, genre: "animated" },
  }),
  new Document({
    pageContent: "Three men walk into the Zone, three men walk out of the Zone",
    metadata: {
      year: 1979,
      director: "Andrei Tarkovsky",
      genre: "science fiction",
      rating: 9.9,
    },
  }),
];

/**
 * Next, we define the attributes we want to be able to query on.
 * in this case, we want to be able to query on the genre, year, director, rating, and length of the movie.
 * We also provide a description of each attribute and the type of the attribute.
 * This is used to generate the query prompts.
 */
const attributeInfo: AttributeInfo[] = [
  {
    name: "genre",
    description: "The genre of the movie",
    type: "string or array of strings",
  },
  {
    name: "year",
    description: "The year the movie was released",
    type: "number",
  },
  {
    name: "director",
    description: "The director of the movie",
    type: "string",
  },
  {
    name: "rating",
    description: "The rating of the movie (1-10)",
    type: "number",
  },
  {
    name: "length",
    description: "The length of the movie in minutes",
    type: "number",
  },
];

/**
 * Next, we instantiate a vector store. This is where we store the embeddings of the documents.
 * We also need to provide an embeddings object. This is used to embed the documents.
 */
const embeddings = new OpenAIEmbeddings();
const vectorStore = await HNSWLib.fromDocuments(docs, embeddings);

<ChatModelTabs customVarName="llm" />

// @lc-docs-hide-cell

import { ChatOpenAI } from "@langchain/openai";

const llm = new ChatOpenAI({
  model: "gpt-4o",
  temperature: 0,
});

import { SelfQueryRetriever } from "langchain/retrievers/self_query";
import { FunctionalTranslator } from "@langchain/core/structured_query";

const selfQueryRetriever = SelfQueryRetriever.fromLLM({
  llm: llm,
  vectorStore: vectorStore,
  /** A short summary of what the document contents represent. */
  documentContents: "Brief summary of a movie",
  attributeInfo: attributeInfo,
  /**
   * We need to create a basic translator that translates the queries into a
   * filter format that the vector store can understand. We provide a basic translator
   * translator here, but you can create your own translator by extending BaseTranslator
   * abstract class. Note that the vector store needs to support filtering on the metadata
   * attributes you want to query on.
   */
  structuredQueryTranslator: new FunctionalTranslator(),
});

await selfQueryRetriever.invoke(
  "Which movies are rated higher than 8.5?"
);

[
  Document {
    pageContent: 'A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea',
    metadata: { year: 2006, director: 'Satoshi Kon', rating: 8.6 },
    id: undefined
  },
  Document {
    pageContent: 'Three men walk into the Zone, three men walk out of the Zone',
    metadata: {
      year: 1979,
      director: 'Andrei Tarkovsky',
      genre: 'science fiction',
      rating: 9.9
    },
    id: undefined
  }
]

import { ChatPromptTemplate } from "@langchain/core/prompts";
import { RunnablePassthrough, RunnableSequence } from "@langchain/core/runnables";
import { StringOutputParser } from "@langchain/core/output_parsers";

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

const prompt = ChatPromptTemplate.fromTemplate(`
Answer the question based only on the context provided.

Context: {context}

Question: {question}`);

const formatDocs = (docs: Document[]) => {
  return docs.map((doc) => JSON.stringify(doc)).join("\n\n");
}

// See https://js.langchain.com/docs/tutorials/rag
const ragChain = RunnableSequence.from([
  {
    context: selfQueryRetriever.pipe(formatDocs),
    question: new RunnablePassthrough(),
  },
  prompt,
  llm,
  new StringOutputParser(),
]);

await ragChain.invoke("Which movies are rated higher than 8.5?")

The movies rated higher than 8.5 are:

1. The movie directed by Satoshi Kon in 2006, which has a rating of 8.6.
2. The movie directed by Andrei Tarkovsky in 1979, which has a rating of 9.9.

const selfQueryRetrieverWithDefaults = SelfQueryRetriever.fromLLM({
  llm,
  vectorStore,
  documentContents: "Brief summary of a movie",
  attributeInfo,
  structuredQueryTranslator: new FunctionalTranslator(),
  searchParams: {
    filter: (doc: Document) => doc.metadata && doc.metadata.rating > 8.5,
    mergeFiltersOperator: "and",
  },
});