test_parser.py — langchain Source File

Source Code

"""Test LLM-generated structured query parsing."""

from typing import Any, cast

import lark
import pytest
from langchain_core.structured_query import (
    Comparator,
    Comparison,
    Operation,
    Operator,
)

from langchain_classic.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((ValueError, lark.exceptions.UnexpectedToken)):
        DEFAULT_PARSER.parse(x)


def test_parse_comparison() -> None:
    comp = 'gte("foo", 2)'
    expected = Comparison(comparator=Comparator.GTE, attribute="foo", value=2)
    for text in (
        comp,
        comp.replace('"', "'"),
        comp.replace(" ", ""),
        comp.replace(" ", "  "),
        comp.replace("(", " ("),
        comp.replace(",", ", "),
        comp.replace("2", "2.0"),
    ):
        actual = DEFAULT_PARSER.parse(text)
        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 text in (
        op,
        op.replace('"', "'"),
        op.replace(" ", ""),
        op.replace(" ", "  "),
        op.replace("(", " ("),
        op.replace(",", ", "),
        op.replace("25", "250"),
    ):
        actual = DEFAULT_PARSER.parse(text)
        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")))'
// ... (93 more lines)