# Table of Contents
- [LlamaIndex - LlamaIndex](#llamaindex-llamaindex)
- [API Reference - LlamaIndex](#api-reference-llamaindex)
- [Argilla - LlamaIndex](#argilla-llamaindex)
- [Docs - LlamaIndex](#docs-llamaindex)
- [Coa - LlamaIndex](#coa-llamaindex)
- [Core Agent Classes - LlamaIndex](#core-agent-classes-llamaindex)
- [Arize phoenix - LlamaIndex](#arize-phoenix-llamaindex)
- [Langfuse - LlamaIndex](#langfuse-llamaindex)
- [Agentops - LlamaIndex](#agentops-llamaindex)
- [Opik - LlamaIndex](#opik-llamaindex)
- [Token counter - LlamaIndex](#token-counter-llamaindex)
- [Code - LlamaIndex](#code-llamaindex)
- [Openai - LlamaIndex](#openai-llamaindex)
---
# LlamaIndex - LlamaIndex
Welcome to LlamaIndex 🦙 
=======================================================================
LlamaIndex is a framework for building context-augmented generative AI applications with [LLMs](https://en.wikipedia.org/wiki/Large_language_model)
including [agents](./understanding/agent/basic_agent/)
and [workflows](./understanding/workflows/)
.
* [Introduction](#introduction)
What is context augmentation? What are agents and workflows? How does LlamaIndex help build them?
* [Use cases](#use-cases)
What kind of apps can you build with LlamaIndex? Who should use it?
* [Getting started](#getting-started)
Get started in Python or TypeScript in just 5 lines of code!
* [LlamaCloud](#llamacloud)
Managed services for LlamaIndex including [LlamaParse](https://docs.cloud.llamaindex.ai/llamaparse/getting_started)
, the world's best document parser.
* [Community](#community)
Get help and meet collaborators on Discord, Twitter, LinkedIn, and learn how to contribute to the project.
* [Related projects](#related-projects)
Check out our library of connectors, readers, and other integrations at [LlamaHub](https://llamahub.ai)
as well as demos and starter apps like [create-llama](https://www.npmjs.com/package/create-llama)
.
Introduction[#](#introduction "Permanent link")
------------------------------------------------
### What is context augmentation?[#](#what-is-context-augmentation "Permanent link")
LLMs offer a natural language interface between humans and data. LLMs come pre-trained on huge amounts of publicly available data, but they are not trained on **your** data. Your data may be private or specific to the problem you're trying to solve. It's behind APIs, in SQL databases, or trapped in PDFs and slide decks.
Context augmentation makes your data available to the LLM to solve the problem at hand. LlamaIndex provides the tools to build any of context-augmentation use case, from prototype to production. Our tools allow you to ingest, parse, index and process your data and quickly implement complex query workflows combining data access with LLM prompting.
The most popular example of context-augmentation is [Retrieval-Augmented Generation or RAG](getting_started/concepts/)
, which combines context with LLMs at inference time.
### What are agents?[#](#what-are-agents "Permanent link")
[Agents](./understanding/agent/basic_agent/)
are LLM-powered knowledge assistants that use tools to perform tasks like research, data extraction, and more. Agents range from simple question-answering to being able to sense, decide and take actions in order to complete tasks.
LlamaIndex provides a framework for building agents including the ability to use RAG pipelines as one of many tools to complete a task.
### What are workflows?[#](#what-are-workflows "Permanent link")
[Workflows](./understanding/workflows/)
are multi-step processes that combine one or more agents, data connectors, and other tools to complete a task. They are event-driven software that allows you to combine RAG data sources and multiple agents to create a complex application that can perform a wide variety of tasks with reflection, error-correction, and other hallmarks of advanced LLM applications. You can then [deploy these agentic workflows](./module_guides/workflow/deployment.md)
as production microservices.
### LlamaIndex is the framework for Context-Augmented LLM Applications[#](#llamaindex-is-the-framework-for-context-augmented-llm-applications "Permanent link")
LlamaIndex imposes no restriction on how you use LLMs. You can use LLMs as auto-complete, chatbots, agents, and more. It just makes using them easier. We provide tools like:
* **Data connectors** ingest your existing data from their native source and format. These could be APIs, PDFs, SQL, and (much) more.
* **Data indexes** structure your data in intermediate representations that are easy and performant for LLMs to consume.
* **Engines** provide natural language access to your data. For example:
* Query engines are powerful interfaces for question-answering (e.g. a RAG flow).
* Chat engines are conversational interfaces for multi-message, "back and forth" interactions with your data.
* **Agents** are LLM-powered knowledge workers augmented by tools, from simple helper functions to API integrations and more.
* **Observability/Evaluation** integrations that enable you to rigorously experiment, evaluate, and monitor your app in a virtuous cycle.
* **Workflows** allow you to combine all of the above into an event-driven system far more flexible than other, graph-based approaches.
Use cases[#](#use-cases "Permanent link")
------------------------------------------
Some popular use cases for LlamaIndex and context augmentation in general include:
* [Question-Answering](use_cases/q_and_a/)
(Retrieval-Augmented Generation aka RAG)
* [Chatbots](use_cases/chatbots/)
* [Document Understanding and Data Extraction](use_cases/extraction/)
* [Autonomous Agents](use_cases/agents/)
that can perform research and take actions
* [Multi-modal applications](use_cases/multimodal/)
that combine text, images, and other data types
* [Fine-tuning](use_cases/fine_tuning/)
models on data to improve performance
Check out our [use cases](use_cases/)
documentation for more examples and links to tutorials.
### 👨👩👧👦 Who is LlamaIndex for?[#](#who-is-llamaindex-for "Permanent link")
LlamaIndex provides tools for beginners, advanced users, and everyone in between.
Our high-level API allows beginner users to use LlamaIndex to ingest and query their data in 5 lines of code.
For more complex applications, our lower-level APIs allow advanced users to customize and extend any module -- data connectors, indices, retrievers, query engines, and reranking modules -- to fit their needs.
Getting Started[#](#getting-started "Permanent link")
------------------------------------------------------
LlamaIndex is available in Python (these docs) and [Typescript](https://ts.llamaindex.ai/)
. If you're not sure where to start, we recommend reading [how to read these docs](getting_started/reading/)
which will point you to the right place based on your experience level.
### 30 second quickstart[#](#30-second-quickstart "Permanent link")
Set an environment variable called `OPENAI_API_KEY` with an [OpenAI API key](https://platform.openai.com/api-keys)
. Install the Python library:
`pip install llama-index`
Put some documents in a folder called `data`, then ask questions about them with our famous 5-line starter:
`from llama_index.core import VectorStoreIndex, SimpleDirectoryReader documents = SimpleDirectoryReader("data").load_data() index = VectorStoreIndex.from_documents(documents) query_engine = index.as_query_engine() response = query_engine.query("Some question about the data should go here") print(response)`
If any part of this trips you up, don't worry! Check out our more comprehensive starter tutorials using [remote APIs like OpenAI](getting_started/starter_example/)
or [any model that runs on your laptop](getting_started/starter_example_local/)
.
LlamaCloud[#](#llamacloud "Permanent link")
--------------------------------------------
If you're an enterprise developer, check out [**LlamaCloud**](https://llamaindex.ai/enterprise)
. It is an end-to-end managed service for data parsing, ingestion, indexing, and retrieval, allowing you to get production-quality data for your production LLM application. It's available both hosted on our servers or as a self-hosted solution.
### LlamaParse[#](#llamaparse "Permanent link")
LlamaParse is our state-of-the-art document parsing solution. It's available as part of LlamaCloud and also available as a self-serve API. You can [sign up](https://cloud.llamaindex.ai/)
and parse up to 1000 pages/day for free, or enter a credit card for unlimited parsing. [Learn more](https://llamaindex.ai/enterprise)
.
Community[#](#community "Permanent link")
------------------------------------------
Need help? Have a feature suggestion? Join the LlamaIndex community:
* [Twitter](https://twitter.com/llama_index)
* [Discord](https://discord.gg/dGcwcsnxhU)
* [LinkedIn](https://www.linkedin.com/company/llamaindex/)
### Getting the library[#](#getting-the-library "Permanent link")
* LlamaIndex Python
* [LlamaIndex Python Github](https://github.com/run-llama/llama_index)
* [Python Docs](https://docs.llamaindex.ai/)
(what you're reading now)
* [LlamaIndex on PyPi](https://pypi.org/project/llama-index/)
* LlamaIndex.TS (Typescript/Javascript package):
* [LlamaIndex.TS Github](https://github.com/run-llama/LlamaIndexTS)
* [TypeScript Docs](https://ts.llamaindex.ai/)
* [LlamaIndex.TS on npm](https://www.npmjs.com/package/llamaindex)
### Contributing[#](#contributing "Permanent link")
We are open-source and always welcome contributions to the project! Check out our [contributing guide](CONTRIBUTING/)
for full details on how to extend the core library or add an integration to a third party like an LLM, a vector store, an agent tool and more.
LlamaIndex Ecosystem[#](#llamaindex-ecosystem "Permanent link")
----------------------------------------------------------------
There's more to the LlamaIndex universe! Check out some of our other projects:
* [llama\_deploy](https://github.com/run-llama/llama_deploy)
| Deploy your agentic workflows as production microservices
* [LlamaHub](https://llamahub.ai)
| A large (and growing!) collection of custom data connectors
* [SEC Insights](https://secinsights.ai)
| A LlamaIndex-powered application for financial research
* [create-llama](https://www.npmjs.com/package/create-llama)
| A CLI tool to quickly scaffold LlamaIndex projects
Back to top

---
# API Reference - LlamaIndex
API Reference[#](#api-reference "Permanent link")
==================================================
LlamaIndex provides thorough documentation of modules and integrations used in the framework.
Use the navigation or search to find the classes you are interested in!
Back to top

---
# Argilla - LlamaIndex
Argilla
=======
argilla\_callback\_handler [#](#llama_index.callbacks.argilla.argilla_callback_handler "Permanent link")
---------------------------------------------------------------------------------------------------------
`argilla_callback_handler(**kwargs: Any) -> [BaseCallbackHandler](../#llama_index.core.callbacks.base_handler.BaseCallbackHandler "llama_index.core.callbacks.base_handler.BaseCallbackHandler")`
Source code in `llama-index-integrations/callbacks/llama-index-callbacks-argilla/llama_index/callbacks/argilla/base.py`
| | |
| --- | --- |
| 6
7
8
9
10
11
12 | ``def argilla_callback_handler(**kwargs: Any) -> BaseCallbackHandler: try: # lazy import from argilla_llama_index import ArgillaCallbackHandler except ImportError: raise ImportError("Please install Argilla with `pip install argilla`") return ArgillaCallbackHandler(**kwargs)`` |
Back to top

---
# Docs - LlamaIndex
Documentation[#](#documentation "Permanent link")
==================================================
This directory contains the documentation source code for LlamaIndex, available at https://docs.llamaindex.ai.
This guide is made for anyone who's interested in running LlamaIndex documentation locally, making changes to it and making contributions. LlamaIndex is made by the thriving community behind it, and you're always welcome to make contributions to the project and the documentation.
Build Docs[#](#build-docs "Permanent link")
--------------------------------------------
If you haven't already, clone the LlamaIndex Github repo to a local directory:
`git clone https://github.com/run-llama/llama_index.git && cd llama_index`
Documentation has its own, dedicated Python virtual environment, and all the tools and scripts are available from the `docs` directory:
`cd llama_index/docs`
From now on, we assume all the commands will be executed from the `docs` directory.
Install all dependencies required for building docs (mainly `mkdocs` and its extension):
* [Install poetry](https://python-poetry.org/docs/#installation)
- this will help you manage package dependencies
* `poetry install` - this will install all dependencies needed for building docs
To build the docs and browse them locally run:
`poetry run serve`
During the build, notebooks are converted to documentation pages, and this takes several minutes. If you're not working on the "Examples" section of the documentation, you can run the same command with `--skip-notebooks`:
`poetry run serve --skip-notebooks`
Important
Building the documentation takes a while, so make sure you see the following output before opening the browser:
`... INFO - Documentation built in 53.32 seconds INFO - [16:18:17] Watching paths for changes: 'docs' INFO - [16:18:17] Serving on http://127.0.0.1:8000/en/stable/`
You can now open your browser at http://localhost:8000/ to view the generated docs. The local server will rebuild the docs and refresh your browser every time you make changes to the docs.
Configuration[#](#configuration "Permanent link")
--------------------------------------------------
Part of the configuration in `mkdocs.yml` is generated by a script that takes care of keeping the examples in sync as well as the API reference for all the packages in this repo.
Running the command `poetry run prepare-for-build` from the `docs` folder will update the `mkdocs.yml` with the latest changes, along with writing new api reference files.
Tip
As a contributor, you wouldn't normally need to run this script, feel free to ask for help in the PR.
Back to top

---
# Coa - LlamaIndex
Coa
===
CoAAgentWorker [#](#llama_index.agent.coa.CoAAgentWorker "Permanent link")
---------------------------------------------------------------------------
Bases: `[BaseAgentWorker](../#llama_index.core.agent.types.BaseAgentWorker "llama_index.core.agent.types.BaseAgentWorker") `
Chain-of-abstraction Agent Worker.
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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 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(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]: """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` [#](#llama_index.agent.coa.CoAAgentWorker.from_tools "Permanent link")
`from_tools(tools: Optional[Sequence[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, tool_retriever: Optional[[ObjectRetriever](../../objects/#llama_index.core.objects.ObjectRetriever "llama_index.core.objects.base.ObjectRetriever") [[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, llm: Optional[[LLM](../../llms/#llama_index.core.llms.llm.LLM "llama_index.core.llms.llm.LLM") ] = None, reasoning_prompt_template: Optional[str] = None, refine_reasoning_prompt_template: Optional[str] = None, output_parser: Optional[[BaseOutputParser](../../output_parsers/#llama_index.core.types.BaseOutputParser "llama_index.core.output_parsers.base.BaseOutputParser") ] = None, callback_manager: Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, verbose: bool = False, **kwargs: Any) -> [CoAAgentWorker](#llama_index.agent.coa.CoAAgentWorker "llama_index.agent.coa.step.CoAAgentWorker")`
Convenience constructor method from set of BaseTools (Optional).
Returns:
| Name | Type | Description |
| --- | --- | --- |
| `LLMCompilerAgentWorker` | `[CoAAgentWorker](#llama_index.agent.coa.CoAAgentWorker "llama_index.agent.coa.step.CoAAgentWorker") ` | the LLMCompilerAgentWorker instance |
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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 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 [#](#llama_index.agent.coa.CoAAgentWorker.initialize_step "Permanent link")
`initialize_step(task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStep](../#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep")`
Initialize step from task.
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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(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": ""}, )` |
### get\_tools [#](#llama_index.agent.coa.CoAAgentWorker.get_tools "Permanent link")
`get_tools(query_str: str) -> List[[AsyncBaseTool](../../tools/#llama_index.core.tools.types.AsyncBaseTool "llama_index.core.tools.types.AsyncBaseTool") ]`
Get tools.
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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 [#](#llama_index.agent.coa.CoAAgentWorker.run_step "Permanent link")
`run_step(step: [TaskStep](../#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](../#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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` [#](#llama_index.agent.coa.CoAAgentWorker.arun_step "Permanent link")
`arun_step(step: [TaskStep](../#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](../#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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 [#](#llama_index.agent.coa.CoAAgentWorker.stream_step "Permanent link")
`stream_step(step: [TaskStep](../#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](../#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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` [#](#llama_index.agent.coa.CoAAgentWorker.astream_step "Permanent link")
`astream_step(step: [TaskStep](../#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](../#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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 [#](#llama_index.agent.coa.CoAAgentWorker.finalize_task "Permanent link")
`finalize_task(task: [Task](../#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> None`
Finalize task, after all the steps are completed.
Source code in `llama-index-integrations/agent/llama-index-agent-coa/llama_index/agent/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()` |
Back to top

---
# Core Agent Classes - LlamaIndex
Core Agent Classes[#](#core-agent-classes "Permanent link")
============================================================
Base Types[#](#base-types "Permanent link")
--------------------------------------------
Base agent types.
### BaseAgent [#](#llama_index.core.agent.types.BaseAgent "Permanent link")
Bases: `[BaseChatEngine](../chat_engines/#llama_index.core.chat_engine.types.BaseChatEngine "llama_index.core.chat_engine.types.BaseChatEngine") `, `[BaseQueryEngine](../query_engine/#llama_index.core.base.base_query_engine.BaseQueryEngine "llama_index.core.base.base_query_engine.BaseQueryEngine") `
Base Agent.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 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 | `class BaseAgent(BaseChatEngine, BaseQueryEngine): """Base Agent.""" def _get_prompts(self) -> PromptDictType: """Get prompts.""" # TODO: the ReAct agent does not explicitly specify prompts, would need a # refactor to expose those prompts return {} def _get_prompt_modules(self) -> PromptMixinType: """Get prompt modules.""" return {} def _update_prompts(self, prompts: PromptDictType) -> None: """Update prompts.""" # ===== Query Engine Interface ===== @trace_method("query") def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE: agent_response = self.chat( query_bundle.query_str, chat_history=[], ) return Response( response=str(agent_response), source_nodes=agent_response.source_nodes ) @trace_method("query") async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE: agent_response = await self.achat( query_bundle.query_str, chat_history=[], ) return Response( response=str(agent_response), source_nodes=agent_response.source_nodes ) def stream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None ) -> StreamingAgentChatResponse: raise NotImplementedError("stream_chat not implemented") async def astream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None ) -> StreamingAgentChatResponse: raise NotImplementedError("astream_chat not implemented")` |
### BaseAgentWorker [#](#llama_index.core.agent.types.BaseAgentWorker "Permanent link")
Bases: `PromptMixin`, `DispatcherSpanMixin`
Base agent worker.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 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 | `class BaseAgentWorker(PromptMixin, DispatcherSpanMixin): """Base agent worker.""" def _get_prompts(self) -> PromptDictType: """Get prompts.""" # TODO: the ReAct agent does not explicitly specify prompts, would need a # refactor to expose those prompts return {} def _get_prompt_modules(self) -> PromptMixinType: """Get prompt modules.""" return {} def _update_prompts(self, prompts: PromptDictType) -> None: """Update prompts.""" @abstractmethod def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep: """Initialize step from task.""" @abstractmethod def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" @abstractmethod async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" raise NotImplementedError @abstractmethod def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" # TODO: figure out if we need a different type for TaskStepOutput raise NotImplementedError @abstractmethod async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError @abstractmethod def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) def as_agent(self, **kwargs: Any) -> "AgentRunner": """Return as an agent runner.""" from llama_index.core.agent.runner.base import AgentRunner return AgentRunner(self, **kwargs)` |
#### initialize\_step `abstractmethod` [#](#llama_index.core.agent.types.BaseAgentWorker.initialize_step "Permanent link")
`initialize_step(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep")`
Initialize step from task.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 208
209
210 | `@abstractmethod def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep: """Initialize step from task."""` |
#### run\_step `abstractmethod` [#](#llama_index.core.agent.types.BaseAgentWorker.run_step "Permanent link")
`run_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.base.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 212
213
214 | `@abstractmethod def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step."""` |
#### arun\_step `abstractmethod` `async` [#](#llama_index.core.agent.types.BaseAgentWorker.arun_step "Permanent link")
`arun_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.base.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 216
217
218
219
220
221 | `@abstractmethod async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" raise NotImplementedError` |
#### stream\_step `abstractmethod` [#](#llama_index.core.agent.types.BaseAgentWorker.stream_step "Permanent link")
`stream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.base.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 223
224
225
226
227 | `@abstractmethod def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" # TODO: figure out if we need a different type for TaskStepOutput raise NotImplementedError` |
#### astream\_step `abstractmethod` `async` [#](#llama_index.core.agent.types.BaseAgentWorker.astream_step "Permanent link")
`astream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.base.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 229
230
231
232
233
234 | `@abstractmethod async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError` |
#### finalize\_task `abstractmethod` [#](#llama_index.core.agent.types.BaseAgentWorker.finalize_task "Permanent link")
`finalize_task(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.base.agent.types.Task") , **kwargs: Any) -> None`
Finalize task, after all the steps are completed.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 236
237
238 | `@abstractmethod def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed."""` |
#### set\_callback\_manager [#](#llama_index.core.agent.types.BaseAgentWorker.set_callback_manager "Permanent link")
`set_callback_manager(callback_manager: [CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ) -> None`
Set callback manager.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 240
241 | `def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager."""` |
#### as\_agent [#](#llama_index.core.agent.types.BaseAgentWorker.as_agent "Permanent link")
`as_agent(**kwargs: Any) -> [AgentRunner](#llama_index.core.agent.AgentRunner "llama_index.core.agent.runner.base.AgentRunner")`
Return as an agent runner.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 244
245
246
247
248 | `def as_agent(self, **kwargs: Any) -> "AgentRunner": """Return as an agent runner.""" from llama_index.core.agent.runner.base import AgentRunner return AgentRunner(self, **kwargs)` |
### Task [#](#llama_index.core.agent.types.Task "Permanent link")
Bases: `BaseModel`
Agent Task.
Represents a "run" of an agent given a user input.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `task_id` | `str` | Task ID | `'584a2407-6b5d-4f72-a850-3018d0e043c0'` |
| `input` | `str` | User input | _required_ |
| `memory` | `[BaseMemory](../memory/#llama_index.core.memory.types.BaseMemory "llama_index.core.memory.types.BaseMemory") ` | Conversational Memory. Maintains state before execution of this task. | _required_ |
| `callback_manager` | `[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ` | Callback manager for the task. | `` |
| `extra_state` | `Dict[str, Any]` | Additional user-specified state for a given task. Can be modified throughout the execution of a task. | `{}` |
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 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 | `class Task(BaseModel): """Agent Task. Represents a "run" of an agent given a user input. """ model_config = ConfigDict(arbitrary_types_allowed=True) task_id: str = Field( default_factory=lambda: str(uuid.uuid4()), description="Task ID" ) input: str = Field(..., description="User input") # NOTE: this is state that may be modified throughout the course of execution of the task memory: SerializeAsAny[BaseMemory] = Field( ..., description=( "Conversational Memory. Maintains state before execution of this task." ), ) callback_manager: CallbackManager = Field( default_factory=lambda: CallbackManager([]), exclude=True, description="Callback manager for the task.", ) extra_state: Dict[str, Any] = Field( default_factory=dict, description=( "Additional user-specified state for a given task. " "Can be modified throughout the execution of a task." ), )` |
### TaskStep [#](#llama_index.core.agent.types.TaskStep "Permanent link")
Bases: `BaseModel`
Agent task step.
Represents a single input step within the execution run ("Task") of an agent given a user input.
The output is returned as a `TaskStepOutput`.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `task_id` | `str` | Task ID | _required_ |
| `step_id` | `str` | Step ID | _required_ |
| `input` | `str \| None` | User input | `None` |
| `step_state` | `Dict[str, Any]` | Additional state for a given step. | `{}` |
| `next_steps` | `Dict[str, [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") ]` | Next steps to be executed. | `{}` |
| `prev_steps` | `Dict[str, [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") ]` | Previous steps that were dependencies for this step. | `{}` |
| `is_ready` | `bool` | Is this step ready to be executed? | `True` |
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 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 | ``class TaskStep(BaseModel): """Agent task step. Represents a single input step within the execution run ("Task") of an agent given a user input. The output is returned as a `TaskStepOutput`. """ task_id: str = Field(..., description="Task ID") step_id: str = Field(..., description="Step ID") input: Optional[str] = Field(default=None, description="User input") # memory: BaseMemory = Field( # ..., description="Conversational Memory" # ) step_state: Dict[str, Any] = Field( default_factory=dict, description="Additional state for a given step." ) # NOTE: the state below may change throughout the course of execution # this tracks the relationships to other steps next_steps: Dict[str, "TaskStep"] = Field( default_factory=dict, description="Next steps to be executed." ) prev_steps: Dict[str, "TaskStep"] = Field( default_factory=dict, description="Previous steps that were dependencies for this step.", ) is_ready: bool = Field( default=True, description="Is this step ready to be executed?" ) def get_next_step( self, step_id: str, input: Optional[str] = None, step_state: Optional[Dict[str, Any]] = None, ) -> "TaskStep": """Convenience function to get next step. Preserve task_id, memory, step_state. """ return TaskStep( task_id=self.task_id, step_id=step_id, input=input, # memory=self.memory, step_state=step_state or self.step_state, ) def link_step( self, next_step: "TaskStep", ) -> None: """Link to next step. Add link from this step to next, and from next step to current. """ self.next_steps[next_step.step_id] = next_step next_step.prev_steps[self.step_id] = self`` |
#### get\_next\_step [#](#llama_index.core.agent.types.TaskStep.get_next_step "Permanent link")
`get_next_step(step_id: str, input: Optional[str] = None, step_state: Optional[Dict[str, Any]] = None) -> [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep")`
Convenience function to get next step.
Preserve task\_id, memory, step\_state.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128 | `def get_next_step( self, step_id: str, input: Optional[str] = None, step_state: Optional[Dict[str, Any]] = None, ) -> "TaskStep": """Convenience function to get next step. Preserve task_id, memory, step_state. """ return TaskStep( task_id=self.task_id, step_id=step_id, input=input, # memory=self.memory, step_state=step_state or self.step_state, )` |
#### link\_step [#](#llama_index.core.agent.types.TaskStep.link_step "Permanent link")
`link_step(next_step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") ) -> None`
Link to next step.
Add link from this step to next, and from next step to current.
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 130
131
132
133
134
135
136
137
138
139
140 | `def link_step( self, next_step: "TaskStep", ) -> None: """Link to next step. Add link from this step to next, and from next step to current. """ self.next_steps[next_step.step_id] = next_step next_step.prev_steps[self.step_id] = self` |
### TaskStepOutput [#](#llama_index.core.agent.types.TaskStepOutput "Permanent link")
Bases: `BaseModel`
Agent task step output.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `output` | `Any` | Task step output | _required_ |
| `task_step` | `[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") ` | Task step input | _required_ |
| `next_steps` | `List[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.base.agent.types.TaskStep") ]` | Next steps to be executed. | _required_ |
| `is_last` | `bool` | Is this the last step? | `False` |
Source code in `llama-index-core/llama_index/core/base/agent/types.py`
| | |
| --- | --- |
| 143
144
145
146
147
148
149
150
151
152
153 | `class TaskStepOutput(BaseModel): """Agent task step output.""" output: Any = Field(..., description="Task step output") task_step: TaskStep = Field(..., description="Task step input") next_steps: List[TaskStep] = Field(..., description="Next steps to be executed.") is_last: bool = Field(default=False, description="Is this the last step?") def __str__(self) -> str: """String representation.""" return str(self.output)` |
Runners[#](#runners "Permanent link")
--------------------------------------
### AgentRunner [#](#llama_index.core.agent.AgentRunner "Permanent link")
Bases: `BaseAgentRunner`
Agent runner.
Top-level agent orchestrator that can create tasks, run each step in a task, or run a task e2e. Stores state and keeps track of tasks.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `agent_worker` | `[BaseAgentWorker](#llama_index.core.agent.types.BaseAgentWorker "llama_index.core.agent.types.BaseAgentWorker") ` | step executor | _required_ |
| `chat_history` | `Optional[List[[ChatMessage](../prompts/#llama_index.core.prompts.ChatMessage "llama_index.core.base.llms.types.ChatMessage") ]]` | chat history. Defaults to None. | `None` |
| `state` | `Optional[AgentState]` | agent state. Defaults to None. | `None` |
| `memory` | `Optional[[BaseMemory](../memory/#llama_index.core.memory.types.BaseMemory "llama_index.core.memory.types.BaseMemory") ]` | memory. Defaults to None. | `None` |
| `llm` | `Optional[[LLM](../llms/#llama_index.core.llms.llm.LLM "llama_index.core.llms.llm.LLM") ]` | LLM. Defaults to None. | `None` |
| `callback_manager` | `Optional[[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ]` | callback manager. Defaults to None. | `None` |
| `init_task_state_kwargs` | `Optional[dict]` | init task state kwargs. Defaults to None. | `None` |
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 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
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736 | ``class AgentRunner(BaseAgentRunner): """Agent runner. Top-level agent orchestrator that can create tasks, run each step in a task, or run a task e2e. Stores state and keeps track of tasks. Args: agent_worker (BaseAgentWorker): step executor chat_history (Optional[List[ChatMessage]], optional): chat history. Defaults to None. state (Optional[AgentState], optional): agent state. Defaults to None. memory (Optional[BaseMemory], optional): memory. Defaults to None. llm (Optional[LLM], optional): LLM. Defaults to None. callback_manager (Optional[CallbackManager], optional): callback manager. Defaults to None. init_task_state_kwargs (Optional[dict], optional): init task state kwargs. Defaults to None. """ # # TODO: implement this in Pydantic def __init__( self, agent_worker: BaseAgentWorker, chat_history: Optional[List[ChatMessage]] = None, state: Optional[AgentState] = None, memory: Optional[BaseMemory] = None, llm: Optional[LLM] = None, callback_manager: Optional[CallbackManager] = None, init_task_state_kwargs: Optional[dict] = None, delete_task_on_finish: bool = False, default_tool_choice: str = "auto", verbose: bool = False, ) -> None: """Initialize.""" self.agent_worker = agent_worker self.state = state or AgentState() self.memory = memory or ChatMemoryBuffer.from_defaults(chat_history, llm=llm) # get and set callback manager if callback_manager is not None: self.agent_worker.set_callback_manager(callback_manager) self.callback_manager = callback_manager else: # TODO: This is *temporary* # Stopgap before having a callback on the BaseAgentWorker interface. # Doing that requires a bit more refactoring to make sure existing code # doesn't break. if hasattr(self.agent_worker, "callback_manager"): self.callback_manager = ( self.agent_worker.callback_manager or CallbackManager() ) else: self.callback_manager = CallbackManager() self.init_task_state_kwargs = init_task_state_kwargs or {} self.delete_task_on_finish = delete_task_on_finish self.default_tool_choice = default_tool_choice self.verbose = verbose @staticmethod def from_llm( tools: Optional[List[BaseTool]] = None, llm: Optional[LLM] = None, **kwargs: Any, ) -> "AgentRunner": from llama_index.core.agent import ReActAgent if os.getenv("IS_TESTING"): return ReActAgent.from_tools( tools=tools, llm=llm, **kwargs, ) try: from llama_index.llms.openai import OpenAI # pants: no-infer-dep from llama_index.llms.openai.utils import ( is_function_calling_model, ) # pants: no-infer-dep except ImportError: raise ImportError( "`llama-index-llms-openai` package not found. Please " "install by running `pip install llama-index-llms-openai`." ) if isinstance(llm, OpenAI) and is_function_calling_model(llm.model): from llama_index.agent.openai import OpenAIAgent # pants: no-infer-dep return OpenAIAgent.from_tools( tools=tools, llm=llm, **kwargs, ) else: return ReActAgent.from_tools( tools=tools, llm=llm, **kwargs, ) @property def chat_history(self) -> List[ChatMessage]: return self.memory.get_all() def reset(self) -> None: self.memory.reset() self.state.reset() def create_task(self, input: str, **kwargs: Any) -> Task: """Create task.""" if not self.init_task_state_kwargs: extra_state = kwargs.pop("extra_state", {}) else: if "extra_state" in kwargs: raise ValueError( "Cannot specify both `extra_state` and `init_task_state_kwargs`" ) else: extra_state = self.init_task_state_kwargs callback_manager = kwargs.pop("callback_manager", self.callback_manager) task = Task( input=input, memory=self.memory, extra_state=extra_state, callback_manager=callback_manager, **kwargs, ) # # put input into memory # self.memory.put(ChatMessage(content=input, role=MessageRole.USER)) # get initial step from task, and put it in the step queue initial_step = self.agent_worker.initialize_step(task) task_state = TaskState( task=task, step_queue=deque([initial_step]), ) # add it to state self.state.task_dict[task.task_id] = task_state return task def delete_task( self, task_id: str, ) -> None: """Delete task. NOTE: this will not delete any previous executions from memory. """ self.state.task_dict.pop(task_id) def list_tasks(self, **kwargs: Any) -> List[Task]: """List tasks.""" return [task_state.task for task_state in self.state.task_dict.values()] def get_task(self, task_id: str, **kwargs: Any) -> Task: """Get task.""" return self.state.get_task(task_id) def get_upcoming_steps(self, task_id: str, **kwargs: Any) -> List[TaskStep]: """Get upcoming steps.""" return list(self.state.get_step_queue(task_id)) def get_completed_steps(self, task_id: str, **kwargs: Any) -> List[TaskStepOutput]: """Get completed steps.""" return self.state.get_completed_steps(task_id) def get_task_output(self, task_id: str, **kwargs: Any) -> TaskStepOutput: """Get task output.""" completed_steps = self.get_completed_steps(task_id) if len(completed_steps) == 0: raise ValueError(f"No completed steps for task_id: {task_id}") return completed_steps[-1] def get_completed_tasks(self, **kwargs: Any) -> List[Task]: """Get completed tasks.""" task_states = list(self.state.task_dict.values()) completed_tasks = [] for task_state in task_states: completed_steps = self.get_completed_steps(task_state.task.task_id) if len(completed_steps) > 0 and completed_steps[-1].is_last: completed_tasks.append(task_state.task) return completed_tasks @dispatcher.span def _run_step( self, task_id: str, step: Optional[TaskStep] = None, input: Optional[str] = None, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> TaskStepOutput: """Execute step.""" task = self.state.get_task(task_id) step_queue = self.state.get_step_queue(task_id) step = step or step_queue.popleft() if input is not None: step.input = input dispatcher.event( AgentRunStepStartEvent(task_id=task_id, step=step, input=input) ) if self.verbose: print(f"> Running step {step.step_id}. Step input: {step.input}") # TODO: figure out if you can dynamically swap in different step executors # not clear when you would do that by theoretically possible if mode == ChatResponseMode.WAIT: cur_step_output = self.agent_worker.run_step(step, task, **kwargs) elif mode == ChatResponseMode.STREAM: cur_step_output = self.agent_worker.stream_step(step, task, **kwargs) else: raise ValueError(f"Invalid mode: {mode}") # append cur_step_output next steps to queue next_steps = cur_step_output.next_steps step_queue.extend(next_steps) # add cur_step_output to completed steps completed_steps = self.state.get_completed_steps(task_id) completed_steps.append(cur_step_output) dispatcher.event(AgentRunStepEndEvent(step_output=cur_step_output)) return cur_step_output @dispatcher.span async def _arun_step( self, task_id: str, step: Optional[TaskStep] = None, input: Optional[str] = None, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> TaskStepOutput: """Execute step.""" dispatcher.event( AgentRunStepStartEvent(task_id=task_id, step=step, input=input) ) task = self.state.get_task(task_id) step_queue = self.state.get_step_queue(task_id) step = step or step_queue.popleft() if input is not None: step.input = input if self.verbose: print(f"> Running step {step.step_id}. Step input: {step.input}") # TODO: figure out if you can dynamically swap in different step executors # not clear when you would do that by theoretically possible if mode == ChatResponseMode.WAIT: cur_step_output = await self.agent_worker.arun_step(step, task, **kwargs) elif mode == ChatResponseMode.STREAM: cur_step_output = await self.agent_worker.astream_step(step, task, **kwargs) else: raise ValueError(f"Invalid mode: {mode}") # append cur_step_output next steps to queue next_steps = cur_step_output.next_steps step_queue.extend(next_steps) # add cur_step_output to completed steps completed_steps = self.state.get_completed_steps(task_id) completed_steps.append(cur_step_output) dispatcher.event(AgentRunStepEndEvent(step_output=cur_step_output)) return cur_step_output @dispatcher.span def run_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step.""" step = validate_step_from_args(task_id, input, step, **kwargs) return self._run_step( task_id, step, input=input, mode=ChatResponseMode.WAIT, **kwargs ) @dispatcher.span async def arun_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async).""" step = validate_step_from_args(task_id, input, step, **kwargs) return await self._arun_step( task_id, step, input=input, mode=ChatResponseMode.WAIT, **kwargs ) @dispatcher.span def stream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (stream).""" step = validate_step_from_args(task_id, input, step, **kwargs) return self._run_step( task_id, step, input=input, mode=ChatResponseMode.STREAM, **kwargs ) @dispatcher.span async def astream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async stream).""" step = validate_step_from_args(task_id, input, step, **kwargs) return await self._arun_step( task_id, step, input=input, mode=ChatResponseMode.STREAM, **kwargs ) @dispatcher.span def finalize_response( self, task_id: str, step_output: Optional[TaskStepOutput] = None, ) -> AGENT_CHAT_RESPONSE_TYPE: """Finalize response.""" if step_output is None: step_output = self.state.get_completed_steps(task_id)[-1] if not step_output.is_last: raise ValueError( "finalize_response can only be called on the last step output" ) if not isinstance( step_output.output, (AgentChatResponse, StreamingAgentChatResponse), ): raise ValueError( "When `is_last` is True, cur_step_output.output must be " f"AGENT_CHAT_RESPONSE_TYPE: {step_output.output}" ) # finalize task self.agent_worker.finalize_task(self.state.get_task(task_id)) if self.delete_task_on_finish: self.delete_task(task_id) # Attach all sources generated across all steps step_output.output.sources = self.get_task(task_id).extra_state.get( "sources", [] ) step_output.output.set_source_nodes() return cast(AGENT_CHAT_RESPONSE_TYPE, step_output.output) @dispatcher.span def _chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Chat with step executor.""" if chat_history is not None: self.memory.set(chat_history) task = self.create_task(message) result_output = None dispatcher.event(AgentChatWithStepStartEvent(user_msg=message)) while True: # pass step queue in as argument, assume step executor is stateless cur_step_output = self._run_step( task.task_id, mode=mode, tool_choice=tool_choice ) if cur_step_output.is_last: result_output = cur_step_output break # ensure tool_choice does not cause endless loops tool_choice = "auto" result = self.finalize_response( task.task_id, result_output, ) dispatcher.event(AgentChatWithStepEndEvent(response=result)) return result @dispatcher.span async def _achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Chat with step executor.""" if chat_history is not None: self.memory.set(chat_history) task = self.create_task(message) result_output = None dispatcher.event(AgentChatWithStepStartEvent(user_msg=message)) while True: # pass step queue in as argument, assume step executor is stateless cur_step_output = await self._arun_step( task.task_id, mode=mode, tool_choice=tool_choice ) if cur_step_output.is_last: result_output = cur_step_output break # ensure tool_choice does not cause endless loops tool_choice = "auto" result = self.finalize_response( task.task_id, result_output, ) dispatcher.event(AgentChatWithStepEndEvent(response=result)) return result @dispatcher.span @trace_method("chat") def chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Optional[Union[str, dict]] = None, ) -> AgentChatResponse: # override tool choice is provided as input. if tool_choice is None: tool_choice = self.default_tool_choice with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = self._chat( message=message, chat_history=chat_history, tool_choice=tool_choice, mode=ChatResponseMode.WAIT, ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @dispatcher.span @trace_method("chat") async def achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Optional[Union[str, dict]] = None, ) -> AgentChatResponse: # override tool choice is provided as input. if tool_choice is None: tool_choice = self.default_tool_choice with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = await self._achat( message=message, chat_history=chat_history, tool_choice=tool_choice, mode=ChatResponseMode.WAIT, ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @dispatcher.span @trace_method("chat") def stream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Optional[Union[str, dict]] = None, ) -> StreamingAgentChatResponse: # override tool choice is provided as input. if tool_choice is None: tool_choice = self.default_tool_choice with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = self._chat( message, chat_history, tool_choice, mode=ChatResponseMode.STREAM ) assert isinstance(chat_response, StreamingAgentChatResponse) or ( isinstance(chat_response, AgentChatResponse) and chat_response.is_dummy_stream ) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response # type: ignore @dispatcher.span @trace_method("chat") async def astream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Optional[Union[str, dict]] = None, ) -> StreamingAgentChatResponse: # override tool choice is provided as input. if tool_choice is None: tool_choice = self.default_tool_choice with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = await self._achat( message, chat_history, tool_choice, mode=ChatResponseMode.STREAM ) assert isinstance(chat_response, StreamingAgentChatResponse) or ( isinstance(chat_response, AgentChatResponse) and chat_response.is_dummy_stream ) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response # type: ignore def undo_step(self, task_id: str) -> None: """Undo previous step.""" raise NotImplementedError("undo_step not implemented")`` |
#### create\_task [#](#llama_index.core.agent.AgentRunner.create_task "Permanent link")
`create_task(input: str, **kwargs: Any) -> [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task")`
Create task.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338 | ``def create_task(self, input: str, **kwargs: Any) -> Task: """Create task.""" if not self.init_task_state_kwargs: extra_state = kwargs.pop("extra_state", {}) else: if "extra_state" in kwargs: raise ValueError( "Cannot specify both `extra_state` and `init_task_state_kwargs`" ) else: extra_state = self.init_task_state_kwargs callback_manager = kwargs.pop("callback_manager", self.callback_manager) task = Task( input=input, memory=self.memory, extra_state=extra_state, callback_manager=callback_manager, **kwargs, ) # # put input into memory # self.memory.put(ChatMessage(content=input, role=MessageRole.USER)) # get initial step from task, and put it in the step queue initial_step = self.agent_worker.initialize_step(task) task_state = TaskState( task=task, step_queue=deque([initial_step]), ) # add it to state self.state.task_dict[task.task_id] = task_state return task`` |
#### delete\_task [#](#llama_index.core.agent.AgentRunner.delete_task "Permanent link")
`delete_task(task_id: str) -> None`
Delete task.
NOTE: this will not delete any previous executions from memory.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 340
341
342
343
344
345
346
347
348
349 | `def delete_task( self, task_id: str, ) -> None: """Delete task. NOTE: this will not delete any previous executions from memory. """ self.state.task_dict.pop(task_id)` |
#### list\_tasks [#](#llama_index.core.agent.AgentRunner.list_tasks "Permanent link")
`list_tasks(**kwargs: Any) -> List[[Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") ]`
List tasks.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 351
352
353 | `def list_tasks(self, **kwargs: Any) -> List[Task]: """List tasks.""" return [task_state.task for task_state in self.state.task_dict.values()]` |
#### get\_task [#](#llama_index.core.agent.AgentRunner.get_task "Permanent link")
`get_task(task_id: str, **kwargs: Any) -> [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task")`
Get task.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 355
356
357 | `def get_task(self, task_id: str, **kwargs: Any) -> Task: """Get task.""" return self.state.get_task(task_id)` |
#### get\_upcoming\_steps [#](#llama_index.core.agent.AgentRunner.get_upcoming_steps "Permanent link")
`get_upcoming_steps(task_id: str, **kwargs: Any) -> List[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ]`
Get upcoming steps.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 359
360
361 | `def get_upcoming_steps(self, task_id: str, **kwargs: Any) -> List[TaskStep]: """Get upcoming steps.""" return list(self.state.get_step_queue(task_id))` |
#### get\_completed\_steps [#](#llama_index.core.agent.AgentRunner.get_completed_steps "Permanent link")
`get_completed_steps(task_id: str, **kwargs: Any) -> List[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ]`
Get completed steps.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 363
364
365 | `def get_completed_steps(self, task_id: str, **kwargs: Any) -> List[TaskStepOutput]: """Get completed steps.""" return self.state.get_completed_steps(task_id)` |
#### get\_task\_output [#](#llama_index.core.agent.AgentRunner.get_task_output "Permanent link")
`get_task_output(task_id: str, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Get task output.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 367
368
369
370
371
372 | `def get_task_output(self, task_id: str, **kwargs: Any) -> TaskStepOutput: """Get task output.""" completed_steps = self.get_completed_steps(task_id) if len(completed_steps) == 0: raise ValueError(f"No completed steps for task_id: {task_id}") return completed_steps[-1]` |
#### get\_completed\_tasks [#](#llama_index.core.agent.AgentRunner.get_completed_tasks "Permanent link")
`get_completed_tasks(**kwargs: Any) -> List[[Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") ]`
Get completed tasks.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 374
375
376
377
378
379
380
381
382
383 | `def get_completed_tasks(self, **kwargs: Any) -> List[Task]: """Get completed tasks.""" task_states = list(self.state.task_dict.values()) completed_tasks = [] for task_state in task_states: completed_steps = self.get_completed_steps(task_state.task.task_id) if len(completed_steps) > 0 and completed_steps[-1].is_last: completed_tasks.append(task_state.task) return completed_tasks` |
#### run\_step [#](#llama_index.core.agent.AgentRunner.run_step "Permanent link")
`run_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 469
470
471
472
473
474
475
476
477
478
479
480
481 | `@dispatcher.span def run_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step.""" step = validate_step_from_args(task_id, input, step, **kwargs) return self._run_step( task_id, step, input=input, mode=ChatResponseMode.WAIT, **kwargs )` |
#### arun\_step `async` [#](#llama_index.core.agent.AgentRunner.arun_step "Permanent link")
`arun_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 483
484
485
486
487
488
489
490
491
492
493
494
495 | `@dispatcher.span async def arun_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async).""" step = validate_step_from_args(task_id, input, step, **kwargs) return await self._arun_step( task_id, step, input=input, mode=ChatResponseMode.WAIT, **kwargs )` |
#### stream\_step [#](#llama_index.core.agent.AgentRunner.stream_step "Permanent link")
`stream_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 497
498
499
500
501
502
503
504
505
506
507
508
509 | `@dispatcher.span def stream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (stream).""" step = validate_step_from_args(task_id, input, step, **kwargs) return self._run_step( task_id, step, input=input, mode=ChatResponseMode.STREAM, **kwargs )` |
#### astream\_step `async` [#](#llama_index.core.agent.AgentRunner.astream_step "Permanent link")
`astream_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 511
512
513
514
515
516
517
518
519
520
521
522
523 | `@dispatcher.span async def astream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async stream).""" step = validate_step_from_args(task_id, input, step, **kwargs) return await self._arun_step( task_id, step, input=input, mode=ChatResponseMode.STREAM, **kwargs )` |
#### finalize\_response [#](#llama_index.core.agent.AgentRunner.finalize_response "Permanent link")
`finalize_response(task_id: str, step_output: Optional[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ] = None) -> AGENT_CHAT_RESPONSE_TYPE`
Finalize response.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560 | ``@dispatcher.span def finalize_response( self, task_id: str, step_output: Optional[TaskStepOutput] = None, ) -> AGENT_CHAT_RESPONSE_TYPE: """Finalize response.""" if step_output is None: step_output = self.state.get_completed_steps(task_id)[-1] if not step_output.is_last: raise ValueError( "finalize_response can only be called on the last step output" ) if not isinstance( step_output.output, (AgentChatResponse, StreamingAgentChatResponse), ): raise ValueError( "When `is_last` is True, cur_step_output.output must be " f"AGENT_CHAT_RESPONSE_TYPE: {step_output.output}" ) # finalize task self.agent_worker.finalize_task(self.state.get_task(task_id)) if self.delete_task_on_finish: self.delete_task(task_id) # Attach all sources generated across all steps step_output.output.sources = self.get_task(task_id).extra_state.get( "sources", [] ) step_output.output.set_source_nodes() return cast(AGENT_CHAT_RESPONSE_TYPE, step_output.output)`` |
#### undo\_step [#](#llama_index.core.agent.AgentRunner.undo_step "Permanent link")
`undo_step(task_id: str) -> None`
Undo previous step.
Source code in `llama-index-core/llama_index/core/agent/runner/base.py`
| | |
| --- | --- |
| 734
735
736 | `def undo_step(self, task_id: str) -> None: """Undo previous step.""" raise NotImplementedError("undo_step not implemented")` |
### ParallelAgentRunner [#](#llama_index.core.agent.ParallelAgentRunner "Permanent link")
Bases: `BaseAgentRunner`
Parallel agent runner.
Executes steps in queue in parallel. Requires async support.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 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
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496 | ``class ParallelAgentRunner(BaseAgentRunner): """Parallel agent runner. Executes steps in queue in parallel. Requires async support. """ def __init__( self, agent_worker: BaseAgentWorker, chat_history: Optional[List[ChatMessage]] = None, state: Optional[DAGAgentState] = None, memory: Optional[BaseMemory] = None, llm: Optional[LLM] = None, callback_manager: Optional[CallbackManager] = None, init_task_state_kwargs: Optional[dict] = None, delete_task_on_finish: bool = False, ) -> None: """Initialize.""" self.memory = memory or ChatMemoryBuffer.from_defaults(chat_history, llm=llm) self.state = state or DAGAgentState() self.callback_manager = callback_manager or CallbackManager([]) self.init_task_state_kwargs = init_task_state_kwargs or {} self.agent_worker = agent_worker self.delete_task_on_finish = delete_task_on_finish @property def chat_history(self) -> List[ChatMessage]: return self.memory.get_all() def reset(self) -> None: self.memory.reset() def create_task(self, input: str, **kwargs: Any) -> Task: """Create task.""" task = Task( input=input, memory=self.memory, extra_state=self.init_task_state_kwargs, **kwargs, ) # # put input into memory # self.memory.put(ChatMessage(content=input, role=MessageRole.USER)) # add it to state # get initial step from task, and put it in the step queue initial_step = self.agent_worker.initialize_step(task) task_state = DAGTaskState( task=task, root_step=initial_step, step_queue=deque([initial_step]), ) self.state.task_dict[task.task_id] = task_state return task def delete_task( self, task_id: str, ) -> None: """Delete task. NOTE: this will not delete any previous executions from memory. """ self.state.task_dict.pop(task_id) def get_completed_tasks(self, **kwargs: Any) -> List[Task]: """Get completed tasks.""" task_states = list(self.state.task_dict.values()) return [ task_state.task for task_state in task_states if len(task_state.completed_steps) > 0 and task_state.completed_steps[-1].is_last ] def get_task_output(self, task_id: str, **kwargs: Any) -> TaskStepOutput: """Get task output.""" task_state = self.state.task_dict[task_id] if len(task_state.completed_steps) == 0: raise ValueError(f"No completed steps for task_id: {task_id}") return task_state.completed_steps[-1] def list_tasks(self, **kwargs: Any) -> List[Task]: """List tasks.""" task_states = list(self.state.task_dict.values()) return [task_state.task for task_state in task_states] def get_task(self, task_id: str, **kwargs: Any) -> Task: """Get task.""" return self.state.get_task(task_id) def get_upcoming_steps(self, task_id: str, **kwargs: Any) -> List[TaskStep]: """Get upcoming steps.""" return list(self.state.get_step_queue(task_id)) def get_completed_steps(self, task_id: str, **kwargs: Any) -> List[TaskStepOutput]: """Get completed steps.""" return self.state.get_completed_steps(task_id) def run_steps_in_queue( self, task_id: str, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> List[TaskStepOutput]: """Execute steps in queue. Run all steps in queue, clearing it out. Assume that all steps can be run in parallel. """ return asyncio_run(self.arun_steps_in_queue(task_id, mode=mode, **kwargs)) async def arun_steps_in_queue( self, task_id: str, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> List[TaskStepOutput]: """Execute all steps in queue. All steps in queue are assumed to be ready. """ # first pop all steps from step_queue steps: List[TaskStep] = [] while len(self.state.get_step_queue(task_id)) > 0: steps.append(self.state.get_step_queue(task_id).popleft()) # take every item in the queue, and run it tasks = [] for step in steps: tasks.append(self._arun_step(task_id, step=step, mode=mode, **kwargs)) return await asyncio.gather(*tasks) def _run_step( self, task_id: str, step: Optional[TaskStep] = None, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> TaskStepOutput: """Execute step.""" task = self.state.get_task(task_id) task_queue = self.state.get_step_queue(task_id) step = step or task_queue.popleft() if not step.is_ready: raise ValueError(f"Step {step.step_id} is not ready") if mode == ChatResponseMode.WAIT: cur_step_output: TaskStepOutput = self.agent_worker.run_step( step, task, **kwargs ) elif mode == ChatResponseMode.STREAM: cur_step_output = self.agent_worker.stream_step(step, task, **kwargs) else: raise ValueError(f"Invalid mode: {mode}") for next_step in cur_step_output.next_steps: if next_step.is_ready: task_queue.append(next_step) # add cur_step_output to completed steps completed_steps = self.state.get_completed_steps(task_id) completed_steps.append(cur_step_output) return cur_step_output async def _arun_step( self, task_id: str, step: Optional[TaskStep] = None, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> TaskStepOutput: """Execute step.""" task = self.state.get_task(task_id) task_queue = self.state.get_step_queue(task_id) step = step or task_queue.popleft() if not step.is_ready: raise ValueError(f"Step {step.step_id} is not ready") if mode == ChatResponseMode.WAIT: cur_step_output = await self.agent_worker.arun_step(step, task, **kwargs) elif mode == ChatResponseMode.STREAM: cur_step_output = await self.agent_worker.astream_step(step, task, **kwargs) else: raise ValueError(f"Invalid mode: {mode}") for next_step in cur_step_output.next_steps: if next_step.is_ready: task_queue.append(next_step) # add cur_step_output to completed steps completed_steps = self.state.get_completed_steps(task_id) completed_steps.append(cur_step_output) return cur_step_output def run_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step.""" return self._run_step(task_id, step, mode=ChatResponseMode.WAIT, **kwargs) async def arun_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async).""" return await self._arun_step( task_id, step, mode=ChatResponseMode.WAIT, **kwargs ) def stream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (stream).""" return self._run_step(task_id, step, mode=ChatResponseMode.STREAM, **kwargs) async def astream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async stream).""" return await self._arun_step( task_id, step, mode=ChatResponseMode.STREAM, **kwargs ) def finalize_response( self, task_id: str, step_output: Optional[TaskStepOutput] = None, ) -> AGENT_CHAT_RESPONSE_TYPE: """Finalize response.""" if step_output is None: step_output = self.state.get_completed_steps(task_id)[-1] if not step_output.is_last: raise ValueError( "finalize_response can only be called on the last step output" ) if not isinstance( step_output.output, (AgentChatResponse, StreamingAgentChatResponse), ): raise ValueError( "When `is_last` is True, cur_step_output.output must be " f"AGENT_CHAT_RESPONSE_TYPE: {step_output.output}" ) # finalize task self.agent_worker.finalize_task(self.state.get_task(task_id)) if self.delete_task_on_finish: self.delete_task(task_id) return cast(AGENT_CHAT_RESPONSE_TYPE, step_output.output) def _chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Chat with step executor.""" if chat_history is not None: self.memory.set(chat_history) task = self.create_task(message) result_output = None while True: # pass step queue in as argument, assume step executor is stateless cur_step_outputs = self.run_steps_in_queue(task.task_id, mode=mode) # check if a step output is_last is_last = any( cur_step_output.is_last for cur_step_output in cur_step_outputs ) if is_last: if len(cur_step_outputs) > 1: raise ValueError( "More than one step output returned in final step." ) cur_step_output = cur_step_outputs[0] result_output = cur_step_output break return self.finalize_response( task.task_id, result_output, ) async def _achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Chat with step executor.""" if chat_history is not None: self.memory.set(chat_history) task = self.create_task(message) result_output = None while True: # pass step queue in as argument, assume step executor is stateless cur_step_outputs = await self.arun_steps_in_queue(task.task_id, mode=mode) # check if a step output is_last is_last = any( cur_step_output.is_last for cur_step_output in cur_step_outputs ) if is_last: if len(cur_step_outputs) > 1: raise ValueError( "More than one step output returned in final step." ) cur_step_output = cur_step_outputs[0] result_output = cur_step_output break return self.finalize_response( task.task_id, result_output, ) @trace_method("chat") def chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", ) -> AgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = self._chat( message, chat_history, tool_choice, mode=ChatResponseMode.WAIT ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @trace_method("chat") async def achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", ) -> AgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = await self._achat( message, chat_history, tool_choice, mode=ChatResponseMode.WAIT ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @trace_method("chat") def stream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", ) -> StreamingAgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = self._chat( message, chat_history, tool_choice, mode=ChatResponseMode.STREAM ) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response # type: ignore @trace_method("chat") async def astream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, tool_choice: Union[str, dict] = "auto", ) -> StreamingAgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = await self._achat( message, chat_history, tool_choice, mode=ChatResponseMode.STREAM ) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response # type: ignore def undo_step(self, task_id: str) -> None: """Undo previous step.""" raise NotImplementedError("undo_step not implemented")`` |
#### create\_task [#](#llama_index.core.agent.ParallelAgentRunner.create_task "Permanent link")
`create_task(input: str, **kwargs: Any) -> [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task")`
Create task.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127 | `def create_task(self, input: str, **kwargs: Any) -> Task: """Create task.""" task = Task( input=input, memory=self.memory, extra_state=self.init_task_state_kwargs, **kwargs, ) # # put input into memory # self.memory.put(ChatMessage(content=input, role=MessageRole.USER)) # add it to state # get initial step from task, and put it in the step queue initial_step = self.agent_worker.initialize_step(task) task_state = DAGTaskState( task=task, root_step=initial_step, step_queue=deque([initial_step]), ) self.state.task_dict[task.task_id] = task_state return task` |
#### delete\_task [#](#llama_index.core.agent.ParallelAgentRunner.delete_task "Permanent link")
`delete_task(task_id: str) -> None`
Delete task.
NOTE: this will not delete any previous executions from memory.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 129
130
131
132
133
134
135
136
137
138 | `def delete_task( self, task_id: str, ) -> None: """Delete task. NOTE: this will not delete any previous executions from memory. """ self.state.task_dict.pop(task_id)` |
#### get\_completed\_tasks [#](#llama_index.core.agent.ParallelAgentRunner.get_completed_tasks "Permanent link")
`get_completed_tasks(**kwargs: Any) -> List[[Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") ]`
Get completed tasks.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 140
141
142
143
144
145
146
147
148 | `def get_completed_tasks(self, **kwargs: Any) -> List[Task]: """Get completed tasks.""" task_states = list(self.state.task_dict.values()) return [ task_state.task for task_state in task_states if len(task_state.completed_steps) > 0 and task_state.completed_steps[-1].is_last ]` |
#### get\_task\_output [#](#llama_index.core.agent.ParallelAgentRunner.get_task_output "Permanent link")
`get_task_output(task_id: str, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Get task output.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 150
151
152
153
154
155 | `def get_task_output(self, task_id: str, **kwargs: Any) -> TaskStepOutput: """Get task output.""" task_state = self.state.task_dict[task_id] if len(task_state.completed_steps) == 0: raise ValueError(f"No completed steps for task_id: {task_id}") return task_state.completed_steps[-1]` |
#### list\_tasks [#](#llama_index.core.agent.ParallelAgentRunner.list_tasks "Permanent link")
`list_tasks(**kwargs: Any) -> List[[Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") ]`
List tasks.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 157
158
159
160 | `def list_tasks(self, **kwargs: Any) -> List[Task]: """List tasks.""" task_states = list(self.state.task_dict.values()) return [task_state.task for task_state in task_states]` |
#### get\_task [#](#llama_index.core.agent.ParallelAgentRunner.get_task "Permanent link")
`get_task(task_id: str, **kwargs: Any) -> [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task")`
Get task.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 162
163
164 | `def get_task(self, task_id: str, **kwargs: Any) -> Task: """Get task.""" return self.state.get_task(task_id)` |
#### get\_upcoming\_steps [#](#llama_index.core.agent.ParallelAgentRunner.get_upcoming_steps "Permanent link")
`get_upcoming_steps(task_id: str, **kwargs: Any) -> List[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ]`
Get upcoming steps.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 166
167
168 | `def get_upcoming_steps(self, task_id: str, **kwargs: Any) -> List[TaskStep]: """Get upcoming steps.""" return list(self.state.get_step_queue(task_id))` |
#### get\_completed\_steps [#](#llama_index.core.agent.ParallelAgentRunner.get_completed_steps "Permanent link")
`get_completed_steps(task_id: str, **kwargs: Any) -> List[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ]`
Get completed steps.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 170
171
172 | `def get_completed_steps(self, task_id: str, **kwargs: Any) -> List[TaskStepOutput]: """Get completed steps.""" return self.state.get_completed_steps(task_id)` |
#### run\_steps\_in\_queue [#](#llama_index.core.agent.ParallelAgentRunner.run_steps_in_queue "Permanent link")
`run_steps_in_queue(task_id: str, mode: [ChatResponseMode](../chat_engines/#llama_index.core.chat_engine.types.ChatResponseMode "llama_index.core.chat_engine.types.ChatResponseMode") = WAIT, **kwargs: Any) -> List[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ]`
Execute steps in queue.
Run all steps in queue, clearing it out.
Assume that all steps can be run in parallel.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 174
175
176
177
178
179
180
181
182
183
184
185
186
187 | `def run_steps_in_queue( self, task_id: str, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> List[TaskStepOutput]: """Execute steps in queue. Run all steps in queue, clearing it out. Assume that all steps can be run in parallel. """ return asyncio_run(self.arun_steps_in_queue(task_id, mode=mode, **kwargs))` |
#### arun\_steps\_in\_queue `async` [#](#llama_index.core.agent.ParallelAgentRunner.arun_steps_in_queue "Permanent link")
`arun_steps_in_queue(task_id: str, mode: [ChatResponseMode](../chat_engines/#llama_index.core.chat_engine.types.ChatResponseMode "llama_index.core.chat_engine.types.ChatResponseMode") = WAIT, **kwargs: Any) -> List[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ]`
Execute all steps in queue.
All steps in queue are assumed to be ready.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210 | `async def arun_steps_in_queue( self, task_id: str, mode: ChatResponseMode = ChatResponseMode.WAIT, **kwargs: Any, ) -> List[TaskStepOutput]: """Execute all steps in queue. All steps in queue are assumed to be ready. """ # first pop all steps from step_queue steps: List[TaskStep] = [] while len(self.state.get_step_queue(task_id)) > 0: steps.append(self.state.get_step_queue(task_id).popleft()) # take every item in the queue, and run it tasks = [] for step in steps: tasks.append(self._arun_step(task_id, step=step, mode=mode, **kwargs)) return await asyncio.gather(*tasks)` |
#### run\_step [#](#llama_index.core.agent.ParallelAgentRunner.run_step "Permanent link")
`run_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 278
279
280
281
282
283
284
285
286 | `def run_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step.""" return self._run_step(task_id, step, mode=ChatResponseMode.WAIT, **kwargs)` |
#### arun\_step `async` [#](#llama_index.core.agent.ParallelAgentRunner.arun_step "Permanent link")
`arun_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 288
289
290
291
292
293
294
295
296
297
298 | `async def arun_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async).""" return await self._arun_step( task_id, step, mode=ChatResponseMode.WAIT, **kwargs )` |
#### stream\_step [#](#llama_index.core.agent.ParallelAgentRunner.stream_step "Permanent link")
`stream_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 300
301
302
303
304
305
306
307
308 | `def stream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (stream).""" return self._run_step(task_id, step, mode=ChatResponseMode.STREAM, **kwargs)` |
#### astream\_step `async` [#](#llama_index.core.agent.ParallelAgentRunner.astream_step "Permanent link")
`astream_step(task_id: str, input: Optional[str] = None, step: Optional[[TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ] = None, **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 310
311
312
313
314
315
316
317
318
319
320 | `async def astream_step( self, task_id: str, input: Optional[str] = None, step: Optional[TaskStep] = None, **kwargs: Any, ) -> TaskStepOutput: """Run step (async stream).""" return await self._arun_step( task_id, step, mode=ChatResponseMode.STREAM, **kwargs )` |
#### finalize\_response [#](#llama_index.core.agent.ParallelAgentRunner.finalize_response "Permanent link")
`finalize_response(task_id: str, step_output: Optional[[TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput") ] = None) -> AGENT_CHAT_RESPONSE_TYPE`
Finalize response.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350 | ``def finalize_response( self, task_id: str, step_output: Optional[TaskStepOutput] = None, ) -> AGENT_CHAT_RESPONSE_TYPE: """Finalize response.""" if step_output is None: step_output = self.state.get_completed_steps(task_id)[-1] if not step_output.is_last: raise ValueError( "finalize_response can only be called on the last step output" ) if not isinstance( step_output.output, (AgentChatResponse, StreamingAgentChatResponse), ): raise ValueError( "When `is_last` is True, cur_step_output.output must be " f"AGENT_CHAT_RESPONSE_TYPE: {step_output.output}" ) # finalize task self.agent_worker.finalize_task(self.state.get_task(task_id)) if self.delete_task_on_finish: self.delete_task(task_id) return cast(AGENT_CHAT_RESPONSE_TYPE, step_output.output)`` |
#### undo\_step [#](#llama_index.core.agent.ParallelAgentRunner.undo_step "Permanent link")
`undo_step(task_id: str) -> None`
Undo previous step.
Source code in `llama-index-core/llama_index/core/agent/runner/parallel.py`
| | |
| --- | --- |
| 494
495
496 | `def undo_step(self, task_id: str) -> None: """Undo previous step.""" raise NotImplementedError("undo_step not implemented")` |
Workers[#](#workers "Permanent link")
--------------------------------------
### CustomSimpleAgentWorker [#](#llama_index.core.agent.CustomSimpleAgentWorker "Permanent link")
Bases: `BaseModel`, `[BaseAgentWorker](#llama_index.core.agent.types.BaseAgentWorker "llama_index.core.agent.types.BaseAgentWorker") `
Custom simple agent worker.
This is "simple" in the sense that some of the scaffolding is setup already. Assumptions: - assumes that the agent has tools, llm, callback manager, and tool retriever - has a `from_tools` convenience function - assumes that the agent is sequential, and doesn't take in any additional intermediate inputs.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `tools` | `Sequence[[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]` | Tools to use for reasoning | _required_ |
| `llm` | `[LLM](../llms/#llama_index.core.llms.llm.LLM "llama_index.core.llms.llm.LLM") ` | LLM to use | _required_ |
| `callback_manager` | `[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ` | Callback manager | _required_ |
| `tool_retriever` | `Optional[[ObjectRetriever](../objects/#llama_index.core.objects.ObjectRetriever "llama_index.core.objects.base.ObjectRetriever") [[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]]` | Tool retriever | _required_ |
| `verbose` | `bool` | Whether to print out reasoning steps | _required_ |
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 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 | ``class CustomSimpleAgentWorker(BaseModel, BaseAgentWorker): """Custom simple agent worker. This is "simple" in the sense that some of the scaffolding is setup already. Assumptions: - assumes that the agent has tools, llm, callback manager, and tool retriever - has a `from_tools` convenience function - assumes that the agent is sequential, and doesn't take in any additional intermediate inputs. Args: tools (Sequence[BaseTool]): Tools to use for reasoning llm (LLM): LLM to use callback_manager (CallbackManager): Callback manager tool_retriever (Optional[ObjectRetriever[BaseTool]]): Tool retriever verbose (bool): Whether to print out reasoning steps """ model_config = ConfigDict(arbitrary_types_allowed=True) tools: Sequence[BaseTool] = Field(..., description="Tools to use for reasoning") llm: LLM = Field(..., description="LLM to use") callback_manager: CallbackManager = Field( default_factory=lambda: CallbackManager([]), exclude=True ) tool_retriever: Optional[ObjectRetriever[BaseTool]] = Field( default=None, description="Tool retriever" ) verbose: bool = Field(False, description="Whether to print out reasoning steps") _get_tools: Callable[[str], Sequence[BaseTool]] = PrivateAttr() def __init__( self, tools: Sequence[BaseTool], llm: LLM, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, **kwargs: Any, ) -> None: callback_manager = callback_manager or CallbackManager([]) super().__init__( tools=tools, llm=llm, callback_manager=callback_manager or CallbackManager([]), tool_retriever=tool_retriever, verbose=verbose, **kwargs, ) if len(tools) > 0 and tool_retriever is not None: raise ValueError("Cannot specify both tools and tool_retriever") elif len(tools) > 0: self._get_tools = lambda _: tools elif tool_retriever is not None: tool_retriever_c = cast(ObjectRetriever[BaseTool], tool_retriever) self._get_tools = lambda message: tool_retriever_c.retrieve(message) else: self._get_tools = lambda _: [] @classmethod def from_tools( cls, tools: Optional[Sequence[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, llm: Optional[LLM] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, **kwargs: Any, ) -> "CustomSimpleAgentWorker": """Convenience constructor method from set of BaseTools (Optional).""" llm = llm or Settings.llm if callback_manager is not None: llm.callback_manager = callback_manager return cls( tools=tools or [], tool_retriever=tool_retriever, llm=llm, callback_manager=callback_manager or CallbackManager([]), verbose=verbose, **kwargs, ) @abstractmethod def _initialize_state(self, task: Task, **kwargs: Any) -> Dict[str, Any]: """Initialize state.""" 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() # initialize initial state initial_state = { "sources": sources, "memory": new_memory, } step_state = self._initialize_state(task, **kwargs) # if intersecting keys, error if set(step_state.keys()).intersection(set(initial_state.keys())): raise ValueError( f"Step state keys {step_state.keys()} and initial state keys {initial_state.keys()} intersect." f"*NOTE*: initial state keys {initial_state.keys()} are reserved." ) step_state.update(initial_state) return TaskStep( task_id=task.task_id, step_id=str(uuid.uuid4()), input=task.input, step_state=step_state, ) def get_tools(self, input: str) -> List[AsyncBaseTool]: """Get tools.""" return [adapt_to_async_tool(t) for t in self._get_tools(input)] def _get_task_step_response( self, agent_response: AGENT_CHAT_RESPONSE_TYPE, step: TaskStep, is_done: bool ) -> TaskStepOutput: """Get task step response.""" if is_done: new_steps = [] else: new_steps = [ step.get_next_step( step_id=str(uuid.uuid4()), # NOTE: input is unused input=None, ) ] return TaskStepOutput( output=agent_response, task_step=step, is_last=is_done, next_steps=new_steps, ) @abstractmethod def _run_step( self, state: Dict[str, Any], task: Task, input: Optional[str] = None ) -> Tuple[AgentChatResponse, bool]: """Run step. Returns: Tuple of (agent_response, is_done) """ async def _arun_step( self, state: Dict[str, Any], task: Task, input: Optional[str] = None ) -> Tuple[AgentChatResponse, bool]: """Run step (async). Can override this method if you want to run the step asynchronously. Returns: Tuple of (agent_response, is_done) """ raise NotImplementedError( "This agent does not support async." "Please implement _arun_step." ) @trace_method("run_step") def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" agent_response, is_done = self._run_step( step.step_state, task, input=step.input ) response = self._get_task_step_response(agent_response, step, is_done) # sync step state with task state task.extra_state.update(step.step_state) return response @trace_method("run_step") async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" agent_response, is_done = await self._arun_step( step.step_state, task, input=step.input ) response = self._get_task_step_response(agent_response, step, is_done) task.extra_state.update(step.step_state) return response @trace_method("run_step") def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" raise NotImplementedError("This agent does not support streaming.") @trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError("This agent does not support streaming.") @abstractmethod def _finalize_task(self, state: Dict[str, Any], **kwargs: Any) -> None: """Finalize task, after all the steps are completed. State is all the step states. """ def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set(task.memory.get() + task.extra_state["memory"].get_all()) # reset new memory task.extra_state["memory"].reset() self._finalize_task(task.extra_state, **kwargs) def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager`` |
#### from\_tools `classmethod` [#](#llama_index.core.agent.CustomSimpleAgentWorker.from_tools "Permanent link")
`from_tools(tools: Optional[Sequence[[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, tool_retriever: Optional[[ObjectRetriever](../objects/#llama_index.core.objects.ObjectRetriever "llama_index.core.objects.base.ObjectRetriever") [[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, llm: Optional[[LLM](../llms/#llama_index.core.llms.llm.LLM "llama_index.core.llms.llm.LLM") ] = None, callback_manager: Optional[[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, verbose: bool = False, **kwargs: Any) -> [CustomSimpleAgentWorker](#llama_index.core.agent.CustomSimpleAgentWorker "llama_index.core.agent.custom.simple.CustomSimpleAgentWorker")`
Convenience constructor method from set of BaseTools (Optional).
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121 | `@classmethod def from_tools( cls, tools: Optional[Sequence[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, llm: Optional[LLM] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, **kwargs: Any, ) -> "CustomSimpleAgentWorker": """Convenience constructor method from set of BaseTools (Optional).""" llm = llm or Settings.llm if callback_manager is not None: llm.callback_manager = callback_manager return cls( tools=tools or [], tool_retriever=tool_retriever, llm=llm, callback_manager=callback_manager or CallbackManager([]), verbose=verbose, **kwargs, )` |
#### initialize\_step [#](#llama_index.core.agent.CustomSimpleAgentWorker.initialize_step "Permanent link")
`initialize_step(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep")`
Initialize step from task.
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 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 | `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() # initialize initial state initial_state = { "sources": sources, "memory": new_memory, } step_state = self._initialize_state(task, **kwargs) # if intersecting keys, error if set(step_state.keys()).intersection(set(initial_state.keys())): raise ValueError( f"Step state keys {step_state.keys()} and initial state keys {initial_state.keys()} intersect." f"*NOTE*: initial state keys {initial_state.keys()} are reserved." ) step_state.update(initial_state) return TaskStep( task_id=task.task_id, step_id=str(uuid.uuid4()), input=task.input, step_state=step_state, )` |
#### get\_tools [#](#llama_index.core.agent.CustomSimpleAgentWorker.get_tools "Permanent link")
`get_tools(input: str) -> List[[AsyncBaseTool](../tools/#llama_index.core.tools.types.AsyncBaseTool "llama_index.core.tools.types.AsyncBaseTool") ]`
Get tools.
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 155
156
157 | `def get_tools(self, input: str) -> List[AsyncBaseTool]: """Get tools.""" return [adapt_to_async_tool(t) for t in self._get_tools(input)]` |
#### run\_step [#](#llama_index.core.agent.CustomSimpleAgentWorker.run_step "Permanent link")
`run_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 207
208
209
210
211
212
213
214
215
216 | `@trace_method("run_step") def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" agent_response, is_done = self._run_step( step.step_state, task, input=step.input ) response = self._get_task_step_response(agent_response, step, is_done) # sync step state with task state task.extra_state.update(step.step_state) return response` |
#### arun\_step `async` [#](#llama_index.core.agent.CustomSimpleAgentWorker.arun_step "Permanent link")
`arun_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 218
219
220
221
222
223
224
225
226
227
228 | `@trace_method("run_step") async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" agent_response, is_done = await self._arun_step( step.step_state, task, input=step.input ) response = self._get_task_step_response(agent_response, step, is_done) task.extra_state.update(step.step_state) return response` |
#### stream\_step [#](#llama_index.core.agent.CustomSimpleAgentWorker.stream_step "Permanent link")
`stream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 230
231
232
233 | `@trace_method("run_step") def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" raise NotImplementedError("This agent does not support streaming.")` |
#### astream\_step `async` [#](#llama_index.core.agent.CustomSimpleAgentWorker.astream_step "Permanent link")
`astream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 235
236
237
238
239
240 | `@trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError("This agent does not support streaming.")` |
#### finalize\_task [#](#llama_index.core.agent.CustomSimpleAgentWorker.finalize_task "Permanent link")
`finalize_task(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> None`
Finalize task, after all the steps are completed.
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 250
251
252
253
254
255
256 | `def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set(task.memory.get() + task.extra_state["memory"].get_all()) # reset new memory task.extra_state["memory"].reset() self._finalize_task(task.extra_state, **kwargs)` |
#### set\_callback\_manager [#](#llama_index.core.agent.CustomSimpleAgentWorker.set_callback_manager "Permanent link")
`set_callback_manager(callback_manager: [CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ) -> None`
Set callback manager.
Source code in `llama-index-core/llama_index/core/agent/custom/simple.py`
| | |
| --- | --- |
| 258
259
260
261 | `def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager` |
### MultimodalReActAgentWorker [#](#llama_index.core.agent.MultimodalReActAgentWorker "Permanent link")
Bases: `[BaseAgentWorker](#llama_index.core.agent.types.BaseAgentWorker "llama_index.core.agent.types.BaseAgentWorker") `
Multimodal ReAct Agent worker.
**NOTE**: This is a BETA feature.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 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
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531 | ``class MultimodalReActAgentWorker(BaseAgentWorker): """Multimodal ReAct Agent worker. **NOTE**: This is a BETA feature. """ def __init__( self, tools: Sequence[BaseTool], multi_modal_llm: MultiModalLLM, max_iterations: int = 10, react_chat_formatter: Optional[ReActChatFormatter] = None, output_parser: Optional[ReActOutputParser] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, generate_chat_message_fn: Optional[ChatMessageCallable] = None, ) -> None: self._multi_modal_llm = multi_modal_llm self.callback_manager = callback_manager or CallbackManager([]) self._max_iterations = max_iterations self._react_chat_formatter = react_chat_formatter or ReActChatFormatter( system_header=REACT_MM_CHAT_SYSTEM_HEADER ) self._output_parser = output_parser or ReActOutputParser() self._verbose = verbose self._generate_chat_message_fn = generate_chat_message_fn if self._generate_chat_message_fn is None: try: from llama_index.multi_modal_llms.openai.utils import ( generate_openai_multi_modal_chat_message, ) # pants: no-infer-dep self._generate_chat_message_fn = ( generate_openai_multi_modal_chat_message ) except ImportError: raise ImportError( "`llama-index-multi-modal-llms-openai` package cannot be found. " "Please install it by using `pip install `llama-index-multi-modal-llms-openai`" ) self._add_user_step_to_reasoning = partial( add_user_step_to_reasoning, generate_chat_message_fn=self._generate_chat_message_fn, # type: ignore ) if len(tools) > 0 and tool_retriever is not None: raise ValueError("Cannot specify both tools and tool_retriever") elif len(tools) > 0: self._get_tools = lambda _: tools elif tool_retriever is not None: tool_retriever_c = cast(ObjectRetriever[BaseTool], tool_retriever) self._get_tools = lambda message: tool_retriever_c.retrieve(message) else: self._get_tools = lambda _: [] @classmethod def from_tools( cls, tools: Optional[Sequence[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, multi_modal_llm: Optional[MultiModalLLM] = None, max_iterations: int = 10, react_chat_formatter: Optional[ReActChatFormatter] = None, output_parser: Optional[ReActOutputParser] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, **kwargs: Any, ) -> "MultimodalReActAgentWorker": """Convenience constructor method from set of BaseTools (Optional). NOTE: kwargs should have been exhausted by this point. In other words the various upstream components such as BaseSynthesizer (response synthesizer) or BaseRetriever should have picked up off their respective kwargs in their constructions. Returns: ReActAgent """ if multi_modal_llm is None: try: from llama_index.multi_modal_llms.openai import ( OpenAIMultiModal, ) # pants: no-infer-dep multi_modal_llm = multi_modal_llm or OpenAIMultiModal( model="gpt-4-vision-preview", max_new_tokens=1000 ) except ImportError: raise ImportError( "`llama-index-multi-modal-llms-openai` package cannot be found. " "Please install it by using `pip install `llama-index-multi-modal-llms-openai`" ) return cls( tools=tools or [], tool_retriever=tool_retriever, multi_modal_llm=multi_modal_llm, max_iterations=max_iterations, react_chat_formatter=react_chat_formatter, output_parser=output_parser, callback_manager=callback_manager, verbose=verbose, ) def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep: """Initialize step from task.""" sources: List[ToolOutput] = [] current_reasoning: List[BaseReasoningStep] = [] # temporary memory for new messages new_memory = ChatMemoryBuffer.from_defaults() # validation if "image_docs" not in task.extra_state: raise ValueError("Image docs not found in task extra state.") # initialize task state task_state = { "sources": sources, "current_reasoning": current_reasoning, "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={"is_first": True, "image_docs": task.extra_state["image_docs"]}, ) def get_tools(self, input: str) -> List[AsyncBaseTool]: """Get tools.""" return [adapt_to_async_tool(t) for t in self._get_tools(input)] def _extract_reasoning_step( self, output: ChatResponse, is_streaming: bool = False ) -> Tuple[str, List[BaseReasoningStep], bool]: """ Extracts the reasoning step from the given output. This method parses the message content from the output, extracts the reasoning step, and determines whether the processing is complete. It also performs validation checks on the output and handles possible errors. """ if output.message.content is None: raise ValueError("Got empty message.") message_content = output.message.content current_reasoning = [] try: reasoning_step = self._output_parser.parse(message_content, is_streaming) except BaseException as exc: raise ValueError(f"Could not parse output: {message_content}") from exc if self._verbose: print_text(f"{reasoning_step.get_content()}\n", color="pink") current_reasoning.append(reasoning_step) if reasoning_step.is_done: return message_content, current_reasoning, True reasoning_step = cast(ActionReasoningStep, reasoning_step) if not isinstance(reasoning_step, ActionReasoningStep): raise ValueError(f"Expected ActionReasoningStep, got {reasoning_step}") return message_content, current_reasoning, False def _process_actions( self, task: Task, tools: Sequence[AsyncBaseTool], output: ChatResponse, is_streaming: bool = False, ) -> Tuple[List[BaseReasoningStep], bool]: tools_dict: Dict[str, AsyncBaseTool] = { tool.metadata.get_name(): tool for tool in tools } _, current_reasoning, is_done = self._extract_reasoning_step( output, is_streaming ) if is_done: return current_reasoning, True # call tool with input reasoning_step = cast(ActionReasoningStep, current_reasoning[-1]) tool = tools_dict[reasoning_step.action] with self.callback_manager.event( CBEventType.FUNCTION_CALL, payload={ EventPayload.FUNCTION_CALL: reasoning_step.action_input, EventPayload.TOOL: tool.metadata, }, ) as event: tool_output = tool.call(**reasoning_step.action_input) event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)}) task.extra_state["sources"].append(tool_output) observation_step = ObservationReasoningStep( observation=str(tool_output), return_direct=tool.metadata.return_direct ) current_reasoning.append(observation_step) if self._verbose: print_text(f"{observation_step.get_content()}\n", color="blue") return current_reasoning, tool.metadata.return_direct async def _aprocess_actions( self, task: Task, tools: Sequence[AsyncBaseTool], output: ChatResponse, is_streaming: bool = False, ) -> Tuple[List[BaseReasoningStep], bool]: tools_dict = {tool.metadata.name: tool for tool in tools} _, current_reasoning, is_done = self._extract_reasoning_step( output, is_streaming ) if is_done: return current_reasoning, True # call tool with input reasoning_step = cast(ActionReasoningStep, current_reasoning[-1]) tool = tools_dict[reasoning_step.action] with self.callback_manager.event( CBEventType.FUNCTION_CALL, payload={ EventPayload.FUNCTION_CALL: reasoning_step.action_input, EventPayload.TOOL: tool.metadata, }, ) as event: tool_output = await tool.acall(**reasoning_step.action_input) event.on_end(payload={EventPayload.FUNCTION_OUTPUT: str(tool_output)}) task.extra_state["sources"].append(tool_output) observation_step = ObservationReasoningStep( observation=str(tool_output), return_direct=tool.metadata.return_direct ) current_reasoning.append(observation_step) if self._verbose: print_text(f"{observation_step.get_content()}\n", color="blue") return current_reasoning, tool.metadata.return_direct def _get_response( self, current_reasoning: List[BaseReasoningStep], sources: List[ToolOutput], ) -> AgentChatResponse: """Get response from reasoning steps.""" if len(current_reasoning) == 0: raise ValueError("No reasoning steps were taken.") elif len(current_reasoning) == self._max_iterations: raise ValueError("Reached max iterations.") if isinstance(current_reasoning[-1], ResponseReasoningStep): response_step = cast(ResponseReasoningStep, current_reasoning[-1]) response_str = response_step.response elif ( isinstance(current_reasoning[-1], ObservationReasoningStep) and current_reasoning[-1].return_direct ): response_str = current_reasoning[-1].observation else: response_str = current_reasoning[-1].get_content() # TODO: add sources from reasoning steps return AgentChatResponse(response=response_str, sources=sources) def _get_task_step_response( self, agent_response: AGENT_CHAT_RESPONSE_TYPE, step: TaskStep, is_done: bool ) -> TaskStepOutput: """Get task step response.""" if is_done: new_steps = [] else: new_steps = [ step.get_next_step( step_id=str(uuid.uuid4()), # NOTE: input is unused input=None, ) ] return TaskStepOutput( output=agent_response, task_step=step, is_last=is_done, next_steps=new_steps, ) def _run_step( self, step: TaskStep, task: Task, ) -> TaskStepOutput: """Run step.""" # This is either not None on the first step or if the user specifies # an intermediate step in the middle if step.input is not None: self._add_user_step_to_reasoning( step=step, memory=task.extra_state["new_memory"], current_reasoning=task.extra_state["current_reasoning"], verbose=self._verbose, ) # TODO: see if we want to do step-based inputs tools = self.get_tools(task.input) input_chat = self._react_chat_formatter.format( tools, chat_history=task.memory.get_all() + task.extra_state["new_memory"].get_all(), current_reasoning=task.extra_state["current_reasoning"], ) # send prompt chat_response = self._multi_modal_llm.chat(input_chat) # given react prompt outputs, call tools or return response reasoning_steps, is_done = self._process_actions( task, tools, output=chat_response ) task.extra_state["current_reasoning"].extend(reasoning_steps) agent_response = self._get_response( task.extra_state["current_reasoning"], task.extra_state["sources"] ) if is_done: task.extra_state["new_memory"].put( ChatMessage(content=agent_response.response, role=MessageRole.ASSISTANT) ) return self._get_task_step_response(agent_response, step, is_done) async def _arun_step( self, step: TaskStep, task: Task, ) -> TaskStepOutput: """Run step.""" if step.input is not None: self._add_user_step_to_reasoning( step=step, memory=task.extra_state["new_memory"], current_reasoning=task.extra_state["current_reasoning"], verbose=self._verbose, ) # TODO: see if we want to do step-based inputs tools = self.get_tools(task.input) input_chat = self._react_chat_formatter.format( tools, chat_history=task.memory.get_all() + task.extra_state["new_memory"].get_all(), current_reasoning=task.extra_state["current_reasoning"], ) # send prompt chat_response = await self._multi_modal_llm.achat(input_chat) # given react prompt outputs, call tools or return response reasoning_steps, is_done = await self._aprocess_actions( task, tools, output=chat_response ) task.extra_state["current_reasoning"].extend(reasoning_steps) agent_response = self._get_response( task.extra_state["current_reasoning"], task.extra_state["sources"] ) if is_done: task.extra_state["new_memory"].put( ChatMessage(content=agent_response.response, role=MessageRole.ASSISTANT) ) return self._get_task_step_response(agent_response, step, is_done) def _run_step_stream( self, step: TaskStep, task: Task, ) -> TaskStepOutput: """Run step.""" raise NotImplementedError("Stream step not implemented yet.") async def _arun_step_stream( self, step: TaskStep, task: Task, ) -> TaskStepOutput: """Run step.""" raise NotImplementedError("Stream step not implemented yet.") @trace_method("run_step") def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" return self._run_step(step, task) @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).""" # TODO: figure out if we need a different type for TaskStepOutput return self._run_step_stream(step, task) @trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" return await self._arun_step_stream(step, task) def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set( task.memory.get_all() + task.extra_state["new_memory"].get_all() ) # reset new memory task.extra_state["new_memory"].reset() def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager`` |
#### from\_tools `classmethod` [#](#llama_index.core.agent.MultimodalReActAgentWorker.from_tools "Permanent link")
`from_tools(tools: Optional[Sequence[[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, tool_retriever: Optional[[ObjectRetriever](../objects/#llama_index.core.objects.ObjectRetriever "llama_index.core.objects.base.ObjectRetriever") [[BaseTool](../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, multi_modal_llm: Optional[[MultiModalLLM](../multi_modal_llms/#llama_index.core.multi_modal_llms.base.MultiModalLLM "llama_index.core.multi_modal_llms.base.MultiModalLLM") ] = None, max_iterations: int = 10, react_chat_formatter: Optional[[ReActChatFormatter](react/#llama_index.core.agent.react.ReActChatFormatter "llama_index.core.agent.react.formatter.ReActChatFormatter") ] = None, output_parser: Optional[ReActOutputParser] = None, callback_manager: Optional[[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, verbose: bool = False, **kwargs: Any) -> [MultimodalReActAgentWorker](#llama_index.core.agent.MultimodalReActAgentWorker "llama_index.core.agent.react_multimodal.step.MultimodalReActAgentWorker")`
Convenience constructor method from set of BaseTools (Optional).
NOTE: kwargs should have been exhausted by this point. In other words the various upstream components such as BaseSynthesizer (response synthesizer) or BaseRetriever should have picked up off their respective kwargs in their constructions.
Returns:
| Type | Description |
| --- | --- |
| `[MultimodalReActAgentWorker](#llama_index.core.agent.MultimodalReActAgentWorker "llama_index.core.agent.react_multimodal.step.MultimodalReActAgentWorker") ` | ReActAgent |
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 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 | ``@classmethod def from_tools( cls, tools: Optional[Sequence[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, multi_modal_llm: Optional[MultiModalLLM] = None, max_iterations: int = 10, react_chat_formatter: Optional[ReActChatFormatter] = None, output_parser: Optional[ReActOutputParser] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, **kwargs: Any, ) -> "MultimodalReActAgentWorker": """Convenience constructor method from set of BaseTools (Optional). NOTE: kwargs should have been exhausted by this point. In other words the various upstream components such as BaseSynthesizer (response synthesizer) or BaseRetriever should have picked up off their respective kwargs in their constructions. Returns: ReActAgent """ if multi_modal_llm is None: try: from llama_index.multi_modal_llms.openai import ( OpenAIMultiModal, ) # pants: no-infer-dep multi_modal_llm = multi_modal_llm or OpenAIMultiModal( model="gpt-4-vision-preview", max_new_tokens=1000 ) except ImportError: raise ImportError( "`llama-index-multi-modal-llms-openai` package cannot be found. " "Please install it by using `pip install `llama-index-multi-modal-llms-openai`" ) return cls( tools=tools or [], tool_retriever=tool_retriever, multi_modal_llm=multi_modal_llm, max_iterations=max_iterations, react_chat_formatter=react_chat_formatter, output_parser=output_parser, callback_manager=callback_manager, verbose=verbose, )`` |
#### initialize\_step [#](#llama_index.core.agent.MultimodalReActAgentWorker.initialize_step "Permanent link")
`initialize_step(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep")`
Initialize step from task.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 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 | `def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep: """Initialize step from task.""" sources: List[ToolOutput] = [] current_reasoning: List[BaseReasoningStep] = [] # temporary memory for new messages new_memory = ChatMemoryBuffer.from_defaults() # validation if "image_docs" not in task.extra_state: raise ValueError("Image docs not found in task extra state.") # initialize task state task_state = { "sources": sources, "current_reasoning": current_reasoning, "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={"is_first": True, "image_docs": task.extra_state["image_docs"]}, )` |
#### get\_tools [#](#llama_index.core.agent.MultimodalReActAgentWorker.get_tools "Permanent link")
`get_tools(input: str) -> List[[AsyncBaseTool](../tools/#llama_index.core.tools.types.AsyncBaseTool "llama_index.core.tools.types.AsyncBaseTool") ]`
Get tools.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 235
236
237 | `def get_tools(self, input: str) -> List[AsyncBaseTool]: """Get tools.""" return [adapt_to_async_tool(t) for t in self._get_tools(input)]` |
#### run\_step [#](#llama_index.core.agent.MultimodalReActAgentWorker.run_step "Permanent link")
`run_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 494
495
496
497 | `@trace_method("run_step") def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" return self._run_step(step, task)` |
#### arun\_step `async` [#](#llama_index.core.agent.MultimodalReActAgentWorker.arun_step "Permanent link")
`arun_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 499
500
501
502
503
504 | `@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 [#](#llama_index.core.agent.MultimodalReActAgentWorker.stream_step "Permanent link")
`stream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 506
507
508
509
510 | `@trace_method("run_step") def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" # TODO: figure out if we need a different type for TaskStepOutput return self._run_step_stream(step, task)` |
#### astream\_step `async` [#](#llama_index.core.agent.MultimodalReActAgentWorker.astream_step "Permanent link")
`astream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 512
513
514
515
516
517 | `@trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" return await self._arun_step_stream(step, task)` |
#### finalize\_task [#](#llama_index.core.agent.MultimodalReActAgentWorker.finalize_task "Permanent link")
`finalize_task(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> None`
Finalize task, after all the steps are completed.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 519
520
521
522
523
524
525
526 | `def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set( task.memory.get_all() + task.extra_state["new_memory"].get_all() ) # reset new memory task.extra_state["new_memory"].reset()` |
#### set\_callback\_manager [#](#llama_index.core.agent.MultimodalReActAgentWorker.set_callback_manager "Permanent link")
`set_callback_manager(callback_manager: [CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ) -> None`
Set callback manager.
Source code in `llama-index-core/llama_index/core/agent/react_multimodal/step.py`
| | |
| --- | --- |
| 528
529
530
531 | `def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager` |
### QueryPipelineAgentWorker [#](#llama_index.core.agent.QueryPipelineAgentWorker "Permanent link")
Bases: `BaseModel`, `[BaseAgentWorker](#llama_index.core.agent.types.BaseAgentWorker "llama_index.core.agent.types.BaseAgentWorker") `
Query Pipeline agent worker.
NOTE: This is now deprecated. Use `FnAgentWorker` instead to build a stateful agent.
Barebones agent worker that takes in a query pipeline.
**Default Workflow**: The default workflow assumes that you compose a query pipeline with `StatefulFnComponent` objects. This allows you to store, update and retrieve state throughout the executions of the query pipeline by the agent.
The task and step state of the agent are stored in this `state` variable via a special key. Of course you can choose to store other variables in this state as well.
**Deprecated Workflow**: The deprecated workflow assumes that the first component in the query pipeline is an `AgentInputComponent` and last is `AgentFnComponent`.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `pipeline` | `QueryPipeline` | Query pipeline | _required_ |
| `callback_manager` | `[CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ` | | _required_ |
| `task_key` | `str` | Key to store task in state | `'task'` |
| `step_state_key` | `str` | Key to store step in state | `'step_state'` |
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 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 | ``@deprecated("Use `FnAgentWorker` instead to build a stateful agent.") class QueryPipelineAgentWorker(BaseModel, BaseAgentWorker): """Query Pipeline agent worker. NOTE: This is now deprecated. Use `FnAgentWorker` instead to build a stateful agent. Barebones agent worker that takes in a query pipeline. **Default Workflow**: The default workflow assumes that you compose a query pipeline with `StatefulFnComponent` objects. This allows you to store, update and retrieve state throughout the executions of the query pipeline by the agent. The task and step state of the agent are stored in this `state` variable via a special key. Of course you can choose to store other variables in this state as well. **Deprecated Workflow**: The deprecated workflow assumes that the first component in the query pipeline is an `AgentInputComponent` and last is `AgentFnComponent`. Args: pipeline (QueryPipeline): Query pipeline """ model_config = ConfigDict(arbitrary_types_allowed=True) pipeline: QueryPipeline = Field(..., description="Query pipeline") callback_manager: CallbackManager = Field(..., exclude=True) task_key: str = Field("task", description="Key to store task in state") step_state_key: str = Field("step_state", description="Key to store step in state") def __init__( self, pipeline: QueryPipeline, callback_manager: Optional[CallbackManager] = None, **kwargs: Any, ) -> None: """Initialize.""" if callback_manager is not None: # set query pipeline callback pipeline.set_callback_manager(callback_manager) else: callback_manager = pipeline.callback_manager super().__init__( pipeline=pipeline, callback_manager=callback_manager, **kwargs, ) # validate query pipeline # self.agent_input_component self.agent_components @property def agent_input_component(self) -> AgentInputComponent: """Get agent input component. NOTE: This is deprecated and will be removed in the future. """ root_key = self.pipeline.get_root_keys()[0] if not isinstance(self.pipeline.module_dict[root_key], AgentInputComponent): raise ValueError( "Query pipeline first component must be AgentInputComponent, got " f"{self.pipeline.module_dict[root_key]}" ) return cast(AgentInputComponent, self.pipeline.module_dict[root_key]) @property def agent_components(self) -> Sequence[BaseAgentComponent]: """Get agent output component.""" return _get_agent_components(self.pipeline) def preprocess(self, task: Task, step: TaskStep) -> None: """Preprocessing flow. This runs preprocessing to propagate the task and step as variables to relevant components in the query pipeline. Contains deprecated flow of updating agent components. But also contains main flow of updating StatefulFnComponent components. """ # NOTE: this is deprecated # partial agent output component with task and step for agent_fn_component in self.agent_components: agent_fn_component.partial(task=task, state=step.step_state) # update stateful components self.pipeline.update_state( {self.task_key: task, self.step_state_key: step.step_state} ) 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() # initialize initial state initial_state = { "sources": sources, "memory": new_memory, } return TaskStep( task_id=task.task_id, step_id=str(uuid.uuid4()), input=task.input, step_state=initial_state, ) def _get_task_step_response( self, agent_response: AGENT_CHAT_RESPONSE_TYPE, step: TaskStep, is_done: bool ) -> TaskStepOutput: """Get task step response.""" if is_done: new_steps = [] else: new_steps = [ step.get_next_step( step_id=str(uuid.uuid4()), # NOTE: input is unused input=None, ) ] 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.""" self.preprocess(task, step) # HACK: do a try/except for now. Fine since old agent components are deprecated try: self.agent_input_component uses_deprecated = True except ValueError: uses_deprecated = False if uses_deprecated: agent_response, is_done = self.pipeline.run( state=step.step_state, task=task ) else: agent_response, is_done = self.pipeline.run() response = self._get_task_step_response(agent_response, step, is_done) # sync step state with task state task.extra_state.update(step.step_state) return response @trace_method("run_step") async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" self.preprocess(task, step) # HACK: do a try/except for now. Fine since old agent components are deprecated try: self.agent_input_component uses_deprecated = True except ValueError: uses_deprecated = False if uses_deprecated: agent_response, is_done = await self.pipeline.arun( state=step.step_state, task=task ) else: agent_response, is_done = await self.pipeline.arun() response = self._get_task_step_response(agent_response, step, is_done) task.extra_state.update(step.step_state) return response @trace_method("run_step") def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" raise NotImplementedError("This agent does not support streaming.") @trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError("This agent does not support streaming.") def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set(task.memory.get() + task.extra_state["memory"].get_all()) # reset new memory task.extra_state["memory"].reset() def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager self.pipeline.set_callback_manager(callback_manager)`` |
#### agent\_input\_component `property` [#](#llama_index.core.agent.QueryPipelineAgentWorker.agent_input_component "Permanent link")
`agent_input_component: AgentInputComponent`
Get agent input component.
NOTE: This is deprecated and will be removed in the future.
#### agent\_components `property` [#](#llama_index.core.agent.QueryPipelineAgentWorker.agent_components "Permanent link")
`agent_components: Sequence[BaseAgentComponent]`
Get agent output component.
#### preprocess [#](#llama_index.core.agent.QueryPipelineAgentWorker.preprocess "Permanent link")
`preprocess(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") ) -> None`
Preprocessing flow.
This runs preprocessing to propagate the task and step as variables to relevant components in the query pipeline.
Contains deprecated flow of updating agent components. But also contains main flow of updating StatefulFnComponent components.
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143 | `def preprocess(self, task: Task, step: TaskStep) -> None: """Preprocessing flow. This runs preprocessing to propagate the task and step as variables to relevant components in the query pipeline. Contains deprecated flow of updating agent components. But also contains main flow of updating StatefulFnComponent components. """ # NOTE: this is deprecated # partial agent output component with task and step for agent_fn_component in self.agent_components: agent_fn_component.partial(task=task, state=step.step_state) # update stateful components self.pipeline.update_state( {self.task_key: task, self.step_state_key: step.step_state} )` |
#### initialize\_step [#](#llama_index.core.agent.QueryPipelineAgentWorker.initialize_step "Permanent link")
`initialize_step(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep")`
Initialize step from task.
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162 | `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() # initialize initial state initial_state = { "sources": sources, "memory": new_memory, } return TaskStep( task_id=task.task_id, step_id=str(uuid.uuid4()), input=task.input, step_state=initial_state, )` |
#### run\_step [#](#llama_index.core.agent.QueryPipelineAgentWorker.run_step "Permanent link")
`run_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step.
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207 | `@trace_method("run_step") def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step.""" self.preprocess(task, step) # HACK: do a try/except for now. Fine since old agent components are deprecated try: self.agent_input_component uses_deprecated = True except ValueError: uses_deprecated = False if uses_deprecated: agent_response, is_done = self.pipeline.run( state=step.step_state, task=task ) else: agent_response, is_done = self.pipeline.run() response = self._get_task_step_response(agent_response, step, is_done) # sync step state with task state task.extra_state.update(step.step_state) return response` |
#### arun\_step `async` [#](#llama_index.core.agent.QueryPipelineAgentWorker.arun_step "Permanent link")
`arun_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async).
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231 | `@trace_method("run_step") async def arun_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async).""" self.preprocess(task, step) # HACK: do a try/except for now. Fine since old agent components are deprecated try: self.agent_input_component uses_deprecated = True except ValueError: uses_deprecated = False if uses_deprecated: agent_response, is_done = await self.pipeline.arun( state=step.step_state, task=task ) else: agent_response, is_done = await self.pipeline.arun() response = self._get_task_step_response(agent_response, step, is_done) task.extra_state.update(step.step_state) return response` |
#### stream\_step [#](#llama_index.core.agent.QueryPipelineAgentWorker.stream_step "Permanent link")
`stream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (stream).
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 233
234
235
236 | `@trace_method("run_step") def stream_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput: """Run step (stream).""" raise NotImplementedError("This agent does not support streaming.")` |
#### astream\_step `async` [#](#llama_index.core.agent.QueryPipelineAgentWorker.astream_step "Permanent link")
`astream_step(step: [TaskStep](#llama_index.core.agent.types.TaskStep "llama_index.core.agent.types.TaskStep") , task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> [TaskStepOutput](#llama_index.core.agent.types.TaskStepOutput "llama_index.core.agent.types.TaskStepOutput")`
Run step (async stream).
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 238
239
240
241
242
243 | `@trace_method("run_step") async def astream_step( self, step: TaskStep, task: Task, **kwargs: Any ) -> TaskStepOutput: """Run step (async stream).""" raise NotImplementedError("This agent does not support streaming.")` |
#### finalize\_task [#](#llama_index.core.agent.QueryPipelineAgentWorker.finalize_task "Permanent link")
`finalize_task(task: [Task](#llama_index.core.agent.types.Task "llama_index.core.agent.types.Task") , **kwargs: Any) -> None`
Finalize task, after all the steps are completed.
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 245
246
247
248
249
250 | `def finalize_task(self, task: Task, **kwargs: Any) -> None: """Finalize task, after all the steps are completed.""" # add new messages to memory task.memory.set(task.memory.get() + task.extra_state["memory"].get_all()) # reset new memory task.extra_state["memory"].reset()` |
#### set\_callback\_manager [#](#llama_index.core.agent.QueryPipelineAgentWorker.set_callback_manager "Permanent link")
`set_callback_manager(callback_manager: [CallbackManager](../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ) -> None`
Set callback manager.
Source code in `llama-index-core/llama_index/core/agent/custom/pipeline_worker.py`
| | |
| --- | --- |
| 252
253
254
255
256 | `def set_callback_manager(self, callback_manager: CallbackManager) -> None: """Set callback manager.""" # TODO: make this abstractmethod (right now will break some agent impls) self.callback_manager = callback_manager self.pipeline.set_callback_manager(callback_manager)` |
Back to top

---
# Arize phoenix - LlamaIndex
Arize phoenix
=============
arize\_phoenix\_callback\_handler [#](#llama_index.callbacks.arize_phoenix.arize_phoenix_callback_handler "Permanent link")
----------------------------------------------------------------------------------------------------------------------------
`arize_phoenix_callback_handler(**kwargs: Any) -> [BaseCallbackHandler](../#llama_index.core.callbacks.base_handler.BaseCallbackHandler "llama_index.core.callbacks.base_handler.BaseCallbackHandler")`
Source code in `llama-index-integrations/callbacks/llama-index-callbacks-arize-phoenix/llama_index/callbacks/arize_phoenix/base.py`
| | |
| --- | --- |
| 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 | ``def arize_phoenix_callback_handler(**kwargs: Any) -> BaseCallbackHandler: # newer versions of arize, v2.x try: from openinference.instrumentation.llama_index import LlamaIndexInstrumentor from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter, ) from opentelemetry.sdk import trace as trace_sdk from opentelemetry.sdk.trace.export import SimpleSpanProcessor endpoint = kwargs.get("endpoint", "http://127.0.0.1:6006/v1/traces") tracer_provider = trace_sdk.TracerProvider() tracer_provider.add_span_processor( SimpleSpanProcessor(OTLPSpanExporter(endpoint)) ) return LlamaIndexInstrumentor().instrument( tracer_provider=kwargs.get("tracer_provider", tracer_provider) ) except ImportError: # using an older version of arize pass # older versions of arize, v1.x try: from phoenix.trace.llama_index import OpenInferenceTraceCallbackHandler except ImportError: raise ImportError( "Please install Arize Phoenix with `pip install -q arize-phoenix`" ) return OpenInferenceTraceCallbackHandler(**kwargs)`` |
Back to top

---
# Langfuse - LlamaIndex
Langfuse
========
langfuse\_callback\_handler [#](#llama_index.callbacks.langfuse.langfuse_callback_handler "Permanent link")
------------------------------------------------------------------------------------------------------------
`langfuse_callback_handler(**eval_params: Any) -> [BaseCallbackHandler](../#llama_index.core.callbacks.base_handler.BaseCallbackHandler "llama_index.core.callbacks.base_handler.BaseCallbackHandler")`
Source code in `llama-index-integrations/callbacks/llama-index-callbacks-langfuse/llama_index/callbacks/langfuse/base.py`
| | |
| --- | --- |
| 8
9
10
11 | `def langfuse_callback_handler(**eval_params: Any) -> BaseCallbackHandler: return LlamaIndexCallbackHandler( **eval_params, sdk_integration="llama-index_set-global-handler" )` |
Back to top

---
# Agentops - LlamaIndex
Agentops
========
Back to top

---
# Opik - LlamaIndex
Opik
====
opik\_callback\_handler [#](#llama_index.callbacks.opik.opik_callback_handler "Permanent link")
------------------------------------------------------------------------------------------------
`opik_callback_handler(**eval_params: Any) -> [BaseCallbackHandler](../#llama_index.core.callbacks.base_handler.BaseCallbackHandler "llama_index.core.callbacks.base_handler.BaseCallbackHandler")`
Source code in `llama-index-integrations/callbacks/llama-index-callbacks-opik/llama_index/callbacks/opik/base.py`
| | |
| --- | --- |
| 6
7
8
9
10
11
12
13
14
15 | ``def opik_callback_handler(**eval_params: Any) -> BaseCallbackHandler: try: from opik.integrations.llama_index import LlamaIndexCallbackHandler return LlamaIndexCallbackHandler(**eval_params) except ImportError: raise ImportError( "Please install the Opik Python SDK with `pip install -U opik`" )`` |
Back to top

---
# Token counter - LlamaIndex
Token counter
=============
TokenCountingHandler [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler "Permanent link")
-----------------------------------------------------------------------------------------------------------
Bases: `PythonicallyPrintingBaseHandler`
Callback handler for counting tokens in LLM and Embedding events.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `tokenizer` | `Optional[Callable[[str], List]]` | Tokenizer to use. Defaults to the global tokenizer (see llama\_index.core.utils.globals\_helper). | `None` |
| `event_starts_to_ignore` | `Optional[List[[CBEventType](../#llama_index.core.callbacks.schema.CBEventType "llama_index.core.callbacks.schema.CBEventType") ]]` | List of event types to ignore at the start of a trace. | `None` |
| `event_ends_to_ignore` | `Optional[List[[CBEventType](../#llama_index.core.callbacks.schema.CBEventType "llama_index.core.callbacks.schema.CBEventType") ]]` | List of event types to ignore at the end of a trace. | `None` |
Source code in `llama-index-core/llama_index/core/callbacks/token_counting.py`
| | |
| --- | --- |
| 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 | `class TokenCountingHandler(PythonicallyPrintingBaseHandler): """Callback handler for counting tokens in LLM and Embedding events. Args: tokenizer: Tokenizer to use. Defaults to the global tokenizer (see llama_index.core.utils.globals_helper). event_starts_to_ignore: List of event types to ignore at the start of a trace. event_ends_to_ignore: List of event types to ignore at the end of a trace. """ def __init__( self, tokenizer: Optional[Callable[[str], List]] = None, event_starts_to_ignore: Optional[List[CBEventType]] = None, event_ends_to_ignore: Optional[List[CBEventType]] = None, verbose: bool = False, logger: Optional[logging.Logger] = None, ) -> None: self.llm_token_counts: List[TokenCountingEvent] = [] self.embedding_token_counts: List[TokenCountingEvent] = [] self.tokenizer = tokenizer or get_tokenizer() self._token_counter = TokenCounter(tokenizer=self.tokenizer) self._verbose = verbose super().__init__( event_starts_to_ignore=event_starts_to_ignore or [], event_ends_to_ignore=event_ends_to_ignore or [], logger=logger, ) def start_trace(self, trace_id: Optional[str] = None) -> None: return def end_trace( self, trace_id: Optional[str] = None, trace_map: Optional[Dict[str, List[str]]] = None, ) -> None: return def on_event_start( self, event_type: CBEventType, payload: Optional[Dict[str, Any]] = None, event_id: str = "", parent_id: str = "", **kwargs: Any, ) -> str: return event_id def on_event_end( self, event_type: CBEventType, payload: Optional[Dict[str, Any]] = None, event_id: str = "", **kwargs: Any, ) -> None: """Count the LLM or Embedding tokens as needed.""" if ( event_type == CBEventType.LLM and event_type not in self.event_ends_to_ignore and payload is not None ): self.llm_token_counts.append( get_llm_token_counts( token_counter=self._token_counter, payload=payload, event_id=event_id, ) ) if self._verbose: self._print( "LLM Prompt Token Usage: " f"{self.llm_token_counts[-1].prompt_token_count}\n" "LLM Completion Token Usage: " f"{self.llm_token_counts[-1].completion_token_count}", ) elif ( event_type == CBEventType.EMBEDDING and event_type not in self.event_ends_to_ignore and payload is not None ): total_chunk_tokens = 0 for chunk in payload.get(EventPayload.CHUNKS, []): self.embedding_token_counts.append( TokenCountingEvent( event_id=event_id, prompt=chunk, prompt_token_count=self._token_counter.get_string_tokens(chunk), completion="", completion_token_count=0, ) ) total_chunk_tokens += self.embedding_token_counts[-1].total_token_count if self._verbose: self._print(f"Embedding Token Usage: {total_chunk_tokens}") @property def total_llm_token_count(self) -> int: """Get the current total LLM token count.""" return sum([x.total_token_count for x in self.llm_token_counts]) @property def prompt_llm_token_count(self) -> int: """Get the current total LLM prompt token count.""" return sum([x.prompt_token_count for x in self.llm_token_counts]) @property def completion_llm_token_count(self) -> int: """Get the current total LLM completion token count.""" return sum([x.completion_token_count for x in self.llm_token_counts]) @property def total_embedding_token_count(self) -> int: """Get the current total Embedding token count.""" return sum([x.total_token_count for x in self.embedding_token_counts]) def reset_counts(self) -> None: """Reset the token counts.""" self.llm_token_counts = [] self.embedding_token_counts = []` |
### total\_llm\_token\_count `property` [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.total_llm_token_count "Permanent link")
`total_llm_token_count: int`
Get the current total LLM token count.
### prompt\_llm\_token\_count `property` [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.prompt_llm_token_count "Permanent link")
`prompt_llm_token_count: int`
Get the current total LLM prompt token count.
### completion\_llm\_token\_count `property` [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.completion_llm_token_count "Permanent link")
`completion_llm_token_count: int`
Get the current total LLM completion token count.
### total\_embedding\_token\_count `property` [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.total_embedding_token_count "Permanent link")
`total_embedding_token_count: int`
Get the current total Embedding token count.
### on\_event\_end [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.on_event_end "Permanent link")
`on_event_end(event_type: [CBEventType](../#llama_index.core.callbacks.schema.CBEventType "llama_index.core.callbacks.schema.CBEventType") , payload: Optional[Dict[str, Any]] = None, event_id: str = '', **kwargs: Any) -> None`
Count the LLM or Embedding tokens as needed.
Source code in `llama-index-core/llama_index/core/callbacks/token_counting.py`
| | |
| --- | --- |
| 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 | `def on_event_end( self, event_type: CBEventType, payload: Optional[Dict[str, Any]] = None, event_id: str = "", **kwargs: Any, ) -> None: """Count the LLM or Embedding tokens as needed.""" if ( event_type == CBEventType.LLM and event_type not in self.event_ends_to_ignore and payload is not None ): self.llm_token_counts.append( get_llm_token_counts( token_counter=self._token_counter, payload=payload, event_id=event_id, ) ) if self._verbose: self._print( "LLM Prompt Token Usage: " f"{self.llm_token_counts[-1].prompt_token_count}\n" "LLM Completion Token Usage: " f"{self.llm_token_counts[-1].completion_token_count}", ) elif ( event_type == CBEventType.EMBEDDING and event_type not in self.event_ends_to_ignore and payload is not None ): total_chunk_tokens = 0 for chunk in payload.get(EventPayload.CHUNKS, []): self.embedding_token_counts.append( TokenCountingEvent( event_id=event_id, prompt=chunk, prompt_token_count=self._token_counter.get_string_tokens(chunk), completion="", completion_token_count=0, ) ) total_chunk_tokens += self.embedding_token_counts[-1].total_token_count if self._verbose: self._print(f"Embedding Token Usage: {total_chunk_tokens}")` |
### reset\_counts [#](#llama_index.core.callbacks.token_counting.TokenCountingHandler.reset_counts "Permanent link")
`reset_counts() -> None`
Reset the token counts.
Source code in `llama-index-core/llama_index/core/callbacks/token_counting.py`
| | |
| --- | --- |
| 260
261
262
263 | `def reset_counts(self) -> None: """Reset the token counts.""" self.llm_token_counts = [] self.embedding_token_counts = []` |
Back to top

---
# Code - LlamaIndex
🚀 Contributing to LlamaIndex[#](#contributing-to-llamaindex "Permanent link")
===============================================================================
Welcome to **LlamaIndex**! We’re excited that you want to contribute and become part of our growing community. Whether you're interested in building integrations, fixing bugs, or adding exciting new features, we've made it easy for you to get started.
* * *
🎯 Quick Start Guide[#](#quick-start-guide "Permanent link")
-------------------------------------------------------------
If you're ready to dive in, here’s a quick setup guide to get you going:
1. **Fork** the repo and clone your fork.
2. Navigate to the project folder:
`cd llama_index`
3. Set up a new virtual environment with `Poetry`:
`poetry shell`
4. Install development (and/or docs) dependencies:
`poetry install --only dev,docs --no-root`
5. Install the package(s) you want to work on. You will for sure need to install `llama-index-core`:
`pip install -e llama-index-core`
From there, you can install specific integrations that you want to work on:
`pip install -e llama-index-integrations/llms/llama-index-llms-openai`
**That’s it!** If anything seems unclear, scroll down to the [Development Guidelines](#-Development-Guidelines)
for more details.
* * *
🛠️ What Can You Work On?[#](#what-can-you-work-on "Permanent link")
---------------------------------------------------------------------
There’s plenty of ways to contribute—whether you’re a seasoned Python developer or just starting out, your contributions are welcome! Here are some ideas:
### 1\. 🆕 Extend Core Modules[#](#1-extend-core-modules "Permanent link")
Help us extend LlamaIndex's functionality by contributing to any of our core modules. Think of this as unlocking new superpowers for LlamaIndex!
* **New Integrations** (e.g., connecting new LLMs, storage systems, or data sources)
* **Data Loaders**, **Vector Stores**, and more!
Explore the different modules below to get inspired!
### 2\. 📦 Contribute Tools, Readers, Packs, or Datasets[#](#2-contribute-tools-readers-packs-or-datasets "Permanent link")
Create new Packs, Readers, or Tools that simplify how others use LlamaIndex with various platforms.
### 3\. 🧠 Add New Features[#](#3-add-new-features "Permanent link")
Have an idea for a feature that could make LlamaIndex even better? Go for it! We love innovative contributions.
### 4\. 🐛 Fix Bugs[#](#4-fix-bugs "Permanent link")
Fixing bugs is a great way to start contributing. Head over to our [Github Issues](https://github.com/run-llama/llama_index/issues)
page and find bugs tagged as [`good first issue`](https://github.com/run-llama/llama_index/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)
.
### 5\. 🎉 Share Usage Examples[#](#5-share-usage-examples "Permanent link")
If you’ve used LlamaIndex in a unique or creative way, consider sharing guides or notebooks. This helps other developers learn from your experience.
### 6\. 🧪 Experiment[#](#6-experiment "Permanent link")
Got an out-there idea? We’re open to experimental features—test it out and make a PR!
### 7\. 📄 Improve Documentation & Code Quality[#](#7-improve-documentation-code-quality "Permanent link")
Help make the project easier to navigate by refining the docs or cleaning up the codebase. Every improvement counts!
* * *
🔥 How to Extend LlamaIndex’s Core Modules[#](#how-to-extend-llamaindexs-core-modules "Permanent link")
--------------------------------------------------------------------------------------------------------
### Data Loaders[#](#data-loaders "Permanent link")
A **data loader** ingests data from any source and converts it into `Document` objects that LlamaIndex can parse and index.
* **Interface**:
* `load_data`: Returns a list of `Document` objects.
* `lazy_load_data`: Returns an iterable of `Document` objects (useful for large datasets).
**Example**: [MongoDB Reader](https://github.com/run-llama/llama_index/tree/main/llama-index-integrations/readers/llama-index-readers-mongodb)
💡 **Ideas**: Want to load data from a source not yet supported? Build a new data loader and submit a PR!
### Node Parsers[#](#node-parsers "Permanent link")
A **node parser** converts `Document` objects into `Node` objects—atomic chunks of data that LlamaIndex works with.
* **Interface**:
* `get_nodes_from_documents`: Returns a list of `Node` objects.
**Example**: [Hierarchical Node Parser](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/node_parser/relational/hierarchical.py)
💡 **Ideas**: Add new ways to structure hierarchical relationships in documents, like play-act-scene or chapter-section formats.
### Text Splitters[#](#text-splitters "Permanent link")
A **text splitter** breaks down large text blocks into smaller chunks—this is key for working with LLMs that have limited context windows.
* **Interface**:
* `split_text`: Takes a string and returns smaller strings (chunks).
**Example**: [Token Text Splitter](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/node_parser/text/token.py)
💡 **Ideas**: Build specialized text splitters for different content types, like code, dialogues, or dense data!
### Vector Stores[#](#vector-stores "Permanent link")
Store embeddings and retrieve them via similarity search with **vector stores**.
* **Interface**:
* `add`, `delete`, `query`, `get_nodes`, `delete_nodes`, `clear`
**Example**: [Pinecone Vector Store](https://github.com/run-llama/llama_index/tree/main/llama-index-integrations/vector_stores/llama-index-vector-stores-pinecone)
💡 **Ideas**: Create support for vector databases that aren't yet integrated!
### Query Engines & Retrievers[#](#query-engines-retrievers "Permanent link")
* **Query Engines** implement `query` to return structured responses.
* **Retrievers** retrieve relevant nodes based on queries.
💡 **Ideas**: Design fancy query engines that combine retrievers or add intelligent processing layers!
* * *
✨ Steps to Contribute[#](#steps-to-contribute "Permanent link")
----------------------------------------------------------------
1. **Fork** the repository on GitHub.
2. **Clone** your fork to your local machine.
`git clone https://github.com/your-username/llama_index.git`
3. **Create a branch** for your work.
`git checkout -b your-feature-branch`
4. **Set up your environment** (follow the [Quick Start Guide](#quick-start-guide)
).
5. **Work on your feature or bugfix**, ensuring you have unit tests covering your code.
6. **Commit** your changes, then push them to your fork.
`git push origin your-feature-branch`
7. **Open a pull request** on GitHub.
And voilà—your contribution is ready for review!
* * *
🧑💻 Development Guidelines[#](#development-guidelines "Permanent link")
--------------------------------------------------------------------------
### Repo Structure[#](#repo-structure "Permanent link")
LlamaIndex is organized as a **monorepo**, meaning different packages live within this single repository. You can focus on a specific package depending on your contribution:
* **Core package**: [`llama-index-core/`](https://github.com/run-llama/llama_index/tree/main/llama-index-core)
* **Integrations**: e.g., [`llama-index-integrations/`](https://github.com/run-llama/llama_index/tree/main/llama-index-integrations)
### Setting Up Your Environment[#](#setting-up-your-environment "Permanent link")
1. **Install Poetry** (if you don’t already have it):
`curl -sSL https://install.python-poetry.org | python3 -`
2. Activate the environment:
`poetry shell`
3. Install dependencies:
`poetry install --only dev,docs --no-root`
4. Install the package(s) you want to work on. You will for sure need to install `llama-index-core`:
`pip install -e llama-index-core`
From there, you can install specific integrations that you want to work on:
`pip install -e llama-index-integrations/llms/llama-index-llms-openai`
### Running Tests[#](#running-tests "Permanent link")
We use `pytest` for testing. Make sure you run tests in each package you modify:
`pytest`
If you’re integrating with a remote system, **mock** it to prevent test failures from external changes.
By default, CICD will fail if test coverage is less than 50% -- so please do add tests for your code!
* * *
👥 Join the Community[#](#join-the-community "Permanent link")
---------------------------------------------------------------
We’d love to hear from you and collaborate! Join our Discord community to ask questions, share ideas, or just chat with fellow developers.
Join us on Discord [https://discord.gg/dGcwcsnxhU](https://discord.gg/dGcwcsnxhU)
* * *
🌟 Acknowledgements[#](#acknowledgements "Permanent link")
-----------------------------------------------------------
Thank you for considering contributing to LlamaIndex! Every contribution—whether it’s code, documentation, or ideas—helps make this project better for everyone.
Happy coding! 😊
Back to top

---
# Openai - LlamaIndex
Openai
======
OpenAIAgent [#](#llama_index.agent.openai.OpenAIAgent "Permanent link")
------------------------------------------------------------------------
Bases: `[AgentRunner](../#llama_index.core.agent.AgentRunner "llama_index.core.agent.runner.base.AgentRunner") `
OpenAI agent.
Subclasses AgentRunner with a OpenAIAgentWorker.
For the legacy implementation see:
`from llama_index.agent.legacy.openai.base import OpenAIAgent`
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py`
| | |
| --- | --- |
| 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 | ``class OpenAIAgent(AgentRunner): """OpenAI agent. Subclasses AgentRunner with a OpenAIAgentWorker. For the legacy implementation see: ```python from llama_index.agent.legacy.openai.base import OpenAIAgent ``` """ def __init__( self, tools: List[BaseTool], llm: OpenAI, memory: BaseMemory, prefix_messages: List[ChatMessage], verbose: bool = False, max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS, default_tool_choice: str = "auto", callback_manager: Optional[CallbackManager] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None, ) -> None: """Init params.""" callback_manager = callback_manager or llm.callback_manager step_engine = OpenAIAgentWorker.from_tools( tools=tools, tool_retriever=tool_retriever, llm=llm, verbose=verbose, max_function_calls=max_function_calls, callback_manager=callback_manager, prefix_messages=prefix_messages, tool_call_parser=tool_call_parser, ) super().__init__( step_engine, memory=memory, llm=llm, callback_manager=callback_manager, default_tool_choice=default_tool_choice, ) @classmethod def from_tools( cls, tools: Optional[List[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, llm: Optional[LLM] = None, chat_history: Optional[List[ChatMessage]] = None, memory: Optional[BaseMemory] = None, memory_cls: Type[BaseMemory] = ChatMemoryBuffer, verbose: bool = False, max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS, default_tool_choice: str = "auto", callback_manager: Optional[CallbackManager] = None, system_prompt: Optional[str] = None, prefix_messages: Optional[List[ChatMessage]] = None, tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None, **kwargs: Any, ) -> "OpenAIAgent": """Create an OpenAIAgent from a list of tools. Similar to `from_defaults` in other classes, this method will infer defaults for a variety of parameters, including the LLM, if they are not specified. """ tools = tools or [] chat_history = chat_history or [] llm = llm or Settings.llm if not isinstance(llm, OpenAI): raise ValueError("llm must be a OpenAI instance") if callback_manager is not None: llm.callback_manager = callback_manager memory = memory or memory_cls.from_defaults(chat_history, llm=llm) if not llm.metadata.is_function_calling_model: raise ValueError( f"Model name {llm.model} does not support function calling API. " ) if system_prompt is not None: if prefix_messages is not None: raise ValueError( "Cannot specify both system_prompt and prefix_messages" ) prefix_messages = [ChatMessage(content=system_prompt, role="system")] prefix_messages = prefix_messages or [] return cls( tools=tools, tool_retriever=tool_retriever, llm=llm, memory=memory, prefix_messages=prefix_messages, verbose=verbose, max_function_calls=max_function_calls, callback_manager=callback_manager, default_tool_choice=default_tool_choice, tool_call_parser=tool_call_parser, )`` |
### from\_tools `classmethod` [#](#llama_index.agent.openai.OpenAIAgent.from_tools "Permanent link")
`from_tools(tools: Optional[List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, tool_retriever: Optional[[ObjectRetriever](../../objects/#llama_index.core.objects.ObjectRetriever "llama_index.core.objects.base.ObjectRetriever") [[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, llm: Optional[[LLM](../../llms/#llama_index.core.llms.llm.LLM "llama_index.core.llms.llm.LLM") ] = None, chat_history: Optional[List[[ChatMessage](../../prompts/#llama_index.core.prompts.ChatMessage "llama_index.core.base.llms.types.ChatMessage") ]] = None, memory: Optional[[BaseMemory](../../memory/#llama_index.core.memory.types.BaseMemory "llama_index.core.memory.types.BaseMemory") ] = None, memory_cls: Type[[BaseMemory](../../memory/#llama_index.core.memory.types.BaseMemory "llama_index.core.memory.types.BaseMemory") ] = [ChatMemoryBuffer](../../memory/chat_memory_buffer/#llama_index.core.memory.chat_memory_buffer.ChatMemoryBuffer "llama_index.core.memory.chat_memory_buffer.ChatMemoryBuffer") , verbose: bool = False, max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS, default_tool_choice: str = 'auto', callback_manager: Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, system_prompt: Optional[str] = None, prefix_messages: Optional[List[[ChatMessage](../../prompts/#llama_index.core.prompts.ChatMessage "llama_index.core.base.llms.types.ChatMessage") ]] = None, tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None, **kwargs: Any) -> [OpenAIAgent](#llama_index.agent.openai.OpenAIAgent "llama_index.agent.openai.base.OpenAIAgent")`
Create an OpenAIAgent from a list of tools.
Similar to `from_defaults` in other classes, this method will infer defaults for a variety of parameters, including the LLM, if they are not specified.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/base.py`
| | |
| --- | --- |
| 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 | ``@classmethod def from_tools( cls, tools: Optional[List[BaseTool]] = None, tool_retriever: Optional[ObjectRetriever[BaseTool]] = None, llm: Optional[LLM] = None, chat_history: Optional[List[ChatMessage]] = None, memory: Optional[BaseMemory] = None, memory_cls: Type[BaseMemory] = ChatMemoryBuffer, verbose: bool = False, max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS, default_tool_choice: str = "auto", callback_manager: Optional[CallbackManager] = None, system_prompt: Optional[str] = None, prefix_messages: Optional[List[ChatMessage]] = None, tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None, **kwargs: Any, ) -> "OpenAIAgent": """Create an OpenAIAgent from a list of tools. Similar to `from_defaults` in other classes, this method will infer defaults for a variety of parameters, including the LLM, if they are not specified. """ tools = tools or [] chat_history = chat_history or [] llm = llm or Settings.llm if not isinstance(llm, OpenAI): raise ValueError("llm must be a OpenAI instance") if callback_manager is not None: llm.callback_manager = callback_manager memory = memory or memory_cls.from_defaults(chat_history, llm=llm) if not llm.metadata.is_function_calling_model: raise ValueError( f"Model name {llm.model} does not support function calling API. " ) if system_prompt is not None: if prefix_messages is not None: raise ValueError( "Cannot specify both system_prompt and prefix_messages" ) prefix_messages = [ChatMessage(content=system_prompt, role="system")] prefix_messages = prefix_messages or [] return cls( tools=tools, tool_retriever=tool_retriever, llm=llm, memory=memory, prefix_messages=prefix_messages, verbose=verbose, max_function_calls=max_function_calls, callback_manager=callback_manager, default_tool_choice=default_tool_choice, tool_call_parser=tool_call_parser, )`` |
OpenAIAssistantAgent [#](#llama_index.agent.openai.OpenAIAssistantAgent "Permanent link")
------------------------------------------------------------------------------------------
Bases: `[BaseAgent](../#llama_index.core.agent.types.BaseAgent "llama_index.core.agent.types.BaseAgent") `
OpenAIAssistant agent.
Wrapper around OpenAI assistant API: https://platform.openai.com/docs/assistants/overview
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 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
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594 | `class OpenAIAssistantAgent(BaseAgent): """OpenAIAssistant agent. Wrapper around OpenAI assistant API: https://platform.openai.com/docs/assistants/overview """ def __init__( self, client: Any, assistant: Any, tools: Optional[List[BaseTool]], callback_manager: Optional[CallbackManager] = None, thread_id: Optional[str] = None, instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, file_dict: Dict[str, str] = {}, verbose: bool = False, ) -> None: """Init params.""" from openai import OpenAI from openai.types.beta.assistant import Assistant self._client = cast(OpenAI, client) self._assistant = cast(Assistant, assistant) self._tools = tools or [] if thread_id is None: thread = self._client.beta.threads.create() thread_id = thread.id self._thread_id = thread_id self._instructions_prefix = instructions_prefix self._run_retrieve_sleep_time = run_retrieve_sleep_time self._verbose = verbose self.file_dict = file_dict self.callback_manager = callback_manager or CallbackManager([]) @classmethod def from_new( cls, name: str, instructions: str, tools: Optional[List[BaseTool]] = None, openai_tools: Optional[List[Dict]] = None, thread_id: Optional[str] = None, model: str = "gpt-4-1106-preview", instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, files: Optional[List[str]] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, file_ids: Optional[List[str]] = None, api_key: Optional[str] = None, ) -> "OpenAIAssistantAgent": """From new assistant. Args: name: name of assistant instructions: instructions for assistant tools: list of tools openai_tools: list of openai tools thread_id: thread id model: model run_retrieve_sleep_time: run retrieve sleep time files: files instructions_prefix: instructions prefix callback_manager: callback manager verbose: verbose file_ids: list of file ids api_key: OpenAI API key """ from openai import OpenAI # this is the set of openai tools # not to be confused with the tools we pass in for function calling openai_tools = openai_tools or [] tools = tools or [] tool_fns = [t.metadata.to_openai_tool() for t in tools] all_openai_tools = openai_tools + tool_fns # initialize client client = OpenAI(api_key=api_key) # process files files = files or [] file_ids = file_ids or [] file_dict = _process_files(client, files) # TODO: openai's typing is a bit sus all_openai_tools = cast(List[Any], all_openai_tools) assistant = client.beta.assistants.create( name=name, instructions=instructions, tools=cast(List[Any], all_openai_tools), model=model, ) return cls( client, assistant, tools, callback_manager=callback_manager, thread_id=thread_id, instructions_prefix=instructions_prefix, file_dict=file_dict, run_retrieve_sleep_time=run_retrieve_sleep_time, verbose=verbose, ) @classmethod def from_existing( cls, assistant_id: str, tools: Optional[List[BaseTool]] = None, thread_id: Optional[str] = None, instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, callback_manager: Optional[CallbackManager] = None, api_key: Optional[str] = None, verbose: bool = False, ) -> "OpenAIAssistantAgent": """From existing assistant id. Args: assistant_id: id of assistant tools: list of BaseTools Assistant can use thread_id: thread id run_retrieve_sleep_time: run retrieve sleep time instructions_prefix: instructions prefix callback_manager: callback manager api_key: OpenAI API key verbose: verbose """ from openai import OpenAI # initialize client client = OpenAI(api_key=api_key) # get assistant assistant = client.beta.assistants.retrieve(assistant_id) # assistant.tools is incompatible with BaseTools so have to pass from params return cls( client, assistant, tools=tools, callback_manager=callback_manager, thread_id=thread_id, instructions_prefix=instructions_prefix, run_retrieve_sleep_time=run_retrieve_sleep_time, verbose=verbose, ) @property def assistant(self) -> Any: """Get assistant.""" return self._assistant @property def client(self) -> Any: """Get client.""" return self._client @property def thread_id(self) -> str: """Get thread id.""" return self._thread_id @property def files_dict(self) -> Dict[str, str]: """Get files dict.""" return self.file_dict @property def chat_history(self) -> List[ChatMessage]: raw_messages = self._client.beta.threads.messages.list( thread_id=self._thread_id, order="asc" ) return from_openai_thread_messages(list(raw_messages)) def reset(self) -> None: """Delete and create a new thread.""" self._client.beta.threads.delete(self._thread_id) thread = self._client.beta.threads.create() thread_id = thread.id self._thread_id = thread_id def get_tools(self, message: str) -> List[BaseTool]: """Get tools.""" return self._tools def upload_files(self, files: List[str]) -> Dict[str, Any]: """Upload files.""" return _process_files(self._client, files) def add_message( self, message: str, file_ids: Optional[List[str]] = None, tools: Optional[List[Dict[str, Any]]] = None, ) -> Any: """Add message to assistant.""" attachments = format_attachments(file_ids=file_ids, tools=tools) return self._client.beta.threads.messages.create( thread_id=self._thread_id, role="user", content=message, attachments=attachments, ) def _run_function_calling(self, run: Any) -> List[ToolOutput]: """Run function calling.""" tool_calls = run.required_action.submit_tool_outputs.tool_calls tool_output_dicts = [] tool_output_objs: List[ToolOutput] = [] for tool_call in tool_calls: fn_obj = tool_call.function _, tool_output = call_function(self._tools, fn_obj, verbose=self._verbose) tool_output_dicts.append( {"tool_call_id": tool_call.id, "output": str(tool_output)} ) tool_output_objs.append(tool_output) # submit tool outputs # TODO: openai's typing is a bit sus self._client.beta.threads.runs.submit_tool_outputs( thread_id=self._thread_id, run_id=run.id, tool_outputs=cast(List[Any], tool_output_dicts), ) return tool_output_objs async def _arun_function_calling(self, run: Any) -> List[ToolOutput]: """Run function calling.""" tool_calls = run.required_action.submit_tool_outputs.tool_calls tool_output_dicts = [] tool_output_objs: List[ToolOutput] = [] for tool_call in tool_calls: fn_obj = tool_call.function _, tool_output = await acall_function( self._tools, fn_obj, verbose=self._verbose ) tool_output_dicts.append( {"tool_call_id": tool_call.id, "output": str(tool_output)} ) tool_output_objs.append(tool_output) # submit tool outputs self._client.beta.threads.runs.submit_tool_outputs( thread_id=self._thread_id, run_id=run.id, tool_outputs=cast(List[Any], tool_output_dicts), ) return tool_output_objs def run_assistant( self, instructions_prefix: Optional[str] = None ) -> Tuple[Any, Dict]: """Run assistant.""" instructions_prefix = instructions_prefix or self._instructions_prefix run = self._client.beta.threads.runs.create( thread_id=self._thread_id, assistant_id=self._assistant.id, instructions=instructions_prefix, ) from openai.types.beta.threads import Run run = cast(Run, run) sources = [] while run.status in ["queued", "in_progress", "requires_action"]: run = self._client.beta.threads.runs.retrieve( thread_id=self._thread_id, run_id=run.id ) if run.status == "requires_action": cur_tool_outputs = self._run_function_calling(run) sources.extend(cur_tool_outputs) time.sleep(self._run_retrieve_sleep_time) if run.status == "failed": raise ValueError( f"Run failed with status {run.status}.\n" f"Error: {run.last_error}" ) return run, {"sources": sources} async def arun_assistant( self, instructions_prefix: Optional[str] = None ) -> Tuple[Any, Dict]: """Run assistant.""" instructions_prefix = instructions_prefix or self._instructions_prefix run = self._client.beta.threads.runs.create( thread_id=self._thread_id, assistant_id=self._assistant.id, instructions=instructions_prefix, ) from openai.types.beta.threads import Run run = cast(Run, run) sources = [] while run.status in ["queued", "in_progress", "requires_action"]: run = self._client.beta.threads.runs.retrieve( thread_id=self._thread_id, run_id=run.id ) if run.status == "requires_action": cur_tool_outputs = await self._arun_function_calling(run) sources.extend(cur_tool_outputs) await asyncio.sleep(self._run_retrieve_sleep_time) if run.status == "failed": raise ValueError( f"Run failed with status {run.status}.\n" f"Error: {run.last_error}" ) return run, {"sources": sources} @property def latest_message(self) -> ChatMessage: """Get latest message.""" raw_messages = self._client.beta.threads.messages.list( thread_id=self._thread_id, order="desc" ) messages = from_openai_thread_messages(list(raw_messages)) return messages[0] def _chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Main chat interface.""" # TODO: since chat interface doesn't expose additional kwargs # we can't pass in file_ids per message _added_message_obj = self.add_message(message) _run, metadata = self.run_assistant( instructions_prefix=self._instructions_prefix, ) latest_message = self.latest_message # get most recent message content return AgentChatResponse( response=str(latest_message.content), sources=metadata["sources"], ) async def _achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", mode: ChatResponseMode = ChatResponseMode.WAIT, ) -> AGENT_CHAT_RESPONSE_TYPE: """Asynchronous main chat interface.""" self.add_message(message) run, metadata = await self.arun_assistant( instructions_prefix=self._instructions_prefix, ) latest_message = self.latest_message # get most recent message content return AgentChatResponse( response=str(latest_message.content), sources=metadata["sources"], ) @trace_method("chat") def chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", ) -> AgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = self._chat( message, chat_history, function_call, mode=ChatResponseMode.WAIT ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @trace_method("chat") async def achat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", ) -> AgentChatResponse: with self.callback_manager.event( CBEventType.AGENT_STEP, payload={EventPayload.MESSAGES: [message]}, ) as e: chat_response = await self._achat( message, chat_history, function_call, mode=ChatResponseMode.WAIT ) assert isinstance(chat_response, AgentChatResponse) e.on_end(payload={EventPayload.RESPONSE: chat_response}) return chat_response @trace_method("chat") def stream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", ) -> StreamingAgentChatResponse: raise NotImplementedError("stream_chat not implemented") @trace_method("chat") async def astream_chat( self, message: str, chat_history: Optional[List[ChatMessage]] = None, function_call: Union[str, dict] = "auto", ) -> StreamingAgentChatResponse: raise NotImplementedError("astream_chat not implemented")` |
### assistant `property` [#](#llama_index.agent.openai.OpenAIAssistantAgent.assistant "Permanent link")
`assistant: Any`
Get assistant.
### client `property` [#](#llama_index.agent.openai.OpenAIAssistantAgent.client "Permanent link")
`client: Any`
Get client.
### thread\_id `property` [#](#llama_index.agent.openai.OpenAIAssistantAgent.thread_id "Permanent link")
`thread_id: str`
Get thread id.
### files\_dict `property` [#](#llama_index.agent.openai.OpenAIAssistantAgent.files_dict "Permanent link")
`files_dict: Dict[str, str]`
Get files dict.
### latest\_message `property` [#](#llama_index.agent.openai.OpenAIAssistantAgent.latest_message "Permanent link")
`latest_message: [ChatMessage](../../prompts/#llama_index.core.prompts.ChatMessage "llama_index.core.base.llms.types.ChatMessage")`
Get latest message.
### from\_new `classmethod` [#](#llama_index.agent.openai.OpenAIAssistantAgent.from_new "Permanent link")
`from_new(name: str, instructions: str, tools: Optional[List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, openai_tools: Optional[List[Dict]] = None, thread_id: Optional[str] = None, model: str = 'gpt-4-1106-preview', instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, files: Optional[List[str]] = None, callback_manager: Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, verbose: bool = False, file_ids: Optional[List[str]] = None, api_key: Optional[str] = None) -> [OpenAIAssistantAgent](#llama_index.agent.openai.OpenAIAssistantAgent "llama_index.agent.openai.openai_assistant_agent.OpenAIAssistantAgent")`
From new assistant.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `name` | `str` | name of assistant | _required_ |
| `instructions` | `str` | instructions for assistant | _required_ |
| `tools` | `Optional[List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]]` | list of tools | `None` |
| `openai_tools` | `Optional[List[Dict]]` | list of openai tools | `None` |
| `thread_id` | `Optional[str]` | thread id | `None` |
| `model` | `str` | model | `'gpt-4-1106-preview'` |
| `run_retrieve_sleep_time` | `float` | run retrieve sleep time | `0.1` |
| `files` | `Optional[List[str]]` | files | `None` |
| `instructions_prefix` | `Optional[str]` | instructions prefix | `None` |
| `callback_manager` | `Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ]` | callback manager | `None` |
| `verbose` | `bool` | verbose | `False` |
| `file_ids` | `Optional[List[str]]` | list of file ids | `None` |
| `api_key` | `Optional[str]` | OpenAI API key | `None` |
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 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
276
277
278
279
280
281
282 | `@classmethod def from_new( cls, name: str, instructions: str, tools: Optional[List[BaseTool]] = None, openai_tools: Optional[List[Dict]] = None, thread_id: Optional[str] = None, model: str = "gpt-4-1106-preview", instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, files: Optional[List[str]] = None, callback_manager: Optional[CallbackManager] = None, verbose: bool = False, file_ids: Optional[List[str]] = None, api_key: Optional[str] = None, ) -> "OpenAIAssistantAgent": """From new assistant. Args: name: name of assistant instructions: instructions for assistant tools: list of tools openai_tools: list of openai tools thread_id: thread id model: model run_retrieve_sleep_time: run retrieve sleep time files: files instructions_prefix: instructions prefix callback_manager: callback manager verbose: verbose file_ids: list of file ids api_key: OpenAI API key """ from openai import OpenAI # this is the set of openai tools # not to be confused with the tools we pass in for function calling openai_tools = openai_tools or [] tools = tools or [] tool_fns = [t.metadata.to_openai_tool() for t in tools] all_openai_tools = openai_tools + tool_fns # initialize client client = OpenAI(api_key=api_key) # process files files = files or [] file_ids = file_ids or [] file_dict = _process_files(client, files) # TODO: openai's typing is a bit sus all_openai_tools = cast(List[Any], all_openai_tools) assistant = client.beta.assistants.create( name=name, instructions=instructions, tools=cast(List[Any], all_openai_tools), model=model, ) return cls( client, assistant, tools, callback_manager=callback_manager, thread_id=thread_id, instructions_prefix=instructions_prefix, file_dict=file_dict, run_retrieve_sleep_time=run_retrieve_sleep_time, verbose=verbose, )` |
### from\_existing `classmethod` [#](#llama_index.agent.openai.OpenAIAssistantAgent.from_existing "Permanent link")
`from_existing(assistant_id: str, tools: Optional[List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]] = None, thread_id: Optional[str] = None, instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, callback_manager: Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ] = None, api_key: Optional[str] = None, verbose: bool = False) -> [OpenAIAssistantAgent](#llama_index.agent.openai.OpenAIAssistantAgent "llama_index.agent.openai.openai_assistant_agent.OpenAIAssistantAgent")`
From existing assistant id.
Parameters:
| Name | Type | Description | Default |
| --- | --- | --- | --- |
| `assistant_id` | `str` | id of assistant | _required_ |
| `tools` | `Optional[List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]]` | list of BaseTools Assistant can use | `None` |
| `thread_id` | `Optional[str]` | thread id | `None` |
| `run_retrieve_sleep_time` | `float` | run retrieve sleep time | `0.1` |
| `instructions_prefix` | `Optional[str]` | instructions prefix | `None` |
| `callback_manager` | `Optional[[CallbackManager](../../callbacks/#llama_index.core.callbacks.base.CallbackManager "llama_index.core.callbacks.CallbackManager") ]` | callback manager | `None` |
| `api_key` | `Optional[str]` | OpenAI API key | `None` |
| `verbose` | `bool` | verbose | `False` |
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327 | `@classmethod def from_existing( cls, assistant_id: str, tools: Optional[List[BaseTool]] = None, thread_id: Optional[str] = None, instructions_prefix: Optional[str] = None, run_retrieve_sleep_time: float = 0.1, callback_manager: Optional[CallbackManager] = None, api_key: Optional[str] = None, verbose: bool = False, ) -> "OpenAIAssistantAgent": """From existing assistant id. Args: assistant_id: id of assistant tools: list of BaseTools Assistant can use thread_id: thread id run_retrieve_sleep_time: run retrieve sleep time instructions_prefix: instructions prefix callback_manager: callback manager api_key: OpenAI API key verbose: verbose """ from openai import OpenAI # initialize client client = OpenAI(api_key=api_key) # get assistant assistant = client.beta.assistants.retrieve(assistant_id) # assistant.tools is incompatible with BaseTools so have to pass from params return cls( client, assistant, tools=tools, callback_manager=callback_manager, thread_id=thread_id, instructions_prefix=instructions_prefix, run_retrieve_sleep_time=run_retrieve_sleep_time, verbose=verbose, )` |
### reset [#](#llama_index.agent.openai.OpenAIAssistantAgent.reset "Permanent link")
`reset() -> None`
Delete and create a new thread.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 356
357
358
359
360
361 | `def reset(self) -> None: """Delete and create a new thread.""" self._client.beta.threads.delete(self._thread_id) thread = self._client.beta.threads.create() thread_id = thread.id self._thread_id = thread_id` |
### get\_tools [#](#llama_index.agent.openai.OpenAIAssistantAgent.get_tools "Permanent link")
`get_tools(message: str) -> List[[BaseTool](../../tools/#llama_index.core.tools.types.BaseTool "llama_index.core.tools.BaseTool") ]`
Get tools.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 363
364
365 | `def get_tools(self, message: str) -> List[BaseTool]: """Get tools.""" return self._tools` |
### upload\_files [#](#llama_index.agent.openai.OpenAIAssistantAgent.upload_files "Permanent link")
`upload_files(files: List[str]) -> Dict[str, Any]`
Upload files.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 367
368
369 | `def upload_files(self, files: List[str]) -> Dict[str, Any]: """Upload files.""" return _process_files(self._client, files)` |
### add\_message [#](#llama_index.agent.openai.OpenAIAssistantAgent.add_message "Permanent link")
`add_message(message: str, file_ids: Optional[List[str]] = None, tools: Optional[List[Dict[str, Any]]] = None) -> Any`
Add message to assistant.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 371
372
373
374
375
376
377
378
379
380
381
382
383
384 | `def add_message( self, message: str, file_ids: Optional[List[str]] = None, tools: Optional[List[Dict[str, Any]]] = None, ) -> Any: """Add message to assistant.""" attachments = format_attachments(file_ids=file_ids, tools=tools) return self._client.beta.threads.messages.create( thread_id=self._thread_id, role="user", content=message, attachments=attachments, )` |
### run\_assistant [#](#llama_index.agent.openai.OpenAIAssistantAgent.run_assistant "Permanent link")
`run_assistant(instructions_prefix: Optional[str] = None) -> Tuple[Any, Dict]`
Run assistant.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460 | `def run_assistant( self, instructions_prefix: Optional[str] = None ) -> Tuple[Any, Dict]: """Run assistant.""" instructions_prefix = instructions_prefix or self._instructions_prefix run = self._client.beta.threads.runs.create( thread_id=self._thread_id, assistant_id=self._assistant.id, instructions=instructions_prefix, ) from openai.types.beta.threads import Run run = cast(Run, run) sources = [] while run.status in ["queued", "in_progress", "requires_action"]: run = self._client.beta.threads.runs.retrieve( thread_id=self._thread_id, run_id=run.id ) if run.status == "requires_action": cur_tool_outputs = self._run_function_calling(run) sources.extend(cur_tool_outputs) time.sleep(self._run_retrieve_sleep_time) if run.status == "failed": raise ValueError( f"Run failed with status {run.status}.\n" f"Error: {run.last_error}" ) return run, {"sources": sources}` |
### arun\_assistant `async` [#](#llama_index.agent.openai.OpenAIAssistantAgent.arun_assistant "Permanent link")
`arun_assistant(instructions_prefix: Optional[str] = None) -> Tuple[Any, Dict]`
Run assistant.
Source code in `llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/openai_assistant_agent.py`
| | |
| --- | --- |
| 462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491 | `async def arun_assistant( self, instructions_prefix: Optional[str] = None ) -> Tuple[Any, Dict]: """Run assistant.""" instructions_prefix = instructions_prefix or self._instructions_prefix run = self._client.beta.threads.runs.create( thread_id=self._thread_id, assistant_id=self._assistant.id, instructions=instructions_prefix, ) from openai.types.beta.threads import Run run = cast(Run, run) sources = [] while run.status in ["queued", "in_progress", "requires_action"]: run = self._client.beta.threads.runs.retrieve( thread_id=self._thread_id, run_id=run.id ) if run.status == "requires_action": cur_tool_outputs = await self._arun_function_calling(run) sources.extend(cur_tool_outputs) await asyncio.sleep(self._run_retrieve_sleep_time) if run.status == "failed": raise ValueError( f"Run failed with status {run.status}.\n" f"Error: {run.last_error}" ) return run, {"sources": sources}` |
Back to top

---