Skip to content
Snippets Groups Projects
Unverified Commit 3f0cbb65 authored by Sasha Dog's avatar Sasha Dog Committed by GitHub
Browse files

Add QueryUnderstandingAgent LlamaPack (#11558)

parent ebb81b09
No related branches found
No related tags found
No related merge requests found
Showing
with 912 additions and 1 deletion
......@@ -906,5 +906,6 @@
"run_jobs": "llama_index.core.async_utils",
"DecomposeQueryTransform": "llama_index.core.query.query_transform.base",
"get_eval_results": "llama_index.core.evaluation.eval_utils",
"REPLICATE_MULTI_MODAL_LLM_MODELS": "llama_index.multi_modal_llms.replicate.base"
"REPLICATE_MULTI_MODAL_LLM_MODELS": "llama_index.multi_modal_llms.replicate.base",
"QueryUnderstandingPack": "llama_index.packs.query_understanding_agent"
}
llama_index/_static
.DS_Store
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
bin/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
etc/
include/
lib/
lib64/
parts/
sdist/
share/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
.ruff_cache
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
notebooks/
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
pyvenv.cfg
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# Jetbrains
.idea
modules/
*.swp
# VsCode
.vscode
# pipenv
Pipfile
Pipfile.lock
# pyright
pyrightconfig.json
poetry_requirements(
name="poetry",
)
python_requirements(
name="reqs",
)
GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
help: ## Show all Makefile targets.
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
format: ## Run code autoformatters (black).
pre-commit install
git ls-files | xargs pre-commit run black --files
lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
test: ## Run tests via pytest.
pytest tests
watch-docs: ## Build and watch documentation.
sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
# LlamaIndex Packs Integration: Query Understanding Agent
This LlamaPack implements Query Understanding Agent
Taking inspiration from Humans - when asked a query, humans would clarify what the query means before proceeding if the human sensed the query is unclear. This LlamaPack implements this.
### Installation
```bash
pip install llama-index
```
## CLI Usage
You can download llamapacks directly using `llamaindex-cli`, which comes installed with the `llama-index` python package:
```bash
llamaindex-cli download-llamapack QueryUnderstandingAgent --download-dir ./query_understanding_agent
```
You can then inspect the files at `./query_understanding_agent` and use them as a template for your own project.
## Code Usage
You can download the pack to a the `./query_understanding_agent` directory:
```python
from llama_index.core.llama_pack import download_llama_pack
# download and install dependencies
QueryUnderstandingAgentPack = download_llama_pack(
"QueryUnderstandingAgent", "./query_understanding_agent"
)
# You can use any llama-hub loader to get documents!
```
From here, you can use the pack, or inspect and modify the pack in `./query_understanding_agent`.
See example notebook for usage.
%% Cell type:code id: tags:
``` python
from dotenv import load_dotenv
load_dotenv()
```
%% Output
True
%% Cell type:markdown id: tags:
# Grab Data
%% Cell type:code id: tags:
``` python
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
```
%% Output
--2024-03-07 21:54:37-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8000::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 75042 (73K) [text/plain]
Saving to: ‘data/paul_graham/paul_graham_essay.txt’
data/paul_graham/pa 100%[===================>] 73.28K --.-KB/s in 0.01s
2024-03-07 21:54:37 (5.91 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]
%% Cell type:code id: tags:
``` python
steve_jobs_text = """1995 was the year when Steve Jobs was approaching the end of his exile from Apple after being kicked out of the company a decade earlier. Him as his new company, NeXT, were brought back in to save Apple from near bankruptcy.Jul 25, 2018"""
```
%% Cell type:markdown id: tags:
# Load Data
%% Cell type:code id: tags:
``` python
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.response.pprint_utils import pprint_response
from llama_index.llms.openai import OpenAI
from llama_index.core.schema import Document
# Tool 1
llm = OpenAI()
data = SimpleDirectoryReader(input_dir="./data/paul_graham/").load_data()
index = VectorStoreIndex.from_documents(data)
# Tool 2
steve_index = VectorStoreIndex.from_documents([Document(text=steve_jobs_text)])
```
%% Cell type:markdown id: tags:
# Agents
%% Cell type:code id: tags:
``` python
from llama_index.core.agent import AgentRunner
from llama_index.packs.query_understanding_agent.step import (
QueryUnderstandingAgentWorker,
HumanInputRequiredException,
)
from llama_index.core.tools import QueryEngineTool
llm = OpenAI(model="gpt-4")
tools = [
QueryEngineTool.from_defaults(
query_engine=index.as_query_engine(),
description="A tool that is useful for retrieving specific snippets from the Paul Graham's life",
),
QueryEngineTool.from_defaults(
query_engine=steve_index.as_query_engine(),
description="A tool that is useful for retrieving specific snippets from the Steve Jobs's life",
),
]
```
%% Cell type:markdown id: tags:
# Baseline
%% Cell type:markdown id: tags:
This does not programmatically pause to detect ambiguity to allow user to provide input. Not only that, sometimes this hallucinate to a random subject.
%% Cell type:code id: tags:
``` python
from llama_index.core.agent import ReActAgentWorker
callback_manager = llm.callback_manager
agent_worker = ReActAgentWorker.from_tools(
tools,
llm=llm,
verbose=True,
callback_manager=callback_manager,
)
agent = AgentRunner(agent_worker, callback_manager=callback_manager)
orig_question = "What did the author do in the summer of 1995?"
response = agent.chat(orig_question)
```
%% Output
Thought: The user is asking about the author but hasn't specified which author. I have tools to query information about Paul Graham and Steve Jobs. I need to ask the user to specify the author.
Action: None
Answer: Could you please specify which author you are referring to? I have information about Paul Graham and Steve Jobs.

