Skip to content

Agents coa

CoAAgentPack #

Bases: BaseLlamaPack

Chain-of-abstraction Agent Pack.

Parameters:

Name Type Description Default
tools List[BaseTool]

List of tools to use.

required
llm Optional[LLM]

LLM to use. Defaults to gpt-4.

None
Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/base.py
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
class CoAAgentPack(BaseLlamaPack):
    """Chain-of-abstraction Agent Pack.

    Args:
        tools (List[BaseTool]): List of tools to use.
        llm (Optional[LLM]): LLM to use. Defaults to gpt-4.
    """

    def __init__(
        self,
        tools: List[BaseTool],
        llm: Optional[LLM] = None,
        callback_manager: Optional[CallbackManager] = None,
        agent_worker_kwargs: Optional[Dict[str, Any]] = None,
        agent_runner_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Init params."""
        self.llm = llm or Settings.llm
        self.callback_manager = callback_manager or self.llm.callback_manager
        self.agent_worker = CoAAgentWorker.from_tools(
            tools=tools,
            llm=llm,
            verbose=True,
            callback_manager=self.callback_manager,
            **(agent_worker_kwargs or {})
        )
        self.agent = AgentRunner(
            self.agent_worker,
            callback_manager=self.callback_manager,
            **(agent_runner_kwargs or {})
        )

    def get_modules(self) -> Dict[str, Any]:
        """Get modules."""
        return {
            "llm": self.llm,
            "callback_manager": self.callback_manager,
            "agent_worker": self.agent_worker,
            "agent": self.agent,
        }

    def run(self, *args: Any, **kwargs: Any) -> Any:
        """Run the pipeline."""
        return self.agent.chat(*args, **kwargs)

get_modules #

get_modules() -> Dict[str, Any]

Get modules.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/base.py
47
48
49
50
51
52
53
54
def get_modules(self) -> Dict[str, Any]:
    """Get modules."""
    return {
        "llm": self.llm,
        "callback_manager": self.callback_manager,
        "agent_worker": self.agent_worker,
        "agent": self.agent,
    }

run #

run(*args: Any, **kwargs: Any) -> Any

Run the pipeline.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/base.py
56
57
58
def run(self, *args: Any, **kwargs: Any) -> Any:
    """Run the pipeline."""
    return self.agent.chat(*args, **kwargs)

CoAAgentWorker #

Bases: BaseAgentWorker

Chain-of-abstraction Agent Worker.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
 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
class CoAAgentWorker(BaseAgentWorker):
    """Chain-of-abstraction Agent Worker."""

    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":
        """Convenience constructor method from set of of BaseTools (Optional).

        Returns:
            LLMCompilerAgentWorker: the LLMCompilerAgentWorker instance

        """
        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:
        """Initialize step from task."""
        sources: List[ToolOutput] = []
        # temporary memory for new messages
        new_memory = ChatMemoryBuffer.from_defaults()

        # put current history in new memory
        messages = task.memory.get()
        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]:
        """Get tools."""
        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:
        """Run step."""
        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:
        """Run step."""
        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:
        """Run step (async)."""
        return await self._arun_step(step, task)

    @trace_method("run_step")
    def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
        """Run step (stream)."""
        # 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:
        """Run step (async stream)."""
        # 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:
        """Finalize task, after all the steps are completed."""
        # 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()

from_tools classmethod #

from_tools(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

Convenience constructor method from set of of BaseTools (Optional).

Returns:

Name Type Description
LLMCompilerAgentWorker CoAAgentWorker

the LLMCompilerAgentWorker instance

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
 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
@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":
    """Convenience constructor method from set of of BaseTools (Optional).

    Returns:
        LLMCompilerAgentWorker: the LLMCompilerAgentWorker instance

    """
    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,
    )

initialize_step #

initialize_step(task: Task, **kwargs: Any) -> TaskStep

Initialize step from task.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep:
    """Initialize step from task."""
    sources: List[ToolOutput] = []
    # temporary memory for new messages
    new_memory = ChatMemoryBuffer.from_defaults()

    # put current history in new memory
    messages = task.memory.get()
    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": ""},
    )

get_tools #

get_tools(query_str: str) -> List[AsyncBaseTool]

Get tools.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
141
142
143
144
145
146
147
148
def get_tools(self, query_str: str) -> List[AsyncBaseTool]:
    """Get tools."""
    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]

run_step #

run_step(step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput

Run step.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
244
245
246
247
@trace_method("run_step")
def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
    """Run step."""
    return asyncio.run(self.arun_step(step=step, task=task, **kwargs))

arun_step async #

arun_step(step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput

Run step (async).

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
249
250
251
252
253
254
@trace_method("run_step")
async def arun_step(
    self, step: TaskStep, task: Task, **kwargs: Any
) -> TaskStepOutput:
    """Run step (async)."""
    return await self._arun_step(step, task)

stream_step #

stream_step(step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput

Run step (stream).

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
256
257
258
259
260
@trace_method("run_step")
def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
    """Run step (stream)."""
    # Streaming isn't really possible, because we need the full response to know if we are done
    raise NotImplementedError

astream_step async #

astream_step(step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput

Run step (async stream).

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
262
263
264
265
266
267
268
@trace_method("run_step")
async def astream_step(
    self, step: TaskStep, task: Task, **kwargs: Any
) -> TaskStepOutput:
    """Run step (async stream)."""
    # Streaming isn't really possible, because we need the full response to know if we are done
    raise NotImplementedError

finalize_task #

finalize_task(task: Task, **kwargs: Any) -> None

Finalize task, after all the steps are completed.

Source code in llama-index-packs/llama-index-packs-agents-coa/llama_index/packs/agents_coa/step.py
270
271
272
273
274
275
def finalize_task(self, task: Task, **kwargs: Any) -> None:
    """Finalize task, after all the steps are completed."""
    # 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()