langchain[patch]: Allow OpenSearch Query Translator to correctly work with Date types (#16022)

**Description:**

Fixes an issue where the Date type in an OpenSearch Self Querying
Retriever would fail to generate a valid query

**Issue:**
https://github.com/langchain-ai/langchain/issues/14225
This commit is contained in:
Ryan French 2024-01-20 01:57:18 +00:00 committed by GitHub
parent ffae98d371
commit 3d23a5eb36
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 111 additions and 4 deletions

View File

@ -58,9 +58,23 @@ class OpenSearchTranslator(Visitor):
Comparator.GT,
Comparator.GTE,
]:
if isinstance(comparison.value, dict):
if "date" in comparison.value:
return {
"range": {
field: {self._format_func(comparison.comparator): comparison.value}
field: {
self._format_func(
comparison.comparator
): comparison.value["date"]
}
}
}
else:
return {
"range": {
field: {
self._format_func(comparison.comparator): comparison.value
}
}
}
@ -70,8 +84,13 @@ class OpenSearchTranslator(Visitor):
field: {"value": comparison.value}
}
}
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(
@ -81,4 +100,5 @@ class OpenSearchTranslator(Visitor):
kwargs = {}
else:
kwargs = {"filter": structured_query.filter.accept(self)}
return structured_query.query, kwargs

View File

@ -85,3 +85,90 @@ def test_visit_structured_query() -> 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