Spaces:
Runtime error
Runtime error
File size: 3,283 Bytes
129cd69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
from __future__ import annotations
from typing import Any, Tuple
from langchain.chains.query_constructor.ir import (
Comparator,
Comparison,
Operation,
Operator,
StructuredQuery,
Visitor,
)
from langchain.vectorstores.redis import Redis
from langchain.vectorstores.redis.filters import (
RedisFilterExpression,
RedisFilterField,
RedisFilterOperator,
RedisNum,
RedisTag,
RedisText,
)
from langchain.vectorstores.redis.schema import RedisModel
_COMPARATOR_TO_BUILTIN_METHOD = {
Comparator.EQ: "__eq__",
Comparator.NE: "__ne__",
Comparator.LT: "__lt__",
Comparator.GT: "__gt__",
Comparator.LTE: "__le__",
Comparator.GTE: "__ge__",
Comparator.CONTAIN: "__eq__",
Comparator.LIKE: "__mod__",
}
class RedisTranslator(Visitor):
"""Translate"""
allowed_comparators = (
Comparator.EQ,
Comparator.NE,
Comparator.LT,
Comparator.LTE,
Comparator.GT,
Comparator.GTE,
Comparator.CONTAIN,
Comparator.LIKE,
)
"""Subset of allowed logical comparators."""
allowed_operators = (Operator.AND, Operator.OR)
"""Subset of allowed logical operators."""
def __init__(self, schema: RedisModel) -> None:
self._schema = schema
def _attribute_to_filter_field(self, attribute: str) -> RedisFilterField:
if attribute in [tf.name for tf in self._schema.text]:
return RedisText(attribute)
elif attribute in [tf.name for tf in self._schema.tag or []]:
return RedisTag(attribute)
elif attribute in [tf.name for tf in self._schema.numeric or []]:
return RedisNum(attribute)
else:
raise ValueError(
f"Invalid attribute {attribute} not in vector store schema. Schema is:"
f"\n{self._schema.as_dict()}"
)
def visit_comparison(self, comparison: Comparison) -> RedisFilterExpression:
filter_field = self._attribute_to_filter_field(comparison.attribute)
comparison_method = _COMPARATOR_TO_BUILTIN_METHOD[comparison.comparator]
return getattr(filter_field, comparison_method)(comparison.value)
def visit_operation(self, operation: Operation) -> Any:
left = operation.arguments[0].accept(self)
if len(operation.arguments) > 2:
right = self.visit_operation(
Operation(
operator=operation.operator, arguments=operation.arguments[1:]
)
)
else:
right = operation.arguments[1].accept(self)
redis_operator = (
RedisFilterOperator.OR
if operation.operator == Operator.OR
else RedisFilterOperator.AND
)
return RedisFilterExpression(operator=redis_operator, left=left, right=right)
def visit_structured_query(
self, structured_query: StructuredQuery
) -> Tuple[str, dict]:
if structured_query.filter is None:
kwargs = {}
else:
kwargs = {"filter": structured_query.filter.accept(self)}
return structured_query.query, kwargs
@classmethod
def from_vectorstore(cls, vectorstore: Redis) -> RedisTranslator:
return cls(vectorstore._schema)
|