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
155
156
157
158
159 | class LoadAndSearchToolSpec(BaseToolSpec):
"""Load and Search Tool.
This tool can be used with other tools that load large amounts of
information. Compared to OndemandLoaderTool this returns two tools,
one to retrieve data to an index and another to allow the Agent to search
the retrieved data with a natural language query string.
"""
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:
"""Init params."""
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":
"""From defaults."""
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)
|