diff --git a/libs/langchain/langchain/retrievers/self_query/elasticsearch.py b/libs/langchain/langchain/retrievers/self_query/elasticsearch.py index c908faa643f..7c2f7671a5c 100644 --- a/libs/langchain/langchain/retrievers/self_query/elasticsearch.py +++ b/libs/langchain/langchain/retrievers/self_query/elasticsearch.py @@ -61,11 +61,10 @@ class ElasticsearchTranslator(Visitor): ] if is_range_comparator: - return { - "range": { - field: {self._format_func(comparison.comparator): comparison.value} - } - } + value = comparison.value + if isinstance(comparison.value, dict) and "date" in comparison.value: + value = comparison.value["date"] + return {"range": {field: {self._format_func(comparison.comparator): value}}} if comparison.comparator == Comparator.CONTAIN: return { @@ -85,6 +84,10 @@ class ElasticsearchTranslator(Visitor): # we want to use the keyword field field = f"{field}.keyword" if isinstance(comparison.value, str) else field + if isinstance(comparison.value, dict): + if "date" in comparison.value: + comparison.value = comparison.value["date"] + return {self._format_func(comparison.comparator): {field: comparison.value}} def visit_structured_query( diff --git a/libs/langchain/tests/unit_tests/retrievers/self_query/test_elasticsearch.py b/libs/langchain/tests/unit_tests/retrievers/self_query/test_elasticsearch.py index 25c4526c801..519f366f031 100644 --- a/libs/langchain/tests/unit_tests/retrievers/self_query/test_elasticsearch.py +++ b/libs/langchain/tests/unit_tests/retrievers/self_query/test_elasticsearch.py @@ -218,3 +218,98 @@ def test_visit_structured_query_complex() -> None: ) actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query) assert expected == actual + + +def test_visit_structured_query_with_date_range() -> None: + query = "Who was the president of France in 1995?" + operation = Operation( + operator=Operator.AND, + arguments=[ + Comparison(comparator=Comparator.EQ, attribute="foo", value="20"), + Operation( + operator=Operator.AND, + arguments=[ + Comparison( + comparator=Comparator.GTE, + attribute="timestamp", + value={"date": "1995-01-01", "type": "date"}, + ), + Comparison( + comparator=Comparator.LT, + attribute="timestamp", + value={"date": "1996-01-01", "type": "date"}, + ), + ], + ), + ], + ) + structured_query = StructuredQuery(query=query, filter=operation, limit=None) + expected = ( + query, + { + "filter": [ + { + "bool": { + "must": [ + {"term": {"metadata.foo.keyword": "20"}}, + { + "bool": { + "must": [ + { + "range": { + "metadata.timestamp": { + "gte": "1995-01-01" + } + } + }, + { + "range": { + "metadata.timestamp": { + "lt": "1996-01-01" + } + } + }, + ] + } + }, + ] + } + } + ] + }, + ) + actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query) + assert expected == actual + + +def test_visit_structured_query_with_date() -> None: + query = "Who was the president of France on 1st of January 1995?" + operation = Operation( + operator=Operator.AND, + arguments=[ + Comparison(comparator=Comparator.EQ, attribute="foo", value="20"), + Comparison( + comparator=Comparator.EQ, + attribute="timestamp", + value={"date": "1995-01-01", "type": "date"}, + ), + ], + ) + structured_query = StructuredQuery(query=query, filter=operation, limit=None) + expected = ( + query, + { + "filter": [ + { + "bool": { + "must": [ + {"term": {"metadata.foo.keyword": "20"}}, + {"term": {"metadata.timestamp": "1995-01-01"}}, + ] + } + } + ] + }, + ) + actual = DEFAULT_TRANSLATOR.visit_structured_query(structured_query) + assert expected == actual