forked from langchain-ai/langchainjs
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix: Chroma and Pinecone integration tests & standalone translator cl…
…asses (langchain-ai#1542) * added chroma and pinecone integration test, also created pinecone and chroma standalone translators * replaced neq with ne * Removed in and nin as comparator * Separate into entrypoints --------- Co-authored-by: jacoblee93 <[email protected]>
- Loading branch information
1 parent
c4f4b49
commit 9af90c9
Showing
10 changed files
with
296 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| import { Comparators, Operators } from "../../chains/query_constructor/ir.js"; | ||
| import { BasicTranslator } from "./base.js"; | ||
|
|
||
| export class ChromaTranslator extends BasicTranslator { | ||
| constructor() { | ||
| super({ | ||
| allowedOperators: [Operators.and, Operators.or], | ||
| allowedComparators: [ | ||
| Comparators.eq, | ||
| Comparators.ne, | ||
| Comparators.gt, | ||
| Comparators.gte, | ||
| Comparators.lt, | ||
| Comparators.lte, | ||
| ], | ||
| }); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| import { Comparators, Operators } from "../../chains/query_constructor/ir.js"; | ||
| import { BasicTranslator } from "./base.js"; | ||
|
|
||
| export class PineconeTranslator extends BasicTranslator { | ||
| constructor() { | ||
| super({ | ||
| allowedOperators: [Operators.and, Operators.or], | ||
| allowedComparators: [ | ||
| Comparators.eq, | ||
| Comparators.ne, | ||
| Comparators.gt, | ||
| Comparators.gte, | ||
| Comparators.lt, | ||
| Comparators.lte, | ||
| ], | ||
| }); | ||
| } | ||
| } |
103 changes: 103 additions & 0 deletions
103
langchain/src/retrievers/self_query/tests/chroma_self_query.int.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| import { test } from "@jest/globals"; | ||
| import { Document } from "../../../document.js"; | ||
| import { AttributeInfo } from "../../../schema/query_constructor.js"; | ||
| import { OpenAIEmbeddings } from "../../../embeddings/openai.js"; | ||
| import { SelfQueryRetriever } from "../index.js"; | ||
| import { ChromaTranslator } from "../chroma.js"; | ||
| import { OpenAI } from "../../../llms/openai.js"; | ||
| import { Chroma } from "../../../vectorstores/chroma.js"; | ||
|
|
||
| test("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: 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 = await 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); | ||
| }); |
122 changes: 122 additions & 0 deletions
122
langchain/src/retrievers/self_query/tests/pinecone_self_query.int.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,122 @@ | ||
| /* eslint-disable no-process-env */ | ||
| import { test } from "@jest/globals"; | ||
| import { PineconeClient } from "@pinecone-database/pinecone"; | ||
| import { Document } from "../../../document.js"; | ||
| import { AttributeInfo } from "../../../schema/query_constructor.js"; | ||
| import { OpenAIEmbeddings } from "../../../embeddings/openai.js"; | ||
| import { SelfQueryRetriever } from "../index.js"; | ||
| import { PineconeTranslator } from "../pinecone.js"; | ||
| import { OpenAI } from "../../../llms/openai.js"; | ||
| import { PineconeStore } from "../../../vectorstores/pinecone.js"; | ||
|
|
||
| test("Pinecone 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: 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", | ||
| }, | ||
| ]; | ||
|
|
||
| if ( | ||
| !process.env.PINECONE_API_KEY || | ||
| !process.env.PINECONE_ENVIRONMENT || | ||
| !process.env.PINECONE_INDEX | ||
| ) { | ||
| throw new Error( | ||
| "PINECONE_ENVIRONMENT and PINECONE_API_KEY and PINECONE_INDEX must be set" | ||
| ); | ||
| } | ||
|
|
||
| const client = new PineconeClient(); | ||
| await client.init({ | ||
| apiKey: process.env.PINECONE_API_KEY, | ||
| environment: process.env.PINECONE_ENVIRONMENT, | ||
| }); | ||
| const index = client.Index(process.env.PINECONE_INDEX); | ||
|
|
||
| const embeddings = new OpenAIEmbeddings(); | ||
| const llm = new OpenAI(); | ||
| const documentContents = "Brief summary of a movie"; | ||
| const vectorStore = await PineconeStore.fromDocuments(docs, embeddings, { | ||
| pineconeIndex: index, | ||
| }); | ||
| const selfQueryRetriever = await SelfQueryRetriever.fromLLM({ | ||
| llm, | ||
| vectorStore, | ||
| documentContents, | ||
| attributeInfo, | ||
| structuredQueryTranslator: new PineconeTranslator(), | ||
| }); | ||
|
|
||
| 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); | ||
| }); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters