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 | class CoAAgentWorker(BaseAgentWorker):
"""链式抽象 代理 工作者。"""
def __init__(
self,
llm: LLM,
reasoning_prompt_template: str,
refine_reasoning_prompt_template: str,
output_parser: BaseOutputParser,
tools: Optional[Sequence[BaseTool]] = None,
tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
callback_manager: Optional[CallbackManager] = None,
verbose: bool = False,
) -> None:
self.llm = llm
self.callback_manager = callback_manager or llm.callback_manager
if tools is None and tool_retriever is None:
raise ValueError("Either tools or tool_retriever must be provided.")
self.tools = tools
self.tool_retriever = tool_retriever
self.reasoning_prompt_template = reasoning_prompt_template
self.refine_reasoning_prompt_template = refine_reasoning_prompt_template
self.output_parser = output_parser
self.verbose = verbose
@classmethod
def from_tools(
cls,
tools: Optional[Sequence[BaseTool]] = None,
tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
llm: Optional[LLM] = None,
reasoning_prompt_template: Optional[str] = None,
refine_reasoning_prompt_template: Optional[str] = None,
output_parser: Optional[BaseOutputParser] = None,
callback_manager: Optional[CallbackManager] = None,
verbose: bool = False,
**kwargs: Any,
) -> "CoAAgentWorker":
"""方便的构造方法，从一组BaseTools中选择（可选）。
返回：
LLMCompilerAgentWorker：LLMCompilerAgentWorker实例
"""
llm = llm or Settings.llm
if callback_manager is not None:
llm.callback_manager = callback_manager
reasoning_prompt_template = (
reasoning_prompt_template or REASONING_PROMPT_TEMPALTE
)
refine_reasoning_prompt_template = (
refine_reasoning_prompt_template or REFINE_REASONING_PROMPT_TEMPALTE
)
output_parser = output_parser or ChainOfAbstractionParser(verbose=verbose)
return cls(
llm,
reasoning_prompt_template,
refine_reasoning_prompt_template,
output_parser,
tools=tools,
tool_retriever=tool_retriever,
callback_manager=callback_manager,
verbose=verbose,
)
def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep:
"""从任务中初始化步骤。"""
sources: List[ToolOutput] = []
# temporary memory for new messages
new_memory = ChatMemoryBuffer.from_defaults()
# put current history in new memory
messages = task.memory.get(input=task.input)
for message in messages:
new_memory.put(message)
# initialize task state
task_state = {
"sources": sources,
"new_memory": new_memory,
}
task.extra_state.update(task_state)
return TaskStep(
task_id=task.task_id,
step_id=str(uuid.uuid4()),
input=task.input,
step_state={"prev_reasoning": ""},
)
def get_tools(self, query_str: str) -> List[AsyncBaseTool]:
"""获取工具。"""
if self.tool_retriever:
tools = self.tool_retriever.retrieve(query_str)
else:
tools = self.tools
return [adapt_to_async_tool(t) for t in tools]
async def _arun_step(
self,
step: TaskStep,
task: Task,
) -> TaskStepOutput:
"""运行步骤。"""
tools = self.get_tools(task.input)
tools_by_name = {tool.metadata.name: tool for tool in tools}
tools_strs = []
for tool in tools:
if isinstance(tool, FunctionTool):
description = tool.metadata.description
# remove function def, we will make our own
if "def " in description:
description = "\n".join(description.split("\n")[1:])
else:
description = tool.metadata.description
tool_str = json_schema_to_python(
tool.metadata.fn_schema_str, tool.metadata.name, description=description
)
tools_strs.append(tool_str)
prev_reasoning = step.step_state.get("prev_reasoning", "")
# show available functions if first step
if self.verbose and not prev_reasoning:
print(f"==== Available Parsed Functions ====")
for tool_str in tools_strs:
print(tool_str)
if not prev_reasoning:
# get the reasoning prompt
reasoning_prompt = self.reasoning_prompt_template.format(
functions="\n".join(tools_strs), question=step.input
)
else:
# get the refine reasoning prompt
reasoning_prompt = self.refine_reasoning_prompt_template.format(
question=step.input, prev_reasoning=prev_reasoning
)
messages = task.extra_state["new_memory"].get()
reasoning_message = ChatMessage(role="user", content=reasoning_prompt)
messages.append(reasoning_message)
# run the reasoning prompt
response = await self.llm.achat(messages)
# print the chain of abstraction if first step
if self.verbose and not prev_reasoning:
print(f"==== Generated Chain of Abstraction ====")
print(str(response.message.content))
# parse the output, run functions
parsed_response, tool_sources = await self.output_parser.aparse(
response.message.content, tools_by_name
)
if len(tool_sources) == 0 or prev_reasoning:
is_done = True
new_steps = []
# only add to memory when we are done
task.extra_state["new_memory"].put(
ChatMessage(role="user", content=task.input)
)
task.extra_state["new_memory"].put(
ChatMessage(role="assistant", content=parsed_response)
)
else:
is_done = False
new_steps = [
TaskStep(
task_id=task.task_id,
step_id=str(uuid.uuid4()),
input=task.input,
step_state={
"prev_reasoning": parsed_response,
},
)
]
agent_response = AgentChatResponse(
response=parsed_response, sources=tool_sources
)
return TaskStepOutput(
output=agent_response,
task_step=step,
is_last=is_done,
next_steps=new_steps,
)
@trace_method("run_step")
def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
"""运行步骤。"""
return asyncio.run(self.arun_step(step=step, task=task, **kwargs))
@trace_method("run_step")
async def arun_step(
self, step: TaskStep, task: Task, **kwargs: Any
) -> TaskStepOutput:
"""运行步骤（异步）。"""
return await self._arun_step(step, task)
@trace_method("run_step")
def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
"""运行步骤（流式）。"""
# Streaming isn't really possible, because we need the full response to know if we are done
raise NotImplementedError
@trace_method("run_step")
async def astream_step(
self, step: TaskStep, task: Task, **kwargs: Any
) -> TaskStepOutput:
"""运行步骤（异步流）。"""
# Streaming isn't really possible, because we need the full response to know if we are done
raise NotImplementedError
def finalize_task(self, task: Task, **kwargs: Any) -> None:
"""完成任务，在所有步骤都完成之后。"""
# add new messages to memory
task.memory.put_messages(task.extra_state["new_memory"].get_all())
# reset new memory
task.extra_state["new_memory"].reset()
|