Skip to content

Commit

Permalink
Fix self-query retriever tracing (#6515)
Browse files Browse the repository at this point in the history
  • Loading branch information
jacoblee93 authored Aug 13, 2024
1 parent 54608ba commit caf620b
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 16 deletions.
2 changes: 1 addition & 1 deletion langchain/src/retrievers/self_query/functional.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
export {
FunctionFilter,
type FunctionFilter,
FunctionalTranslator,
} from "@langchain/core/structured_query";
17 changes: 10 additions & 7 deletions langchain/src/retrievers/self_query/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,10 @@ export class SelfQueryRetriever<T extends VectorStore>
): Promise<Document<Record<string, unknown>>[]> {
const generatedStructuredQuery = await this.queryConstructor.invoke(
{ query },
runManager?.getChild("query_constructor")
{
callbacks: runManager?.getChild("query_constructor"),
runName: "query_constructor",
}
);

const nextArg = this.structuredQueryTranslator.visitStructuredQuery(
Expand All @@ -122,12 +125,12 @@ export class SelfQueryRetriever<T extends VectorStore>
myQuery = generatedQuery;
}

return this.vectorStore.similaritySearch(
myQuery,
this.searchParams?.k,
filter,
runManager?.getChild("vectorstore")
);
return this.vectorStore
.asRetriever({
k: this.searchParams?.k,
filter,
})
.invoke(myQuery, { callbacks: runManager?.getChild("retriever") });
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { test } from "@jest/globals";
import { Document } from "@langchain/core/documents";
import { OpenAIEmbeddings, OpenAI } from "@langchain/openai";
import { OpenAIEmbeddings, ChatOpenAI } from "@langchain/openai";
import { AttributeInfo } from "../../../chains/query_constructor/index.js";
import { SelfQueryRetriever } from "../index.js";
import { FunctionalTranslator } from "../functional.js";
Expand Down Expand Up @@ -73,8 +73,8 @@ test("Memory Vector Store Self Query Retriever Test", async () => {
];

const embeddings = new OpenAIEmbeddings();
const llm = new OpenAI({
modelName: "gpt-3.5-turbo",
const llm = new ChatOpenAI({
model: "gpt-4o-mini",
});
const documentContents = "Brief summary of a movie";
const vectorStore = await MemoryVectorStore.fromDocuments(docs, embeddings);
Expand Down Expand Up @@ -203,8 +203,8 @@ test("Memory Vector Store Self Query Retriever Test With Default Filter Or Merge
];

const embeddings = new OpenAIEmbeddings();
const llm = new OpenAI({
modelName: "gpt-3.5-turbo",
const llm = new ChatOpenAI({
model: "gpt-4o-mini",
});
const documentContents = "Brief summary of a movie";
const vectorStore = await MemoryVectorStore.fromDocuments(docs, embeddings);
Expand Down Expand Up @@ -344,8 +344,8 @@ test("Memory Vector Store Self Query Retriever Test With Default Filter And Merg
];

const embeddings = new OpenAIEmbeddings();
const llm = new OpenAI({
modelName: "gpt-3.5-turbo",
const llm = new ChatOpenAI({
model: "gpt-4o-mini",
});
const documentContents = "Brief summary of a movie";
const vectorStore = await MemoryVectorStore.fromDocuments(docs, embeddings);
Expand Down Expand Up @@ -384,5 +384,5 @@ test("Memory Vector Store Self Query Retriever Test With Default Filter And Merg
"Awawawa au au au wawawawa hello?"
);
// console.log(query4);
expect(query4.length).toEqual(0); // this one should return documents since default filter takes over
expect(query4.length).toEqual(7); // this one should return documents since default filter takes over
});

0 comments on commit caf620b

Please sign in to comment.