-
Notifications
You must be signed in to change notification settings - Fork 29
Expand file tree
/
Copy pathagent_tool.py
More file actions
275 lines (224 loc) · 9.41 KB
/
Copy pathagent_tool.py
File metadata and controls
275 lines (224 loc) · 9.41 KB
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
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
from .utils import DSSTaggableObjectListItem, DSSTaggableObjectSettings, AnyLoc
from .knowledgebank import DSSKnowledgeBank, DSSKnowledgeBankListItem
from .llm_tracing import prepare_query_for_nested_llm_mesh_call
import json
class DSSAgentToolListItem(DSSTaggableObjectListItem):
"""
.. important::
Do not instantiate this class directly, instead use :meth:`dataikuapi.dss.project.DSSProject.list_agent_tools`.
"""
def __init__(self, client, project_key, data):
super(DSSAgentToolListItem, self).__init__(data)
self.project_key = project_key
self.client = client
def to_agent_tool(self):
"""
Convert the current item.
"""
return DSSAgentTool(self.client, self.project_key, self._data["id"], "descriptor" in self._data and self._data["descriptor"] or None)
@property
def id(self):
"""
:returns: The id of the tool.
:rtype: string
"""
return self._data["id"]
@property
def type(self):
"""
:returns: The type of the tool
:rtype: string
"""
return self._data["type"]
@property
def name(self):
"""
:returns: The name of the tool
:rtype: string
"""
return self._data["name"]
class DSSAgentTool(object):
"""
.. important::
Do not instantiate this class directly, instead use :meth:`dataikuapi.dss.project.DSSProject.get_agent_tool`.
"""
def __init__(self, client, project_key, tool_id, descriptor=None):
self.client = client
self.project_key = project_key
self.tool_id = tool_id
self._descriptor = descriptor
@property
def id(self):
"""
:returns: The id of the tool.
:rtype: string
"""
return self.tool_id
def get_descriptor(self, context=None):
"""
Get the descriptor of the tool
:return: a descriptor of the tool
:rtype: dict
"""
if self._descriptor is None:
if context is None:
self._descriptor = self.client._perform_json("GET", "/projects/%s/agents/tools/%s/descriptor" % (self.project_key, self.tool_id))
else:
self._descriptor = self.client._perform_json("POST", "/projects/%s/agents/tools/%s/descriptor" % (self.project_key, self.tool_id), body={"context": context})
return self._descriptor
def get_settings(self):
"""
Get the agent tools' settings
:return: a handle on the tool settings
:rtype: :class:`dataikuapi.dss.agent_tool.DSSAgentToolSettings` or a subclass
"""
settings = self.client._perform_json(
"GET", "/projects/%s/agents/tools/%s" % (self.project_key, self.id))
if settings["type"] == "VectorStoreSearch":
return DSSVectorStoreSearchAgentToolSettings(self, settings)
else:
return DSSAgentToolSettings(self, settings)
def delete(self):
"""
Delete the agent tool
"""
return self.client._perform_empty("DELETE", "/projects/%s/agents/tools/%s" % (self.project_key, self.id))
def as_langchain_structured_tool(self, context = None):
"""
:returns: this tool as a LangChain StructuredTool
:rtype: langchain_core.tools.StructuredTool
"""
from dataikuapi.dss.langchain.tool import convert_to_langchain_structured_tool
return convert_to_langchain_structured_tool(self, context)
def run(self, input, context=None, subtool_name=None, memory_fragment=None, tool_validation_responses=None, tool_validation_requests=None):
"""
Execute a tool call
:param str input: Text input
:param dict context: Additional request context
:param str subtool_name: Name of the sub-tool, if applicable (e.g., for a MCP tool)
:rtype: dict
:returns: The result of running this tool
"""
invocation = {
"toolId" : self.tool_id,
"input" : {
"input" : input
}
}
if tool_validation_responses:
invocation["input"]["toolValidationResponses"] = tool_validation_responses
if tool_validation_requests:
invocation["input"]["toolValidationRequests"] = tool_validation_requests
if memory_fragment:
invocation["input"]["memoryFragment"] = memory_fragment
if subtool_name is not None:
invocation["input"]["subtoolName"] = subtool_name
if context is not None:
invocation["input"]["context"] = context
# Note that 'prepare_query_for_nested_llm_mesh_call' throws an exception when the max LLM mesh stack depth is reached
invocation["input"] = prepare_query_for_nested_llm_mesh_call(invocation["input"])
return self.client._perform_json("POST", "/projects/%s/agents/tools/%s/invocations" % (self.project_key, self.tool_id), body=invocation)
def describe_tool_call(self, input, descriptor, context=None, subtool_name=None):
"""
Get a description for a tool call before it is executed
:return: a string description of the tool call
:rtype: Optional[str]
"""
description_request = {
"input" : {
"input" : input
},
"descriptor": descriptor
}
if subtool_name is not None:
description_request["input"]["subtoolName"] = subtool_name
if context is not None:
description_request["input"]["context"] = context
description_request["input"] = prepare_query_for_nested_llm_mesh_call(description_request["input"])
tool_call_descriptor = self.client._perform_json("POST", "/projects/%s/agents/tools/%s/describe-tool-call" % (self.project_key, self.tool_id), body=description_request)
if tool_call_descriptor is None:
return None
return tool_call_descriptor.get("description", None)
#####################################################
# Creation and Edition - Base Classes
#####################################################
class DSSAgentToolCreator(object):
"""
Helper to create new agent tools
.. important::
Do not instantiate directly, use :meth:`dataikuapi.dss.project.DSSProject.new_agent_tool()` instead.
"""
def __init__(self, project, type, name, id):
self.project = project
self.proto = {
"type": type,
"name": name,
"id":id,
"creationParams" : {}
}
def create(self):
"""
Creates the new agent tool in the project, and return a handle to interact with it.
:rtype: :class:`dataikuapi.dss.agent_tool.DSSAgentTool`
"""
self._finish_creation()
id = self.project.client._perform_json("POST", "/projects/%s/agents/tools" % self.project.project_key, body=self.proto)
return DSSAgentTool(self.project.client, self.project.project_key, id["id"])
def _finish_creation(self):
pass
class DSSAgentToolSettings(DSSTaggableObjectSettings):
def __init__(self, agent_tool, settings):
super(DSSAgentToolSettings, self).__init__(settings)
self.agent_tool = agent_tool
self._settings = settings
def get_raw(self):
"""
:returns: the raw settings dict for this agent tool
:rtype: dict
"""
return self._settings
@property
def params(self):
"""
The parameters of the tool, as a dict. Changes to the dict will be reflected when saving
"""
return self._settings["params"]
def save(self):
"""
Saves the settings of the agent tool
"""
self.agent_tool.client._perform_empty(
"PUT", "/projects/%s/agents/tools/%s" % (self.agent_tool.project_key, self.agent_tool.id), body=self._settings)
#####################################################
# Creation and Edition - Per-type
#####################################################
def _kb_to_loc(context_project_key, kb):
if isinstance(kb, DSSKnowledgeBank):
return AnyLoc(kb.project_key, kb.id)
elif isinstance(kb, DSSKnowledgeBankListItem):
return AnyLoc(kb.project_key, kb.id)
elif isinstance(kb, str):
return AnyLoc.from_ref(context_project_key, kb)
else:
raise Exception("Invalid kb object: %s" % kb)
class DSSVectorStoreSearchAgentToolCreator(DSSAgentToolCreator):
def __init__(self, project, type, name, id):
DSSAgentToolCreator.__init__(self, project, type, name, id)
def with_knowledge_bank(self, kb):
"""
:param kb: Knowledge Bank (object, list item, or identifier) to use in this tool
:type kb: DSSKnowledgeBank | DSSKnowledgeBankListItem | str
"""
loc = _kb_to_loc(self.project.project_key, kb)
self.proto["creationParams"]["knowledgeBankRef"] = loc.to_ref(self.project.project_key)
return self
class DSSVectorStoreSearchAgentToolSettings(DSSAgentToolSettings):
def __init__(self, agent_tool, settings):
DSSAgentToolSettings.__init__(self, agent_tool, settings)
def set_knowledge_bank(self, kb):
"""
:param kb: Knowledge Bank (object, list item, or identifier) to use in this tool
:type kb: DSSKnowledgeBank | DSSKnowledgeBankListItem | str
"""
loc = _kb_to_loc(self.project.project_key, kb)
self.settings["params"]["knowledgeBankRef"] = loc.to_ref(self.agent_tool.project_key)