Spaces:
Runtime error
Runtime error
File size: 4,139 Bytes
ed4d993 |
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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
# flake8: noqa
"""Tools for working with JSON specs."""
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]]:
"""Parse input of the form data["key1"][0]["key2"] into a list of keys."""
_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):
"""Base class for JSON spec."""
dict_: Dict
max_value_length: int = 200
@classmethod
def from_file(cls, path: Path) -> JsonSpec:
"""Create a JsonSpec from a file."""
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:
"""Return the keys of the dict at the given path.
Args:
text: Python representation of the path to the dict (e.g. 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:
"""Return the value of the dict at the given path.
Args:
text: Python representation of the path to the dict (e.g. 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):
"""Tool for listing keys in a JSON spec."""
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):
"""Tool for getting a value in a JSON spec."""
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)
|