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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154 | class LoadAndSearchToolSpec(BaseToolSpec):
"""加载和搜索工具。
这个工具可以与其他加载大量信息的工具一起使用。与OndemandLoaderTool相比,这个工具返回两个工具,一个用于将数据检索到索引中,另一个用于允许代理使用自然语言查询字符串搜索检索到的数据。"""
loader_prompt = """
Use this tool to load data from the following function. It must then be read from
the corresponding read_{} function.
{}
"""
# TODO, more general read prompt, not always natural language?
reader_prompt = """
Once data has been loaded from {} it can then be read using a natural
language query from this function.
You are required to pass the natural language query argument when calling this endpoint
Args:
query (str): The natural language query used to retreieve information from the index
"""
def __init__(
self,
tool: FunctionTool,
index_cls: Type[BaseIndex],
index_kwargs: Dict,
metadata: ToolMetadata,
index: Optional[BaseIndex] = None,
) -> None:
"""初始化参数。"""
self._index_cls = index_cls
self._index_kwargs = index_kwargs
self._index = index
self._metadata = metadata
self._tool = tool
if self._metadata.name is None:
raise ValueError("Tool name cannot be None")
self.spec_functions = [
self._metadata.name,
f"read_{self._metadata.name}",
]
self._tool_list = [
FunctionTool.from_defaults(
fn=self.load,
name=self._metadata.name,
description=self.loader_prompt.format(
self._metadata.name, self._metadata.description
),
fn_schema=self._metadata.fn_schema,
),
FunctionTool.from_defaults(
fn=self.read,
name=str(f"read_{self._metadata.name}"),
description=self.reader_prompt.format(metadata.name),
fn_schema=create_schema_from_function("ReadData", self.read),
),
]
@property
def metadata(self) -> ToolMetadata:
return self._metadata
@classmethod
def from_defaults(
cls,
tool: FunctionTool,
index_cls: Optional[Type[BaseIndex]] = None,
index_kwargs: Optional[Dict] = None,
name: Optional[str] = None,
description: Optional[str] = None,
fn_schema: Optional[Type[BaseModel]] = None,
) -> "LoadAndSearchToolSpec":
"""从默认值。"""
index_cls = index_cls or VectorStoreIndex
index_kwargs = index_kwargs or {}
if name is None:
name = tool.metadata.name
if description is None:
description = tool.metadata.description
if fn_schema is None:
fn_schema = tool.metadata.fn_schema
metadata = ToolMetadata(name=name, description=description, fn_schema=fn_schema)
return cls(
tool=tool,
index_cls=index_cls,
index_kwargs=index_kwargs,
metadata=metadata,
)
def to_tool_list(
self,
spec_functions: Optional[List[SPEC_FUNCTION_TYPE]] = None,
func_to_metadata_mapping: Optional[Dict[str, ToolMetadata]] = None,
) -> List[FunctionTool]:
return self._tool_list
def load(self, *args: Any, **kwargs: Any) -> Any:
# Call the wrapped tool and save the result in the index
docs = self._tool(*args, **kwargs).raw_output
# convert to Document if necessary
if isinstance(docs, list):
for i, doc in enumerate(docs):
if not isinstance(doc, Document):
docs[i] = Document(text=str(doc))
elif isinstance(docs, str):
docs = [Document(text=docs)]
elif isinstance(docs, Document):
docs = [docs]
else:
docs = [Document(text=str(docs))]
if self._index:
for doc in docs:
self._index.insert(doc, **self._index_kwargs)
else:
self._index = self._index_cls.from_documents(docs, **self._index_kwargs)
return (
"Content loaded! You can now search the information using read_{}".format(
self._metadata.name
)
)
def read(self, query: str) -> Any:
# Query the index for the result
if not self._index:
return (
"Error: No content has been loaded into the index. "
f"You must call {self._metadata.name} first"
)
query_engine = self._index.as_query_engine()
response = query_engine.query(query)
return str(response)
|