@langchain/community/dist/structured_query/tests/chroma_self_query.int.test.js

Third-party integrations for LangChain.js

import { test } from "@jest/globals";
import { Document } from "@langchain/core/documents";
import { OpenAIEmbeddings, OpenAI } from "@langchain/openai";
import { SelfQueryRetriever } from "langchain/retrievers/self_query";
import { Chroma } from "../../vectorstores/chroma.js";
import { ChromaTranslator } from "../chroma.js";
test.skip("Chroma Store Self Query Retriever Test", async () => {
    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,
            },
        }),
    ];
    const 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",
        },
    ];
    const embeddings = new OpenAIEmbeddings();
    const llm = new OpenAI();
    const documentContents = "Brief summary of a movie";
    const vectorStore = await Chroma.fromDocuments(docs, embeddings, {
        collectionName: "a-movie-collection",
    });
    const selfQueryRetriever = SelfQueryRetriever.fromLLM({
        llm,
        vectorStore,
        documentContents,
        attributeInfo,
        structuredQueryTranslator: new ChromaTranslator(),
    });
    const query1 = await selfQueryRetriever.getRelevantDocuments("Which movies are less than 90 minutes?");
    const query2 = await selfQueryRetriever.getRelevantDocuments("Which movies are rated higher than 8.5?");
    const query3 = await selfQueryRetriever.getRelevantDocuments("Which movies are directed by Greta Gerwig?");
    const query4 = await selfQueryRetriever.getRelevantDocuments("Which movies are either comedy or drama and are less than 90 minutes?");
    console.log(query1, query2, query3, query4);
});