Add query parsing unit tests

langchain/chains/query_constructor/base.py

@@ -33,7 +33,7 @@ def parse(self, text: str) -> StructuredQuery:
             parsed = parse_json_markdown(text, expected_keys)
             if len(parsed["query"]) == 0:
                 parsed["query"] = " "
-            if parsed["filter"] == "NO_FILTER":
+            if parsed["filter"] == "NO_FILTER" or not parsed["filter"]:
                 parsed["filter"] = None
             else:
                 parsed["filter"] = self.ast_parse(parsed["filter"])

langchain/chains/query_constructor/parser.py

@@ -20,19 +20,21 @@
 
     func_call: CNAME "(" [args] ")"
 
-    ?value: SIGNED_NUMBER -> number
+    ?value: SIGNED_INT -> int
+        | SIGNED_FLOAT -> float
         | list
         | string
-        | "false" -> false
-        | "true" -> true
+        | ("false" | "False" | "FALSE") -> false
+        | ("true" | "True" | "TRUE") -> true
 
     args: expr ("," expr)*
-    string: ESCAPED_STRING
+    string: /'[^']*'/ | ESCAPED_STRING
     list: "[" [args] "]"
 
     %import common.CNAME
-    %import common.SIGNED_NUMBER
     %import common.ESCAPED_STRING
+    %import common.SIGNED_FLOAT
+    %import common.SIGNED_INT
     %import common.WS
     %ignore WS
 """
@@ -44,7 +46,7 @@ def __init__(
         self,
         *args: Any,
         allowed_comparators: Optional[Sequence[Comparator]] = None,
-        allowed_operators: Optional[Sequence[Operator]],
+        allowed_operators: Optional[Sequence[Operator]] = None,
         **kwargs: Any,
     ):
         super().__init__(*args, **kwargs)
@@ -93,9 +95,14 @@ def true(self) -> bool:
         return True
 
     def list(self, item: Any) -> list:
+        if item is None:
+            return []
         return list(item)
 
-    def number(self, item: Any) -> float:
+    def int(self, item: Any) -> int:
+        return int(item)
+
+    def float(self, item: Any) -> float:
         return float(item)
 
     def string(self, item: Any) -> str:

langchain/chains/query_constructor/prompt.py

@@ -32,7 +32,7 @@
 {{
     "query": "teenager love",
     "filter": "and(or(eq(\\"artist\\", \\"Taylor Swift\\"), eq(\\"artist\\", \\"Katy Perry\\")), \
-lt(\\"length\\", 180), eq(\\"genre\\", \\"pop\gg\"))"
+lt(\\"length\\", 180), eq(\\"genre\\", \\"pop\\"))"
 }}"""
 
 NO_FILTER_ANSWER = """\

tests/unit_tests/chains/query_constructor/test_parser.py

@@ -0,0 +1,116 @@
+"""Test LLM-generated structured query parsing."""
+from typing import Any, Tuple, cast
+
+import pytest
+
+from langchain.chains.query_constructor.ir import (
+    Comparator,
+    Comparison,
+    Operation,
+    Operator,
+)
+from langchain.chains.query_constructor.parser import get_parser
+
+DEFAULT_PARSER = get_parser()
+
+
+@pytest.mark.parametrize("x", ("", "foo", 'foo("bar", "baz")'))
+def test_parse_invalid_grammar(x: str) -> None:
+    with pytest.raises(Exception):
+        DEFAULT_PARSER.parse(x)
+
+
+def test_parse_comparison() -> None:
+    comp = 'gte("foo", 2)'
+    expected = Comparison(comparator=Comparator.GTE, attribute="foo", value=2)
+    for input in (
+        comp,
+        comp.replace('"', "'"),
+        comp.replace(" ", ""),
+        comp.replace(" ", "  "),
+        comp.replace("(", " ("),
+        comp.replace(",", ", "),
+        comp.replace("2", "2.0"),
+    ):
+        actual = DEFAULT_PARSER.parse(input)
+        assert expected == actual
+
+
+def test_parse_operation() -> None:
+    op = 'and(eq("foo", "bar"), lt("baz", 1995.25))'
+    eq = Comparison(comparator=Comparator.EQ, attribute="foo", value="bar")
+    lt = Comparison(comparator=Comparator.LT, attribute="baz", value=1995.25)
+    expected = Operation(operator=Operator.AND, arguments=[eq, lt])
+    for input in (
+        op,
+        op.replace('"', "'"),
+        op.replace(" ", ""),
+        op.replace(" ", "  "),
+        op.replace("(", " ("),
+        op.replace(",", ", "),
+        op.replace("25", "250"),
+    ):
+        actual = DEFAULT_PARSER.parse(input)
+        assert expected == actual
+
+
+def test_parse_nested_operation() -> None:
+    op = 'and(or(eq("a", "b"), eq("a", "c"), eq("a", "d")), not(eq("z", "foo")))'
+    eq1 = Comparison(comparator=Comparator.EQ, attribute="a", value="b")
+    eq2 = Comparison(comparator=Comparator.EQ, attribute="a", value="c")
+    eq3 = Comparison(comparator=Comparator.EQ, attribute="a", value="d")
+    eq4 = Comparison(comparator=Comparator.EQ, attribute="z", value="foo")
+    _not = Operation(operator=Operator.NOT, arguments=[eq4])
+    _or = Operation(operator=Operator.OR, arguments=[eq1, eq2, eq3])
+    expected = Operation(operator=Operator.AND, arguments=[_or, _not])
+    actual = DEFAULT_PARSER.parse(op)
+    assert expected == actual
+
+
+def test_parse_disallowed_comparator() -> None:
+    parser = get_parser(allowed_comparators=[Comparator.EQ])
+    with pytest.raises(ValueError):
+        parser.parse('gt("a", 2)')
+
+
+def test_parse_disallowed_operator() -> None:
+    parser = get_parser(allowed_operators=[Operator.AND])
+    with pytest.raises(ValueError):
+        parser.parse('not(gt("a", 2))')
+
+
+def _test_parse_value(x: Any) -> None:
+    parsed = cast(Comparison, (DEFAULT_PARSER.parse(f'eq("x", {x})')))
+    actual = parsed.value
+    assert actual == x
+
+
+@pytest.mark.parametrize("x", (-1, 0, 1_000_000))
+def test_parse_int_value(x: int) -> None:
+    _test_parse_value(x)
+
+
+@pytest.mark.parametrize("x", (-1.001, 0.00000002, 1_234_567.6543210))
+def test_parse_float_value(x: float) -> None:
+    _test_parse_value(x)
+
+
+@pytest.mark.parametrize("x", ([], [1, "b", "true"]))
+def test_parse_list_value(x: list) -> None:
+    _test_parse_value(x)
+
+
+def test_parse_string_value() -> None:
+    for x in ('""', '" "', '"foo"', "'foo'"):
+        parsed = cast(Comparison, DEFAULT_PARSER.parse(f'eq("x", {x})'))
+        actual = parsed.value
+        assert actual == x[1:-1]
+
+
+def test_parse_bool_value() -> None:
+    for y in ("true", "false"):
+        for x in (y, y.upper(), y.title()):
+            parsed = cast(Comparison, DEFAULT_PARSER.parse(f'eq("x", {x})'))
+            actual = parsed.value
+            expected = y == "true"
+            assert actual == expected