Source code for langchain_community.tools.json.tool

# flake8: noqa
"""用于处理JSON规范的工具。"""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Dict, List, Optional, Union

from langchain_core.pydantic_v1 import BaseModel

from langchain_core.callbacks import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)
from langchain_core.tools import BaseTool


def _parse_input(text: str) -> List[Union[str, int]]:
    """将形式为data["key1"][0]["key2"]的输入解析为键列表。"""
    _res = re.findall(r"\[.*?]", text)
    # strip the brackets and quotes, convert to int if possible
    res = [i[1:-1].replace('"', "").replace("'", "") for i in _res]
    res = [int(i) if i.isdigit() else i for i in res]
    return res


class JsonSpec(BaseModel):
    """JSON规范的基类。"""

    dict_: Dict
    max_value_length: int = 200

    @classmethod
    def from_file(cls, path: Path) -> JsonSpec:
        """从文件中创建一个JsonSpec。"""
        if not path.exists():
            raise FileNotFoundError(f"File not found: {path}")
        dict_ = json.loads(path.read_text())
        return cls(dict_=dict_)

    def keys(self, text: str) -> str:
        """返回给定路径上字典的键。

参数：
    text：字典路径的Python表示（例如 data["key1"][0]["key2"]）。
"""
        try:
            items = _parse_input(text)
            val = self.dict_
            for i in items:
                if i:
                    val = val[i]
            if not isinstance(val, dict):
                raise ValueError(
                    f"Value at path `{text}` is not a dict, get the value directly."
                )
            return str(list(val.keys()))
        except Exception as e:
            return repr(e)

    def value(self, text: str) -> str:
        """返回给定路径的字典的值。

参数：
    text：字典路径的Python表示（例如 data["key1"][0]["key2"]）。
"""
        try:
            items = _parse_input(text)
            val = self.dict_
            for i in items:
                val = val[i]

            if isinstance(val, dict) and len(str(val)) > self.max_value_length:
                return "Value is a large dictionary, should explore its keys directly"
            str_val = str(val)
            if len(str_val) > self.max_value_length:
                str_val = str_val[: self.max_value_length] + "..."
            return str_val
        except Exception as e:
            return repr(e)


class JsonListKeysTool(BaseTool):
    """用于列出JSON规范中键的工具。"""

    name: str = "json_spec_list_keys"
    description: str = """
    Can be used to list all keys at a given path. 
    Before calling this you should be SURE that the path to this exists.
    The input is a text representation of the path to the dict in Python syntax (e.g. data["key1"][0]["key2"]).
    """
    spec: JsonSpec

    def _run(
        self,
        tool_input: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return self.spec.keys(tool_input)

    async def _arun(
        self,
        tool_input: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return self._run(tool_input)


class JsonGetValueTool(BaseTool):
    """用于从JSON规范中获取值的工具。"""

    name: str = "json_spec_get_value"
    description: str = """
    Can be used to see value in string format at a given path.
    Before calling this you should be SURE that the path to this exists.
    The input is a text representation of the path to the dict in Python syntax (e.g. data["key1"][0]["key2"]).
    """
    spec: JsonSpec

    def _run(
        self,
        tool_input: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        return self.spec.value(tool_input)

    async def _arun(
        self,
        tool_input: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        return self._run(tool_input)