159 lines
5.3 KiB
JavaScript
159 lines
5.3 KiB
JavaScript
import { isFilterEmpty, castValue, isInt, isFloat, BaseTranslator, Comparators, Operators, } from "@langchain/core/structured_query";
|
|
/**
|
|
* A class that translates or converts `StructuredQuery` to equivalent Qdrant filters.
|
|
* @example
|
|
* ```typescript
|
|
* const selfQueryRetriever = new SelfQueryRetriever({
|
|
* llm: new ChatOpenAI(),
|
|
* vectorStore: new QdrantVectorStore(...),
|
|
* documentContents: "Brief summary of a movie",
|
|
* attributeInfo: [],
|
|
* structuredQueryTranslator: new QdrantTranslator(),
|
|
* });
|
|
*
|
|
* const relevantDocuments = await selfQueryRetriever.getRelevantDocuments(
|
|
* "Which movies are rated higher than 8.5?",
|
|
* );
|
|
* ```
|
|
*/
|
|
export class QdrantTranslator extends BaseTranslator {
|
|
constructor() {
|
|
super(...arguments);
|
|
Object.defineProperty(this, "allowedOperators", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: [Operators.and, Operators.or, Operators.not]
|
|
});
|
|
Object.defineProperty(this, "allowedComparators", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: [
|
|
Comparators.eq,
|
|
Comparators.ne,
|
|
Comparators.lt,
|
|
Comparators.lte,
|
|
Comparators.gt,
|
|
Comparators.gte,
|
|
]
|
|
});
|
|
}
|
|
/**
|
|
* Visits an operation and returns a QdrantFilter.
|
|
* @param operation The operation to visit.
|
|
* @returns A QdrantFilter.
|
|
*/
|
|
visitOperation(operation) {
|
|
const args = operation.args?.map((arg) => arg.accept(this));
|
|
const operator = {
|
|
[Operators.and]: "must",
|
|
[Operators.or]: "should",
|
|
[Operators.not]: "must_not",
|
|
}[operation.operator];
|
|
return {
|
|
[operator]: args,
|
|
};
|
|
}
|
|
/**
|
|
* Visits a comparison and returns a QdrantCondition.
|
|
* The value is casted to the correct type.
|
|
* The attribute is prefixed with "metadata.",
|
|
* since metadata is nested in the Qdrant payload.
|
|
* @param comparison The comparison to visit.
|
|
* @returns A QdrantCondition.
|
|
*/
|
|
visitComparison(comparison) {
|
|
const attribute = `metadata.${comparison.attribute}`;
|
|
const value = castValue(comparison.value);
|
|
if (comparison.comparator === "eq") {
|
|
return {
|
|
key: attribute,
|
|
match: {
|
|
value,
|
|
},
|
|
};
|
|
}
|
|
else if (comparison.comparator === "ne") {
|
|
return {
|
|
key: attribute,
|
|
match: {
|
|
except: [value],
|
|
},
|
|
};
|
|
}
|
|
if (!isInt(value) && !isFloat(value)) {
|
|
throw new Error("Value for gt, gte, lt, lte must be a number");
|
|
}
|
|
// For gt, gte, lt, lte, we need to use the range filter
|
|
return {
|
|
key: attribute,
|
|
range: {
|
|
[comparison.comparator]: value,
|
|
},
|
|
};
|
|
}
|
|
/**
|
|
* Visits a structured query and returns a VisitStructuredQueryOutput.
|
|
* If the query has a filter, it is visited.
|
|
* @param query The structured query to visit.
|
|
* @returns An instance of VisitStructuredQueryOutput.
|
|
*/
|
|
visitStructuredQuery(query) {
|
|
let nextArg = {};
|
|
if (query.filter) {
|
|
nextArg = {
|
|
filter: { must: [query.filter.accept(this)] },
|
|
};
|
|
}
|
|
return nextArg;
|
|
}
|
|
/**
|
|
* Merges two filters into one. If both filters are empty, returns
|
|
* undefined. If one filter is empty or the merge type is 'replace',
|
|
* returns the other filter. If the merge type is 'and' or 'or', returns a
|
|
* new filter with the merged results. Throws an error for unknown merge
|
|
* types.
|
|
* @param defaultFilter The default filter to merge.
|
|
* @param generatedFilter The generated filter to merge.
|
|
* @param mergeType The type of merge to perform. Can be 'and', 'or', or 'replace'. Defaults to 'and'.
|
|
* @param forceDefaultFilter If true, the default filter is always returned if the generated filter is empty. Defaults to false.
|
|
* @returns A merged QdrantFilter, or undefined if both filters are empty.
|
|
*/
|
|
mergeFilters(defaultFilter, generatedFilter, mergeType = "and", forceDefaultFilter = false) {
|
|
if (isFilterEmpty(defaultFilter) && isFilterEmpty(generatedFilter)) {
|
|
return undefined;
|
|
}
|
|
if (isFilterEmpty(defaultFilter) || mergeType === "replace") {
|
|
if (isFilterEmpty(generatedFilter)) {
|
|
return undefined;
|
|
}
|
|
return generatedFilter;
|
|
}
|
|
if (isFilterEmpty(generatedFilter)) {
|
|
if (forceDefaultFilter) {
|
|
return defaultFilter;
|
|
}
|
|
if (mergeType === "and") {
|
|
return undefined;
|
|
}
|
|
return defaultFilter;
|
|
}
|
|
if (mergeType === "and") {
|
|
return {
|
|
must: [defaultFilter, generatedFilter],
|
|
};
|
|
}
|
|
else if (mergeType === "or") {
|
|
return {
|
|
should: [defaultFilter, generatedFilter],
|
|
};
|
|
}
|
|
else {
|
|
throw new Error("Unknown merge type");
|
|
}
|
|
}
|
|
formatFunction() {
|
|
throw new Error("Not implemented");
|
|
}
|
|
}
|