Source code for langchain_community.query_constructors.tencentvectordb

from __future__ import annotations

from typing import Optional, Sequence, Tuple

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


class TencentVectorDBTranslator(Visitor):
    """将StructuredQuery翻译为腾讯VectorDB查询。"""

    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,
    ]

    def __init__(self, meta_keys: Optional[Sequence[str]] = None):
        """初始化翻译器。

参数：
    meta_keys：要在查询中使用的元键列表。默认值：[]。
"""
        self.meta_keys = meta_keys or []

    def visit_operation(self, operation: Operation) -> str:
        """访问一个操作节点并返回翻译后的查询。

参数:
    operation: 要访问的操作节点。

返回:
    翻译后的查询。
"""
        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)})"

    def visit_comparison(self, comparison: Comparison) -> str:
        """访问一个比较节点并返回翻译后的查询。

参数:
    comparison: 要访问的比较节点。

返回:
    翻译后的查询。
"""
        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}")

    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        """访问一个结构化查询节点并返回翻译后的查询。

参数:
    structured_query: 要访问的StructuredQuery节点。

返回:
    翻译后的查询和查询参数。
"""
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"expr": structured_query.filter.accept(self)}
        return structured_query.query, kwargs