langchain[patch]: Ensure that the Elasticsearch Query Translator functions accurately w… (#17044)

Description:
Addresses a problem where the Date type within an Elasticsearch
SelfQueryRetriever would encounter difficulties in generating a valid
query.

Issue: #17042

---------

Co-authored-by: Max Jakob <max.jakob@elastic.co>
Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Taha Khabouss 2024-02-13 19:54:24 +01:00 committed by GitHub
parent e5c76f9dbd
commit 15baffc484
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 103 additions and 5 deletions

View File

@ -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(

View File

@ -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