mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-15 14:36:54 +00:00
Implemented MongoDB Atlas Self-Query Retriever (#13321)
# Description This PR implements Self-Query Retriever for MongoDB Atlas vector store. I've implemented the comparators and operators that are supported by MongoDB Atlas vector store according to the section titled "Atlas Vector Search Pre-Filter" from https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/. Namely: ``` allowed_comparators = [ Comparator.EQ, Comparator.NE, Comparator.GT, Comparator.GTE, Comparator.LT, Comparator.LTE, Comparator.IN, Comparator.NIN, ] """Subset of allowed logical operators.""" allowed_operators = [ Operator.AND, Operator.OR ] ``` Translations from comparators/operators to MongoDB Atlas filter operators(you can find the syntax in the "Atlas Vector Search Pre-Filter" section from the previous link) are done using the following dictionary: ``` map_dict = { Operator.AND: "$and", Operator.OR: "$or", Comparator.EQ: "$eq", Comparator.NE: "$ne", Comparator.GTE: "$gte", Comparator.LTE: "$lte", Comparator.LT: "$lt", Comparator.GT: "$gt", Comparator.IN: "$in", Comparator.NIN: "$nin", } ``` In visit_structured_query() the filters are passed as "pre_filter" and not "filter" as in the MongoDB link above since langchain's implementation of MongoDB atlas vector store(libs\langchain\langchain\vectorstores\mongodb_atlas.py) in _similarity_search_with_score() sets the "filter" key to have the value of the "pre_filter" argument. ``` params["filter"] = pre_filter ``` Test cases and documentation have also been added. # Issue #11616 # Dependencies No new dependencies have been added. # Documentation I have created the notebook mongodb_atlas_self_query.ipynb outlining the steps to get the self-query mechanism working. I worked closely with [@Farhan-Faisal](https://github.com/Farhan-Faisal) on this PR. --------- Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
@@ -21,6 +21,7 @@ from langchain.retrievers.self_query.dashvector import DashvectorTranslator
|
||||
from langchain.retrievers.self_query.deeplake import DeepLakeTranslator
|
||||
from langchain.retrievers.self_query.elasticsearch import ElasticsearchTranslator
|
||||
from langchain.retrievers.self_query.milvus import MilvusTranslator
|
||||
from langchain.retrievers.self_query.mongodb_atlas import MongoDBAtlasTranslator
|
||||
from langchain.retrievers.self_query.myscale import MyScaleTranslator
|
||||
from langchain.retrievers.self_query.opensearch import OpenSearchTranslator
|
||||
from langchain.retrievers.self_query.pinecone import PineconeTranslator
|
||||
@@ -36,6 +37,7 @@ from langchain.vectorstores import (
|
||||
DeepLake,
|
||||
ElasticsearchStore,
|
||||
Milvus,
|
||||
MongoDBAtlasVectorSearch,
|
||||
MyScale,
|
||||
OpenSearchVectorSearch,
|
||||
Pinecone,
|
||||
@@ -66,6 +68,7 @@ def _get_builtin_translator(vectorstore: VectorStore) -> Visitor:
|
||||
SupabaseVectorStore: SupabaseVectorTranslator,
|
||||
TimescaleVector: TimescaleVectorTranslator,
|
||||
OpenSearchVectorSearch: OpenSearchTranslator,
|
||||
MongoDBAtlasVectorSearch: MongoDBAtlasTranslator,
|
||||
}
|
||||
if isinstance(vectorstore, Qdrant):
|
||||
return QdrantTranslator(metadata_key=vectorstore.metadata_payload_key)
|
||||
|
@@ -0,0 +1,74 @@
|
||||
"""Logic for converting internal query language to a valid MongoDB Atlas query."""
|
||||
from typing import Dict, Tuple, Union
|
||||
|
||||
from langchain.chains.query_constructor.ir import (
|
||||
Comparator,
|
||||
Comparison,
|
||||
Operation,
|
||||
Operator,
|
||||
StructuredQuery,
|
||||
Visitor,
|
||||
)
|
||||
|
||||
MULTIPLE_ARITY_COMPARATORS = [Comparator.IN, Comparator.NIN]
|
||||
|
||||
|
||||
class MongoDBAtlasTranslator(Visitor):
|
||||
"""Translate Mongo internal query language elements to valid filters."""
|
||||
|
||||
"""Subset of allowed logical comparators."""
|
||||
allowed_comparators = [
|
||||
Comparator.EQ,
|
||||
Comparator.NE,
|
||||
Comparator.GT,
|
||||
Comparator.GTE,
|
||||
Comparator.LT,
|
||||
Comparator.LTE,
|
||||
Comparator.IN,
|
||||
Comparator.NIN,
|
||||
]
|
||||
|
||||
"""Subset of allowed logical operators."""
|
||||
allowed_operators = [Operator.AND, Operator.OR]
|
||||
|
||||
## Convert a operator or a comparator to Mongo Query Format
|
||||
def _format_func(self, func: Union[Operator, Comparator]) -> str:
|
||||
self._validate_func(func)
|
||||
map_dict = {
|
||||
Operator.AND: "$and",
|
||||
Operator.OR: "$or",
|
||||
Comparator.EQ: "$eq",
|
||||
Comparator.NE: "$ne",
|
||||
Comparator.GTE: "$gte",
|
||||
Comparator.LTE: "$lte",
|
||||
Comparator.LT: "$lt",
|
||||
Comparator.GT: "$gt",
|
||||
Comparator.IN: "$in",
|
||||
Comparator.NIN: "$nin",
|
||||
}
|
||||
return map_dict[func]
|
||||
|
||||
def visit_operation(self, operation: Operation) -> Dict:
|
||||
args = [arg.accept(self) for arg in operation.arguments]
|
||||
return {self._format_func(operation.operator): args}
|
||||
|
||||
def visit_comparison(self, comparison: Comparison) -> Dict:
|
||||
if comparison.comparator in MULTIPLE_ARITY_COMPARATORS and not isinstance(
|
||||
comparison.value, list
|
||||
):
|
||||
comparison.value = [comparison.value]
|
||||
|
||||
comparator = self._format_func(comparison.comparator)
|
||||
|
||||
attribute = comparison.attribute
|
||||
|
||||
return {attribute: {comparator: comparison.value}}
|
||||
|
||||
def visit_structured_query(
|
||||
self, structured_query: StructuredQuery
|
||||
) -> Tuple[str, dict]:
|
||||
if structured_query.filter is None:
|
||||
kwargs = {}
|
||||
else:
|
||||
kwargs = {"pre_filter": structured_query.filter.accept(self)}
|
||||
return structured_query.query, kwargs
|
@@ -0,0 +1,131 @@
|
||||
from typing import Dict, Tuple
|
||||
|
||||
from langchain.chains.query_constructor.ir import (
|
||||
Comparator,
|
||||
Comparison,
|
||||
Operation,
|
||||
Operator,
|
||||
StructuredQuery,
|
||||
)
|
||||
from langchain.retrievers.self_query.mongodb_atlas import MongoDBAtlasTranslator
|
||||
|
||||
DEFAULT_TRANSLATOR = MongoDBAtlasTranslator()
|
||||
|
||||
|
||||
def test_visit_comparison_lt() -> None:
|
||||
comp = Comparison(comparator=Comparator.LT, attribute="qty", value=20)
|
||||
expected = {"qty": {"$lt": 20}}
|
||||
actual = DEFAULT_TRANSLATOR.visit_comparison(comp)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_comparison_eq() -> None:
|
||||
comp = Comparison(comparator=Comparator.EQ, attribute="qty", value=10)
|
||||
expected = {"qty": {"$eq": 10}}
|
||||
actual = DEFAULT_TRANSLATOR.visit_comparison(comp)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_comparison_ne() -> None:
|
||||
comp = Comparison(comparator=Comparator.NE, attribute="name", value="foo")
|
||||
expected = {"name": {"$ne": "foo"}}
|
||||
actual = DEFAULT_TRANSLATOR.visit_comparison(comp)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_comparison_in() -> None:
|
||||
comp = Comparison(comparator=Comparator.IN, attribute="name", value="foo")
|
||||
expected = {"name": {"$in": ["foo"]}}
|
||||
actual = DEFAULT_TRANSLATOR.visit_comparison(comp)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_comparison_nin() -> None:
|
||||
comp = Comparison(comparator=Comparator.NIN, attribute="name", value="foo")
|
||||
expected = {"name": {"$nin": ["foo"]}}
|
||||
actual = DEFAULT_TRANSLATOR.visit_comparison(comp)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_operation() -> None:
|
||||
op = Operation(
|
||||
operator=Operator.AND,
|
||||
arguments=[
|
||||
Comparison(comparator=Comparator.GTE, attribute="qty", value=10),
|
||||
Comparison(comparator=Comparator.LTE, attribute="qty", value=20),
|
||||
Comparison(comparator=Comparator.EQ, attribute="name", value="foo"),
|
||||
],
|
||||
)
|
||||
expected = {
|
||||
"$and": [
|
||||
{"qty": {"$gte": 10}},
|
||||
{"qty": {"$lte": 20}},
|
||||
{"name": {"$eq": "foo"}},
|
||||
]
|
||||
}
|
||||
actual = DEFAULT_TRANSLATOR.visit_operation(op)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_structured_query_no_filter() -> None:
|
||||
query = "What is the capital of France?"
|
||||
structured_query = StructuredQuery(
|
||||
query=query,
|
||||
filter=None,
|
||||
)
|
||||
expected: Tuple[str, Dict] = (query, {})
|
||||
actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_structured_query_one_attr() -> None:
|
||||
query = "What is the capital of France?"
|
||||
comp = Comparison(comparator=Comparator.IN, attribute="qty", value=[5, 15, 20])
|
||||
structured_query = StructuredQuery(
|
||||
query=query,
|
||||
filter=comp,
|
||||
)
|
||||
expected = (
|
||||
query,
|
||||
{"pre_filter": {"qty": {"$in": [5, 15, 20]}}},
|
||||
)
|
||||
actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query)
|
||||
assert expected == actual
|
||||
|
||||
|
||||
def test_visit_structured_query_deep_nesting() -> None:
|
||||
query = "What is the capital of France?"
|
||||
op = Operation(
|
||||
operator=Operator.AND,
|
||||
arguments=[
|
||||
Comparison(comparator=Comparator.EQ, attribute="name", value="foo"),
|
||||
Operation(
|
||||
operator=Operator.OR,
|
||||
arguments=[
|
||||
Comparison(comparator=Comparator.GT, attribute="qty", value=6),
|
||||
Comparison(
|
||||
comparator=Comparator.NIN,
|
||||
attribute="tags",
|
||||
value=["bar", "foo"],
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
structured_query = StructuredQuery(
|
||||
query=query,
|
||||
filter=op,
|
||||
)
|
||||
expected = (
|
||||
query,
|
||||
{
|
||||
"pre_filter": {
|
||||
"$and": [
|
||||
{"name": {"$eq": "foo"}},
|
||||
{"$or": [{"qty": {"$gt": 6}}, {"tags": {"$nin": ["bar", "foo"]}}]},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query)
|
||||
assert expected == actual
|
Reference in New Issue
Block a user