mirror of https://github.com/hwchase17/langchain
community[minor]: Add AstraDB self query retriever (#15738)
- **Description:** this change adds a self-query retriever for AstraDB - **Twitter handle:** cbornet_pull/14969/head^2
parent
fb676d8a9b
commit
15c2b4a47e
@ -0,0 +1,70 @@
|
||||
"""Logic for converting internal query language to a valid AstraDB 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 AstraDBTranslator(Visitor):
|
||||
"""Translate AstraDB 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]
|
||||
|
||||
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)
|
||||
return {comparison.attribute: {comparator: comparison.value}}
|
||||
|
||||
def visit_structured_query(
|
||||
self, structured_query: StructuredQuery
|
||||
) -> Tuple[str, dict]:
|
||||
if structured_query.filter is None:
|
||||
kwargs = {}
|
||||
else:
|
||||
kwargs = {"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.astradb import AstraDBTranslator
|
||||
|
||||
DEFAULT_TRANSLATOR = AstraDBTranslator()
|
||||
|
||||
|
||||
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,
|
||||
{"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,
|
||||
{
|
||||
"filter": {
|
||||
"$and": [
|
||||
{"name": {"$eq": "foo"}},
|
||||
{"$or": [{"qty": {"$gt": 6}}, {"tags": {"$nin": ["bar", "foo"]}}]},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query)
|
||||
assert expected == actual
|
Loading…
Reference in New Issue