%% Cell type:markdown id: tags:
# Query Understanding Agent Worker
%% Cell type:markdown id: tags:
This allows the agent to ask user for clarification if the user query is unclear
%% Cell type:code id: tags:
``` python
# from llama_index.core.agent import ReActAgentWorker
callback_manager = llm.callback_manager
agent_worker = QueryUnderstandingAgentWorker.from_tools(
tools,
llm=llm,
callback_manager=callback_manager,
)
agent = AgentRunner(agent_worker, callback_manager=callback_manager)
```
%% Cell type:code id: tags:
``` python
# from llama_index.llms.openai import OpenAI
orig_question = "what did he do in the summer of 1995?"
llm = OpenAI(model="gpt-4")
clarifying_questions = []
try:
response = agent.chat(orig_question)
except HumanInputRequiredException as e:
response = input(e.message)
clarifying_questions.append((e.message, response))
should_end = False
while not should_end:
clarifying_texts = "\n".join(
[
f"""
Q: {question}
A: {answer}
"""
for question, answer in clarifying_questions
]
)
query_text = f"""
Given a query and a set of clarifying questions, please rewrite the query to be more clear.
Example:
Q: What trajectory is the monthly earning from the three months: April, May and June?
Clarifying Questions:
Q: What year are you referring to?
A: In 2022
Q: What company are you referring to?
A: Uber
Rewrite: What was the trajectory of Uber's monthly earnings for the months of April, May, and June in 2022?
Q:{orig_question}
Clarifying Questions: {clarifying_texts}
Rewrite: """
rewrite_response = llm.complete(query_text)
orig_question = rewrite_response
try:
output = agent.chat(rewrite_response.text)
should_end = True
print(f"response: {output.response}")
except HumanInputRequiredException as er:
response = input(er.message)
clarifying_questions.append((er.message, response))
```
%% Output
response: In the summer of 1995, Steve Jobs was involved in the process of returning to Apple after his departure from the company a decade earlier.
from llama_index.packs.query_understanding_agent.base import QueryUnderstandingAgentPack
__all__ = ["QueryUnderstandingAgentPack"]
"""Query Understanding agent pack."""
from typing import Any, Dict, List, Optional
from llama_index.core.agent import AgentRunner
from llama_index.core.callbacks import CallbackManager
from llama_index.core.llama_pack.base import BaseLlamaPack
from llama_index.core.llms.llm import LLM
from llama_index.core.tools.types import BaseTool
from llama_index.llms.openai import OpenAI
from .step import QueryUnderstandingAgentWorker
class QueryUnderstandingAgentPack(BaseLlamaPack):
"""LLMCompilerAgent pack.
Args:
tools (List[BaseTool]): List of tools to use.
llm (Optional[LLM]): LLM to use.
"""
def __init__(
self,
tools: List[BaseTool],
llm: Optional[LLM] = None,
callback_manager: Optional[CallbackManager] = None,
agent_worker_kwargs: Optional[Dict[str, Any]] = None,
agent_runner_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
"""Init params."""
self.llm = llm or OpenAI(model="gpt-4")
self.callback_manager = callback_manager or self.llm.callback_manager
self.agent_worker = QueryUnderstandingAgentWorker.from_tools(
tools,
llm=llm,
verbose=True,
callback_manager=self.callback_manager,
**(agent_worker_kwargs or {})
)
self.agent = AgentRunner(
self.agent_worker,
callback_manager=self.callback_manager,
**(agent_runner_kwargs or {})
)
def get_modules(self) -> Dict[str, Any]:
"""Get modules."""
return {
"llm": self.llm,
"callback_manager": self.callback_manager,
"agent_worker": self.agent_worker,
"agent": self.agent,
}
def run(self, *args: Any, **kwargs: Any) -> Any:
"""Run the pipeline."""
return self.agent.chat(*args, **kwargs)
from llama_index.core.bridge.pydantic import Field, BaseModel, PrivateAttr
from llama_index.core import PromptTemplate
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field
import uuid
from llama_index.core.agent.types import (
Task,
TaskStep,
TaskStepOutput,
)
from llama_index.core.agent.custom.simple import CustomSimpleAgentWorker
from llama_index.core.callbacks import (
trace_method,
)
from llama_index.core.chat_engine.types import AgentChatResponse
from llama_index.core.agent import CustomSimpleAgentWorker, Task
from typing import Dict, Any, List, Tuple, Optional
from llama_index.core.tools import BaseTool, QueryEngineTool
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.prompts import ChatPromptTemplate
from llama_index.core.llms import ChatMessage, MessageRole
DEFAULT_PROMPT_STR = """
Given the question, the tools provided and response, please determine if the question is clear enough to provide the response (with the help of the tools).
If the question is not clear enough, provide a clarifying_question for user to clarify. Provide a clarifying question that only need plan text answers. DO NOT ASK the user a similar question to the one the user asks you.
If the response is sufficient, then return has_error: false, and requires_human_input: false. No need to make the response perfect, as long as it is sufficient.
Given the questions and tools, here are several ways to clarify questions:
- If there are multiple tools for each timeframe, then if the timeframe is not specified in the question, we need to ask to be clear on the timeframe.
- Be clear on which subjects the user is referring to, ask which subject the user is referring to if there are a tool for each subject
Example 1:
Tools:
A useful tool for financial documents for uber in 2022
A useful tool for financial documents for lyft in 2022
Question: What is the company's financial stats in 2022?
Response: Uber's financial stats is 20k in the year 2022
Answer:
{{"requires_human_input": true, "has_error": true, "clarifying_question": "Which company are you referring to?", "explanation": "Given the tools and the question, it is not clear which company the user is referring to"}}
Example 2:
Tools:
A useful tool for financial documents for uber in 2022
A useful tool for financial documents for lyft in 2022
Question: What is the Uber's financial stats in 2022?
Response: Uber's financial stats is 20k in the year 2022
Answer:
{{"requires_human_input": false, "has_error": false, "clarifying_question": "", "explanation": "It is quite clear that the user is referring to uber and the year 2022"}}
Tools:
{tools}
Question: {query_str}
Response: {response_str}
Please return the evaluation of the response in the following JSON format.
Answer:
"""
@dataclass
class AgentChatComplexResponse(AgentChatResponse):
"""Agent chat response with metadata."""
output: Dict[str, Any] = field(default_factory=dict)
def get_chat_prompt_template(
system_prompt: str, current_reasoning: Tuple[str, str]
) -> ChatPromptTemplate:
system_msg = ChatMessage(role=MessageRole.SYSTEM, content=system_prompt)
messages = [system_msg]
for raw_msg in current_reasoning:
if raw_msg[0] == "user":
messages.append(ChatMessage(role=MessageRole.USER, content=raw_msg[1]))
else:
messages.append(ChatMessage(role=MessageRole.ASSISTANT, content=raw_msg[1]))
return ChatPromptTemplate(message_templates=messages)
class ResponseEval(BaseModel):
"""Evaluation of whether the response has an error."""
clarifying_question: str = Field(
..., description="The clarifying question, if human input is required"
)
explanation: str = Field(
...,
description=(
"The explanation for the error OR for the clarifying question."
"Can include the direct stack trace as well."
),
)
has_error: bool = Field(..., description="Whether the response has an error")
requires_human_input: bool = Field(
...,
description="Whether the response needs human input. If true, the clarifying question should be provided.",
)
class HumanInputRequiredException(Exception):
"""Exception raised when human input is required."""
def __init__(
self,
message="Human input is required",
task_id: Optional[str] = None,
step: TaskStep = None,
):
self.message = message
self.task_id = task_id
self.step = step
super().__init__(self.message)
class QueryUnderstandingAgentWorker(CustomSimpleAgentWorker):
"""Agent worker that adds a retry layer on top of a router.
Continues iterating until there's no errors / task is done.
"""
prompt_str: str = Field(default=DEFAULT_PROMPT_STR)
max_iterations: int = Field(default=10)
_router_query_engine: RouterQueryEngine = PrivateAttr()
def __init__(self, tools: List[BaseTool], **kwargs: Any) -> None:
"""Init params."""
# validate that all tools are query engine tools
for tool in tools:
if not isinstance(tool, QueryEngineTool):
raise ValueError(
f"Tool {tool.metadata.name} is not a query engine tool."
)
self._router_query_engine = RouterQueryEngine.from_defaults(
llm=kwargs.get("llm"),
select_multi=False,
query_engine_tools=tools,
verbose=kwargs.get("verbose", False),
)
super().__init__(
tools=tools,
**kwargs,
)
def _initialize_state(self, task: Task, **kwargs: Any) -> Dict[str, Any]:
"""Initialize state."""
return {"count": 0, "current_reasoning": []}
def _run_llm_program(self, query_str, response_str, tools):
for _ in range(3):
try:
return self.llm.structured_predict(
ResponseEval,
PromptTemplate(self.prompt_str),
query_str=query_str,
response_str=str(response_str),
tools=tools,
)
except Exception as e:
print(f"Attempt failed with error: {e}")
continue
raise Exception("Failed to run LLM program after 3 attempts")
def _run_step(
self, state: Dict[str, Any], task: Task, input: Optional[str] = None
) -> Tuple[AgentChatComplexResponse, bool]:
"""Run step.
Returns:
Tuple of (agent_response, is_done)
"""
if input is not None:
# if input is specified, override input
new_input = input
elif "new_input" not in state:
new_input = task.input
else:
new_input = state["new_input"]["text"]
if self.verbose:
print(f"> Current Input: {new_input}")
# first run router query engine
response = self._router_query_engine.query(new_input)
# append to current reasoning
state["current_reasoning"].extend(
[("user", new_input), ("assistant", str(response))]
)
# Then, check for errors
# dynamically create pydantic program for structured output extraction based on template
tools = "\n".join([a.description for a in self._router_query_engine._metadatas])
response_eval = self._run_llm_program(
query_str=new_input, response_str=str(response), tools=tools
)
if self.verbose:
print(f"> Question: {new_input}")
print(f"> Response: {response}")
print(f"> Response eval: {response_eval.dict()}")
# return response
if response_eval.requires_human_input:
return (
AgentChatComplexResponse(
response=response,
output={
"type": "requires_human_input",
"clarifying_question": str(response_eval.clarifying_question),
"has_error": response_eval.has_error,
"explanation": response_eval.explanation,
},
),
True,
)
return AgentChatComplexResponse(response=response), not response_eval.has_error
@trace_method("run_step")
def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
"""Run step."""
output, is_done = self._run_step(step.step_state, task, input=step.input)
if output.output and output.output["type"] == "requires_human_input":
raise HumanInputRequiredException(
message=output.output["clarifying_question"],
task_id=task.task_id,
step=step,
)
response = self._get_task_step_response(output, step, is_done)
# sync step state with task state
task.extra_state.update(step.step_state)
return response
def _get_task_step_response(
self,
output: Dict,
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=output,
task_step=step,
is_last=is_done,
next_steps=new_steps,
)
def _finalize_task(self, state: Dict[str, Any], **kwargs) -> None:
"""Finalize task."""
# nothing to finalize here
# this is usually if you want to modify any sort of
# internal state beyond what is set in `_initialize_state`
[build-system]
build-backend = "poetry.core.masonry.api"
requires = ["poetry-core"]
[tool.codespell]
check-filenames = true
check-hidden = true
# Feel free to un-skip examples, and experimental, you will just need to
# work through many typos (--write-changes and --interactive will help)
skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
[tool.llamahub]
contains_example = false
import_path = "llama_index.packs.query_understanding_agent"
[tool.llamahub.class_authors]
QueryUnderstandingPack = "llama-index"
[tool.mypy]
disallow_untyped_defs = true
# Remove venv skip when integrated with pre-commit
exclude = ["_static", "build", "examples", "notebooks", "venv"]
ignore_missing_imports = true
python_version = "3.8"
[tool.poetry]
authors = ["Sasha Sheng <hackgoofer@gmail.com>"]
description = "llama-index packs query understanding agent integration"
exclude = ["**/BUILD"]
license = "MIT"
name = "llama-index-packs-query-understanding-agent"
packages = [{include = "llama_index/"}]
readme = "README.md"
version = "0.1.0"
[tool.poetry.dependencies]
python = ">=3.8.1,<4.0"
llama-index-core = "^0.10.0"
llama-index-llms-openai = "^0.1.7"
[tool.poetry.group.dev.dependencies]
black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
codespell = {extras = ["toml"], version = ">=v2.2.6"}
ipython = "8.10.0"
jupyter = "^1.0.0"
mypy = "0.991"
pre-commit = "3.2.0"
pylint = "2.15.10"
pytest = "7.2.1"
pytest-mock = "3.11.1"
ruff = "0.0.292"
tree-sitter-languages = "^1.8.0"
types-Deprecated = ">=0.1.0"
types-PyYAML = "^6.0.12.12"
types-protobuf = "^4.24.0.4"
types-redis = "4.5.5.0"
types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
types-setuptools = "67.1.0.0"
python_tests()
from llama_index.core.llama_pack import BaseLlamaPack
from llama_index.packs.query_understanding_agent import QueryUnderstandingAgentPack
def test_class():
names_of_base_classes = [b.__name__ for b in QueryUnderstandingAgentPack.__mro__]
assert BaseLlamaPack.__name__ in names_of_base_classes
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment