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