Source code for langchain.retrievers.self_query.tencentvectordb

from __future__ import annotations

from typing import Optional, Sequence, Tuple

from langchain_core.structured_query import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)


[docs]class TencentVectorDBTranslator(Visitor): """Translate StructuredQuery to Tencent VectorDB query.""" COMPARATOR_MAP = { Comparator.EQ: "=", Comparator.NE: "!=", Comparator.GT: ">", Comparator.GTE: ">=", Comparator.LT: "<", Comparator.LTE: "<=", Comparator.IN: "in", Comparator.NIN: "not in", } allowed_comparators: Optional[Sequence[Comparator]] = list(COMPARATOR_MAP.keys()) allowed_operators: Optional[Sequence[Operator]] = [ Operator.AND, Operator.OR, Operator.NOT, ]
[docs] def __init__(self, meta_keys: Optional[Sequence[str]] = None): """Initialize the translator. Args: meta_keys: List of meta keys to be used in the query. Default: []. """ self.meta_keys = meta_keys or []
[docs] def visit_operation(self, operation: Operation) -> str: """Visit an operation node and return the translated query. Args: operation: Operation node to be visited. Returns: Translated query. """ if operation.operator in (Operator.AND, Operator.OR): ret = f" {operation.operator.value} ".join( [arg.accept(self) for arg in operation.arguments] ) if operation.operator == Operator.OR: ret = f"({ret})" return ret else: return f"not ({operation.arguments[0].accept(self)})"
[docs] def visit_comparison(self, comparison: Comparison) -> str: """Visit a comparison node and return the translated query. Args: comparison: Comparison node to be visited. Returns: Translated query. """ if self.meta_keys and comparison.attribute not in self.meta_keys: raise ValueError( f"Expr Filtering found Unsupported attribute: {comparison.attribute}" ) if comparison.comparator in self.COMPARATOR_MAP: if comparison.comparator in [Comparator.IN, Comparator.NIN]: value = map( lambda x: f'"{x}"' if isinstance(x, str) else x, comparison.value ) return ( f"{comparison.attribute}" f" {self.COMPARATOR_MAP[comparison.comparator]} " f"({', '.join(value)})" ) if isinstance(comparison.value, str): return ( f"{comparison.attribute} " f"{self.COMPARATOR_MAP[comparison.comparator]}" f' "{comparison.value}"' ) return ( f"{comparison.attribute}" f" {self.COMPARATOR_MAP[comparison.comparator]} " f"{comparison.value}" ) else: raise ValueError(f"Unsupported comparator {comparison.comparator}")
[docs] def visit_structured_query( self, structured_query: StructuredQuery ) -> Tuple[str, dict]: """Visit a structured query node and return the translated query. Args: structured_query: StructuredQuery node to be visited. Returns: Translated query and query kwargs. """ if structured_query.filter is None: kwargs = {} else: kwargs = {"expr": structured_query.filter.accept(self)} return structured_query.query, kwargs