├── .gitignore ├── LICENSE ├── README-zh.md ├── README.md ├── assets ├── ExperimentalResults.png ├── LevelRAG-ch.png └── LevelRAG.png ├── scripts ├── build_dense_wiki.sh ├── build_elastic_wiki.sh ├── run_highlevel.sh ├── run_highlevel_gui.sh ├── run_simple.sh └── run_simple_gui.sh └── searchers ├── __init__.py ├── dense_searcher.py ├── high_level_searcher.py ├── hybrid_searcher.py ├── keyword_searcher.py ├── prompts ├── bm25_refine_emphasize_prompt.json ├── bm25_refine_extend_prompt.json ├── bm25_refine_filter_prompt.json ├── bm25_rewrite_prompt.json ├── decompose_with_context_prompt.json ├── decompose_without_context_prompt.json ├── dense_rewrite_prompt.json ├── lucene_rewrite_prompt.json ├── rewrite_by_answer_with_context_prompt.json ├── rewrite_by_answer_without_context_prompt.json ├── summarize_by_answer_prompt.json ├── verify_prompt.json └── web_rewrite_prompt.json ├── searcher.py └── web_searcher.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .idea 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | applications/DeepSpeed-Chat/data 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | pip-wheel-metadata/ 26 | share/python-wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .nox/ 46 | .coverage 47 | .coverage.* 48 | .cache 49 | nosetests.xml 50 | coverage.xml 51 | *.cover 52 | *.py,cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 63 | db.sqlite3 64 | db.sqlite3-journal 65 | 66 | # Flask stuff: 67 | instance/ 68 | .webassets-cache 69 | 70 | # Scrapy stuff: 71 | .scrapy 72 | 73 | # Sphinx documentation 74 | docs/_build/ 75 | 76 | # PyBuilder 77 | target/ 78 | 79 | # Jupyter Notebook 80 | .ipynb_checkpoints 81 | 82 | # IPython 83 | profile_default/ 84 | ipython_config.py 85 | 86 | # pyenv 87 | .python-version 88 | 89 | # pipenv 90 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 91 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 92 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 93 | # install all needed dependencies. 94 | #Pipfile.lock 95 | 96 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 97 | __pypackages__/ 98 | 99 | # Celery stuff 100 | celerybeat-schedule 101 | celerybeat.pid 102 | 103 | # SageMath parsed files 104 | *.sage.py 105 | 106 | # Environments 107 | .env 108 | .venv 109 | env/ 110 | venv/ 111 | ENV/ 112 | env.bak/ 113 | venv.bak/ 114 | 115 | # Spyder project settings 116 | .spyderproject 117 | .spyproject 118 | 119 | # Rope project settings 120 | .ropeproject 121 | 122 | # mkdocs documentation 123 | /site 124 | 125 | # mypy 126 | .mypy_cache/ 127 | .dmypy.json 128 | dmypy.json 129 | 130 | # Pyre type checker 131 | .pyre/ 132 | 133 | 134 | # vscode 135 | .vscode 136 | 137 | 138 | # third party models 139 | yala/model/third_party_models 140 | 141 | # aim 142 | .aim 143 | 144 | # test files 145 | _test*.py 146 | _test*.ipynb 147 | 148 | # experimental configs 149 | experimental_configs/ 150 | 151 | # hydra logs 152 | outputs/ 153 | 154 | # pytest configs 155 | tests/configs/ 156 | 157 | # cibuildwheel 158 | wheelhouse/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 zhangzhuocheng 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 6 | 7 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 8 | 9 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /README-zh.md: -------------------------------------------------------------------------------- 1 | # LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers 2 | 3 | ![Language](https://img.shields.io/badge/language-python-brightgreen) 4 | [![Code Style](https://img.shields.io/badge/code%20style-black-black)](https://github.com/psf/black) 5 | [![Imports: isort](https://img.shields.io/badge/imports-isort-blue)](https://pycqa.github.io/isort/) 6 | [![github license](https://img.shields.io/github/license/ictnlp/LevelRAG)](LICENSE) 7 | [![arXiv](https://img.shields.io/badge/arXiv-2502.18139-b31b1b)](https://arxiv.org/abs/2502.18139) 8 | 9 | 本项目为论文 [**LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers**](https://arxiv.org/abs/2502.18139) 的源码。 10 | 11 | ## 概览 12 | LevelRAG 是一种两阶段的检索增强生成（RAG）框架，结合了多跳逻辑规划和混合检索，以提高检索过程的完整性和准确性。其中第一阶段采用一个高级搜索器，将用户查询分解为原子查询。第二阶段利用多个低级搜索器，为每个子查询检索最相关的文档，然后将相关信息汇总到高级搜索器中生成最终答案。在每个低级搜索器中，采用大型语言模型（LLMs）对原子查询进行适应性优化，以更好地适应低级搜索器中内置的检索器。 13 | 14 |
15 | 16 |
17 | 18 | 19 | ## 运行 LevelRAG 20 | 21 | ### 环境准备 22 | 本项目是基于 [FlexRAG](https://github.com/ictnlp/FlexRAG) 实现的。请参照以下命令安装 FlexRAG： 23 | 24 | ```bash 25 | pip install flexrag==0.2.0 26 | ``` 27 | 28 | 下载本项目源码： 29 | 30 | ```bash 31 | git clone https://github.com/ictnlp/LevelRAG 32 | ``` 33 | 34 | ### 准备检索器 35 | 36 | > [!TIP] 37 | > 如果您希望以更简单的方式运行LevelRAG，请查看[运行 Simple LevelRAG](#运行-simple-levelrag)一节。 38 | 39 | 在运行 LevelRAG 前，需要构建检索器。LevelRAG 使用了三种不同的检索器，分别是 `DenseRetriever`、`ElasticRetriever`和`WebRetriever` 。除了 `WebRetriever` 不需要构建索引外，`DenseRetriever` 和 `ElasticRetriever` 都需要先构建索引。在我们的实验中，我们使用了 [Atlas](https://github.com/facebookresearch/atlas) 提供的维基百科语料库。您可以通过以下命令下载语料库： 40 | 41 | 42 | ```bash 43 | wget https://dl.fbaipublicfiles.com/atlas/corpora/wiki/enwiki-dec2021/text-list-100-sec.jsonl 44 | wget https://dl.fbaipublicfiles.com/atlas/corpora/wiki/enwiki-dec2021/infobox.jsonl 45 | ``` 46 | 47 | 下载完语料库后，您可以运行以下命令构建 `DenseRetriever`： 48 | 49 | ```bash 50 | DENSE_PATH=wikipedia 51 | 52 | python -m flexrag.entrypoints.prepare_index \ 53 | retriever_type=dense \ 54 | corpus_path=[text-list-100-sec.jsonl,infobox.jsonl] \ 55 | saving_fields=[text] \ 56 | dense_config.database_path=$DENSE_PATH \ 57 | dense_config.passage_encoder_config.encoder_type=hf \ 58 | dense_config.passage_encoder_config.hf_config.model_path=facebook/contriever-msmarco \ 59 | dense_config.passage_encoder_config.hf_config.device_id=[0] \ 60 | dense_config.encode_fields=[text] \ 61 | dense_config.index_type=faiss \ 62 | dense_config.batch_size=1024 \ 63 | dense_config.log_interval=100000 64 | ``` 65 | 66 | 类似的，您可以运行以下命令构建 `ElasticRetriever`： 67 | 68 | ```bash 69 | python -m flexrag.entrypoints.prepare_index \ 70 | retriever_type=elastic \ 71 | corpus_path=[text-list-100-sec.jsonl,infobox.jsonl] \ 72 | saving_fields=[text] \ 73 | elastic_config.host='http://127.0.0.1:9200/' \ 74 | elastic_config.index_name=wikipedia \ 75 | elastic_config.batch_size=512 \ 76 | elastic_config.log_interval=100 \ 77 | reinit=True 78 | ``` 79 | 80 | > **Notice:** 81 | > 在构建 `ElasticRetriever` 前，您需要安装 elasticsearch 。您可以参考[这里](https://www.elastic.co/guide/en/elasticsearch/reference/current/install-elasticsearch.html)的指令来安装 elasticsearch。 82 | 83 | `WebRetriever` 不需要构建索引，但需要您事先准备 Bing Search API 的密钥。您可以访问 [Bing Search API](https://www.microsoft.com/en-us/bing/apis) 来获取密钥。 84 | 85 | 86 | ### 准备生成器 87 | LevelRAG 使用 `Qwen2-7B-Instruct` 作为生成器，您可以通过以下命令使用 `vllm` 来部署生成器： 88 | 89 | ```bash 90 | python -m vllm.entrypoints.openai.api_server \ 91 | --model Qwen2-7B-Instruct \ 92 | --gpu-memory-utilization 0.95 \ 93 | --tensor-parallel-size 4 \ 94 | --port 8000 \ 95 | --host 0.0.0.0 \ 96 | --trust-remote-code 97 | ``` 98 | 99 | 该命令将使用 4 个 GPU 来部署 `Qwen2-7B-Instruct` 模型，您可以根据您的 GPU 数量和显存大小来调整 `--tensor-parallel-size` 和 `--gpu-memory-utilization` 参数。 100 | 101 | 102 | ### 启动 LevelRAG 评估 103 | 在完成检索器的准备后，您可以通过运行 `scripts` 文件夹中的 `run_highlevel.sh` 脚本来运行 LevelRAG 评估脚本。注意替换脚本中的变量： 104 | - LEVELRAG_PATH：您在下载本项目源码时的路径 105 | - DENSE_PATH: 您在构建 `DenseRetriever` 时保存的路径 106 | - BING_KEY：您在准备 `WebRetriever` 时获取的 Bing Search API 密钥 107 | 108 | 109 | ### 启动 LevelRAG 图形界面 110 | 您也可以通过运行 `scripts` 文件夹中的 `run_highlevel_gui.sh` 脚本来启动 LevelRAG 的图形界面，在图形界面中您可以输入查询并查看 LevelRAG 的输出。 111 | 112 | ## 运行 Simple LevelRAG 113 | 如果您不希望使用 `WebRetriever` 及 `ElasticRetriever`。您可以仅使用 `DenseRetriever` 来运行 LevelRAG。得益于良好的多跳问题分解及子查询适应性优化， LevelRAG 在仅使用单一检索器时也能取得不错的效果，且运行速度更快。您可以通过运行 `scripts` 文件夹中的 `run_simple.sh` 脚本来运行 Simple LevelRAG 的评估脚本，或运行 `run_simple_gui.sh` 脚本来启动 Simple LevelRAG 的图形界面。 114 | 115 | 该脚本使用了 FlexRAG 项目提供的 `DenseRetriever` 检索器，因此您**无需构建索引**，直接运行脚本即可。 116 | 117 | > [!NOTE] 118 | > 请确认您已经将 `run_simple.sh` 或 `run_simple_gui.sh` 脚本中的 `API_KEY` 替换为您的 OpenAI API 密钥。 119 | 120 | > [!TIP] 121 | > 您也可以通过替换 `run_simple.sh` 或 `run_simple_gui.sh` 脚本中的 `MODEL_NAME` 来使用其它大模型作为生成器。 122 | 123 | ## 实验结果 124 | 我们在多个单跳及多跳知识密集型问答数据集上进行了实验。实验结果显示，相较于对比方法LevelRAG实现了非常显著的性能提升，实验结果请参考下表。 125 | 126 |
127 | 128 |
129 | 130 | ## 许可 131 | 本项目使用 MIT 许可证。有关更多信息，请参阅 [LICENSE](LICENSE) 文件。 132 | 133 | 134 | ## 引用 135 | 如果您觉得我们的工作对您有所帮助，请考虑引用我们的论文或 Star 本项目。 136 | 137 | ```bibtex 138 | @misc{zhang2025levelragenhancingretrievalaugmentedgeneration, 139 | title={LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers}, 140 | author={Zhuocheng Zhang and Yang Feng and Min Zhang}, 141 | year={2025}, 142 | eprint={2502.18139}, 143 | archivePrefix={arXiv}, 144 | primaryClass={cs.CL}, 145 | url={https://arxiv.org/abs/2502.18139}, 146 | } 147 | ``` 148 | 149 | 如果您对本项目有任何问题，欢迎在 GitHub 上创建 issue 或通过邮件联系我们（zhangzhuocheng20z@ict.acn.cn）。 150 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers 2 | 3 | 4 | ![Language](https://img.shields.io/badge/language-python-brightgreen) 5 | [![Code Style](https://img.shields.io/badge/code%20style-black-black)](https://github.com/psf/black) 6 | [![Imports: isort](https://img.shields.io/badge/imports-isort-blue)](https://pycqa.github.io/isort/) 7 | [![github license](https://img.shields.io/github/license/ictnlp/LevelRAG)](LICENSE) 8 | [![arXiv](https://img.shields.io/badge/arXiv-2502.18139-b31b1b)](https://arxiv.org/abs/2502.18139) 9 | 10 |
11 | 12 | \[ [English](README.md) | [中文](README-zh.md) \] 13 | 14 |
15 | 16 | Source code for paper [**LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers**](https://arxiv.org/abs/2502.18139). 17 | 18 | ## Overview 19 | **LevelRAG** is a two-stage retrieval-augmented generation (RAG) framework that incorporates multi-hop logic planning and hybrid retrieval to enhance both completeness and accuracy of the retrieval process. The first stage involves a high-level searcher that decomposing the user query into atomic sub-queries. The second stage utilizes multiple low-level searchers to retrieve the most relevant documents for each sub-query, which are then used to generate the final answer. In each low-level searcher, large language models (LLMs) are employed to refine the atomic queries to better fit the corresponding retriever. 20 | 21 |
22 | 23 |
24 | 25 | 26 | ## Running LevelRAG 27 | 28 | ### Prepare the Environment 29 | Our code is based on [FlexRAG](https://github.com/ictnlp/FlexRAG) project. Please follow the instruction to install FlexRAG: 30 | ```bash 31 | pip install flexrag==0.2.0 32 | ``` 33 | 34 | Download the source code of this project: 35 | ```bash 36 | git clone https://github.com/ictnlp/LevelRAG 37 | ``` 38 | 39 | ### Prepare the Retriever 40 | 41 | > [!TIP] 42 | > If you want to run LevelRAG in a simpler way, please refer to the [Running the Simple LevelRAG](#running-the-simple-levelrag) section. 43 | 44 | Before running the LevelRAG, preparing the retriever is necessary. LevelRAG employs three kind of retrievers in total, naming `DenseRetriever`, `ElasticRetriever`, and `WebRetriever`, respectively. Except for the `WebRetriever`, which does not require index construction, both the `DenseRetriever` and the `ElasticRetriever` need to prepare the index first. In our experiments, we use the wikipedia corpus provided by [Atlas](https://github.com/facebookresearch/atlas). You can download the corpus by running the following command: 45 | 46 | ```bash 47 | wget https://dl.fbaipublicfiles.com/atlas/corpora/wiki/enwiki-dec2021/text-list-100-sec.jsonl 48 | wget https://dl.fbaipublicfiles.com/atlas/corpora/wiki/enwiki-dec2021/infobox.jsonl 49 | ``` 50 | 51 | After downloading the corpus, you can run the following command to build the `DenseRetriever`: 52 | ```bash 53 | python -m flexrag.entrypoints.prepare_index \ 54 | retriever_type=dense \ 55 | file_paths=[text-list-100-sec.jsonl,infobox.jsonl] \ 56 | saving_fields=[title,section,text] \ 57 | id_field=id \ 58 | dense_config.database_path=wikipedia \ 59 | dense_config.passage_encoder_config.encoder_type=hf \ 60 | dense_config.passage_encoder_config.hf_config.model_path=facebook/contriever-msmarco \ 61 | dense_config.passage_encoder_config.hf_config.device_id=[0] \ 62 | dense_config.encode_fields=[text] \ 63 | dense_config.index_type=faiss \ 64 | dense_config.batch_size=1024 \ 65 | dense_config.log_interval=100000 66 | ``` 67 | 68 | Similarly, you can run the following command to build the `ElasticRetriever`: 69 | 70 | ```bash 71 | python -m flexrag.entrypoints.prepare_index \ 72 | retriever_type=elastic \ 73 | file_paths=[text-list-100-sec.jsonl,infobox.jsonl] \ 74 | saving_fields=[title,section,text] \ 75 | id_field=id \ 76 | elastic_config.host='http://127.0.0.1:9200/' \ 77 | elastic_config.index_name=wikipedia \ 78 | elastic_config.batch_size=512 \ 79 | elastic_config.log_interval=100000 \ 80 | reinit=True 81 | ``` 82 | 83 | > **Notice:** 84 | > Before building the `ElasticRetriever`, you need to setup the elasticsearch server. You can follow the instruction [here](https://www.elastic.co/guide/en/elasticsearch/reference/current/install-elasticsearch.html) to install the elasticsearch server. 85 | 86 | Using the `WebRetriever` does not require index construction. However, you need to prepare the Bing Search API_KEY in advance. You can visit [Bing Search API](https://www.microsoft.com/en-us/bing/apis) to get the API_KEY. 87 | 88 | ### Prepare the Genereator 89 | LevelRAG uses the `Qwen2-7B-Instruct` as the generator. You can deploy the generator using `vllm` by running the following command: 90 | 91 | ```bash 92 | python -m vllm.entrypoints.openai.api_server \ 93 | --model Qwen2-7B-Instruct \ 94 | --gpu-memory-utilization 0.95 \ 95 | --tensor-parallel-size 4 \ 96 | --port 8000 \ 97 | --host 0.0.0.0 \ 98 | --trust-remote-code 99 | ``` 100 | 101 | This command will deploy the `Qwen2-7B-Instruct` using 4 GPUs. You can adjust the `--tensor-parallel-size` and the `--gpu-memory-utilization` according to your own GPU configuration. 102 | 103 | 104 | ### Run the LevelRAG Evaluation Script 105 | After preparing the retriever, you can run the LevelRAG by running the scripts in the `scripts` folder. Before running the scripts, make sure you have substituted the placeholder variables in the scripts with the correct values. 106 | 107 | - LEVELRAG_PATH: The path to the LevelRAG repository. 108 | - DENSE_PATH: The path to the `DenseRetriever`. 109 | - BING_KEY: The Bing Search API_KEY. 110 | 111 | ### Run the LevelRAG GUI Demo 112 | We also provide a GUI demo for LevelRAG. You can run the GUI demo by running the `run_highlevel_gui.sh` script in the `scripts` folder. In the GUI, you can input the query and view the output of LevelRAG. 113 | 114 | ## Running the Simple LevelRAG 115 | If you think building the retriever is too complicated, you can run the simple version of LevelRAG by running the 116 | `run_simple.sh` script in the `scripts` folder. The simple version of LevelRAG only uses the `DenseRetriever` and does not require the `WebRetriever` and the `ElasticRetriever`. Thanks to the good multi-hop problem decomposition and sub-query adaptivity optimization, LevelRAG can achieve good performance even with a single retriever, and the running speed is faster. You can also run the `run_simple_gui.sh` script to start the GUI application of the simple version of LevelRAG. 117 | 118 | > [!NOTE] 119 | > Please make sure you have change the `API_KEY` in the `run_simple.sh` script to your own OpenAI API_KEY. 120 | 121 | ## Experimental Results 122 | We conducted experiments on multiple single-hop and multi-hop knowledge-intensive question answering datasets. The experimental results show that, compared to the baseline method, LevelRAG achieves a significant performance improvement. Please refer to the table below for the experimental results. 123 | 124 |
125 | 126 |
127 | 128 | ## License 129 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. 130 | 131 | ## Citation 132 | If you find our work useful, please consider citing our paper: 133 | 134 | ```bibtex 135 | @misc{zhang2025levelragenhancingretrievalaugmentedgeneration, 136 | title={LevelRAG: Enhancing Retrieval-Augmented Generation with Multi-hop Logic Planning over Rewriting Augmented Searchers}, 137 | author={Zhuocheng Zhang and Yang Feng and Min Zhang}, 138 | year={2025}, 139 | eprint={2502.18139}, 140 | archivePrefix={arXiv}, 141 | primaryClass={cs.CL}, 142 | url={https://arxiv.org/abs/2502.18139}, 143 | } 144 | ``` 145 | 146 | If you have any questions, feel free to create an issue on GitHub or contact us via email (zhangzhuocheng20z@ict.acn.cn). -------------------------------------------------------------------------------- /assets/ExperimentalResults.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ictnlp/LevelRAG/4ad16faaae78d8c08631a69b45f870a18933591d/assets/ExperimentalResults.png -------------------------------------------------------------------------------- /assets/LevelRAG-ch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ictnlp/LevelRAG/4ad16faaae78d8c08631a69b45f870a18933591d/assets/LevelRAG-ch.png -------------------------------------------------------------------------------- /assets/LevelRAG.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ictnlp/LevelRAG/4ad16faaae78d8c08631a69b45f870a18933591d/assets/LevelRAG.png -------------------------------------------------------------------------------- /scripts/build_dense_wiki.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -euo pipefail 4 | 5 | WIKI_FILE=text-list-100-sec.jsonl 6 | WIKI_INFOBOX=infobox.jsonl 7 | DENSE_PATH= 8 | ENCODER_PATH=facebook/contriever-msmarco 9 | 10 | python -m flexrag.entrypoints.prepare_index \ 11 | retriever_type=dense \ 12 | file_paths=[$WIKI_FILE,$WIKI_INFOBOX] \ 13 | saving_fields=[title,section,text] \ 14 | id_field=id \ 15 | dense_config.database_path=$DENSE_PATH \ 16 | dense_config.passage_encoder_config.encoder_type=hf \ 17 | dense_config.passage_encoder_config.hf_config.model_path=$ENCODER_PATH \ 18 | dense_config.passage_encoder_config.hf_config.device_id=[0] \ 19 | dense_config.encode_fields=[text] \ 20 | dense_config.index_type=faiss \ 21 | dense_config.batch_size=1024 \ 22 | dense_config.log_interval=100000 23 | 24 | -------------------------------------------------------------------------------- /scripts/build_elastic_wiki.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -euo pipefail 4 | 5 | WIKI_FILE=text-list-100-sec.jsonl 6 | WIKI_INFOBOX=infobox.jsonl 7 | ELASTIC_HOST=http://127.0.0.1:9200/ 8 | 9 | python -m flexrag.entrypoints.prepare_index \ 10 | retriever_type=elastic \ 11 | file_paths=[$WIKI_FILE,$WIKI_INFOBOX] \ 12 | saving_fields=[title,section,text] \ 13 | id_field=id \ 14 | elastic_config.host=$ELASTIC_HOST \ 15 | elastic_config.index_name=wiki \ 16 | elastic_config.batch_size=512 \ 17 | elastic_config.log_interval=100 \ 18 | reinit=True 19 | -------------------------------------------------------------------------------- /scripts/run_highlevel.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | LEVELRAG_PATH="" 4 | MODEL_NAME="Qwen2-7B-Instruct" 5 | BASE_URL="http://127.0.0.1:8000/v1" 6 | ELASTIC_HOST=http://127.0.0.1:9200/ 7 | DENSE_PATH= 8 | ENCODER_PATH=facebook/contriever-msmarco 9 | BING_KEY="" 10 | 11 | 12 | python -m flexrag.entrypoints.run_assistant \ 13 | user_module=$LEVELRAG_PATH/searchers \ 14 | name=nq \ 15 | split=test \ 16 | assistant_type=highlevel \ 17 | highlevel_config.searchers=[keyword,web,dense] \ 18 | highlevel_config.decompose=True \ 19 | highlevel_config.summarize_for_decompose=True \ 20 | highlevel_config.summarize_for_answer=True \ 21 | highlevel_config.keyword_config.rewrite_query=adaptive \ 22 | highlevel_config.keyword_config.feedback_depth=3 \ 23 | highlevel_config.keyword_config.response_type=short \ 24 | highlevel_config.keyword_config.generator_type=openai \ 25 | highlevel_config.keyword_config.openai_config.model_name=$MODEL_NAME \ 26 | highlevel_config.keyword_config.openai_config.base_url=$BASE_URL \ 27 | highlevel_config.keyword_config.gen_cfg.do_sample=False \ 28 | highlevel_config.keyword_config.host=$ELASTIC_HOST \ 29 | highlevel_config.keyword_config.index_name=wiki_2021 \ 30 | highlevel_config.dense_config.rewrite_query=adaptive \ 31 | highlevel_config.dense_config.response_type=short \ 32 | highlevel_config.dense_config.generator_type=openai \ 33 | highlevel_config.dense_config.openai_config.model_name=$MODEL_NAME \ 34 | highlevel_config.dense_config.openai_config.base_url=$BASE_URL \ 35 | highlevel_config.dense_config.gen_cfg.do_sample=False \ 36 | highlevel_config.dense_config.database_path=$DENSE_PATH \ 37 | highlevel_config.dense_config.index_type=faiss \ 38 | highlevel_config.dense_config.query_encoder_config.encoder_type=hf \ 39 | highlevel_config.dense_config.query_encoder_config.hf_config.model_path=$ENCODER_PATH \ 40 | highlevel_config.dense_config.query_encoder_config.hf_config.device_id=[0] \ 41 | highlevel_config.web_config.search_engine_type=bing \ 42 | highlevel_config.web_config.bing_config.subscription_key=$BING_KEY \ 43 | highlevel_config.web_config.web_reader_type=snippet \ 44 | highlevel_config.web_config.rewrite_query=False \ 45 | highlevel_config.web_config.response_type=short \ 46 | highlevel_config.web_config.generator_type=openai \ 47 | highlevel_config.web_config.openai_config.model_name=$MODEL_NAME \ 48 | highlevel_config.web_config.openai_config.base_url=$BASE_URL \ 49 | highlevel_config.web_config.gen_cfg.do_sample=False \ 50 | highlevel_config.response_type=short \ 51 | highlevel_config.generator_type=openai \ 52 | highlevel_config.openai_config.model_name=$MODEL_NAME \ 53 | highlevel_config.openai_config.base_url=$BASE_URL \ 54 | highlevel_config.gen_cfg.do_sample=False \ 55 | eval_config.metrics_type=[retrieval_success_rate,generation_f1,generation_em] \ 56 | eval_config.retrieval_success_rate_config.eval_field=text \ 57 | eval_config.response_preprocess.processor_type=[simplify_answer] \ 58 | log_interval=10 59 | -------------------------------------------------------------------------------- /scripts/run_highlevel_gui.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | LEVELRAG_PATH="" 4 | MODEL_NAME="Qwen2-7B-Instruct" 5 | BASE_URL="http://127.0.0.1:8000/v1" 6 | ELASTIC_HOST=http://127.0.0.1:9200/ 7 | DENSE_PATH= 8 | ENCODER_PATH=facebook/contriever-msmarco 9 | BING_KEY="" 10 | 11 | 12 | python -m flexrag.entrypoints.run_interactive \ 13 | user_module=$LEVELRAG_PATH/searchers \ 14 | assistant_type=highlevel \ 15 | highlevel_config.searchers=[keyword,web,dense] \ 16 | highlevel_config.decompose=True \ 17 | highlevel_config.summarize_for_decompose=True \ 18 | highlevel_config.summarize_for_answer=True \ 19 | highlevel_config.keyword_config.rewrite_query=adaptive \ 20 | highlevel_config.keyword_config.feedback_depth=3 \ 21 | highlevel_config.keyword_config.response_type=short \ 22 | highlevel_config.keyword_config.generator_type=openai \ 23 | highlevel_config.keyword_config.openai_config.model_name=$MODEL_NAME \ 24 | highlevel_config.keyword_config.openai_config.base_url=$BASE_URL \ 25 | highlevel_config.keyword_config.gen_cfg.do_sample=False \ 26 | highlevel_config.keyword_config.host=$ELASTIC_HOST \ 27 | highlevel_config.keyword_config.index_name=wiki_2021 \ 28 | highlevel_config.dense_config.rewrite_query=adaptive \ 29 | highlevel_config.dense_config.response_type=short \ 30 | highlevel_config.dense_config.generator_type=openai \ 31 | highlevel_config.dense_config.openai_config.model_name=$MODEL_NAME \ 32 | highlevel_config.dense_config.openai_config.base_url=$BASE_URL \ 33 | highlevel_config.dense_config.gen_cfg.do_sample=False \ 34 | highlevel_config.dense_config.database_path=$DENSE_PATH \ 35 | highlevel_config.dense_config.index_type=faiss \ 36 | highlevel_config.dense_config.query_encoder_config.encoder_type=hf \ 37 | highlevel_config.dense_config.query_encoder_config.hf_config.model_path=$ENCODER_PATH \ 38 | highlevel_config.dense_config.query_encoder_config.hf_config.device_id=[0] \ 39 | highlevel_config.web_config.search_engine_type=bing \ 40 | highlevel_config.web_config.bing_config.subscription_key=$BING_KEY \ 41 | highlevel_config.web_config.web_reader_type=snippet \ 42 | highlevel_config.web_config.rewrite_query=False \ 43 | highlevel_config.web_config.response_type=short \ 44 | highlevel_config.web_config.generator_type=openai \ 45 | highlevel_config.web_config.openai_config.model_name=$MODEL_NAME \ 46 | highlevel_config.web_config.openai_config.base_url=$BASE_URL \ 47 | highlevel_config.web_config.gen_cfg.do_sample=False \ 48 | highlevel_config.response_type=short \ 49 | highlevel_config.generator_type=openai \ 50 | highlevel_config.openai_config.model_name=$MODEL_NAME \ 51 | highlevel_config.openai_config.base_url=$BASE_URL \ 52 | highlevel_config.gen_cfg.do_sample=False 53 | -------------------------------------------------------------------------------- /scripts/run_simple.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | LEVELRAG_PATH="" 4 | MODEL_NAME="gpt-4o" 5 | API_KEY="" 6 | 7 | 8 | python -m flexrag.entrypoints.run_assistant \ 9 | user_module=$LEVELRAG_PATH/searchers \ 10 | name=nq \ 11 | split=test \ 12 | assistant_type=highlevel \ 13 | highlevel_config.searchers=[dense] \ 14 | highlevel_config.decompose=True \ 15 | highlevel_config.summarize_for_decompose=True \ 16 | highlevel_config.summarize_for_answer=True \ 17 | highlevel_config.dense_config.rewrite_query=adaptive \ 18 | highlevel_config.dense_config.response_type=short \ 19 | highlevel_config.dense_config.generator_type=openai \ 20 | highlevel_config.dense_config.openai_config.model_name=$MODEL_NAME \ 21 | highlevel_config.dense_config.openai_config.api_key=$API_KEY \ 22 | highlevel_config.dense_config.gen_cfg.do_sample=False \ 23 | highlevel_config.dense_config.hf_repo='FlexRAG/wiki2021_atlas_contriever' \ 24 | highlevel_config.response_type=short \ 25 | highlevel_config.generator_type=openai \ 26 | highlevel_config.openai_config.model_name=$MODEL_NAME \ 27 | highlevel_config.openai_config.api_key=$api_key \ 28 | highlevel_config.gen_cfg.do_sample=False \ 29 | eval_config.metrics_type=[retrieval_success_rate,generation_f1,generation_em] \ 30 | eval_config.retrieval_success_rate_config.eval_field=text \ 31 | eval_config.response_preprocess.processor_type=[simplify_answer] \ 32 | log_interval=10 33 | -------------------------------------------------------------------------------- /scripts/run_simple_gui.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | LEVELRAG_PATH="" 4 | MODEL_NAME="gpt-4o" 5 | API_KEY="" 6 | 7 | 8 | python -m flexrag.entrypoints.run_interactive \ 9 | user_module=$LEVELRAG_PATH/searchers \ 10 | assistant_type=highlevel \ 11 | highlevel_config.searchers=[dense] \ 12 | highlevel_config.decompose=True \ 13 | highlevel_config.summarize_for_decompose=True \ 14 | highlevel_config.summarize_for_answer=True \ 15 | highlevel_config.dense_config.rewrite_query=adaptive \ 16 | highlevel_config.dense_config.response_type=short \ 17 | highlevel_config.dense_config.generator_type=openai \ 18 | highlevel_config.dense_config.openai_config.model_name=$MODEL_NAME \ 19 | highlevel_config.dense_config.openai_config.api_key=$API_KEY \ 20 | highlevel_config.dense_config.gen_cfg.do_sample=False \ 21 | highlevel_config.dense_config.hf_repo='FlexRAG/wiki2021_atlas_contriever' \ 22 | highlevel_config.response_type=short \ 23 | highlevel_config.generator_type=openai \ 24 | highlevel_config.openai_config.model_name=$MODEL_NAME \ 25 | highlevel_config.openai_config.api_key=$API_KEY \ 26 | highlevel_config.gen_cfg.do_sample=False 27 | -------------------------------------------------------------------------------- /searchers/__init__.py: -------------------------------------------------------------------------------- 1 | from .keyword_searcher import KeywordSearcher, KeywordSearcherConfig 2 | from .dense_searcher import DenseSearcher, DenseSearcherConfig 3 | from .high_level_searcher import HighLevalSearcher, HighLevelSearcherConfig 4 | from .hybrid_searcher import HybridSearcher, HybridSearcherConfig 5 | from .searcher import BaseSearcher, BaseSearcherConfig 6 | from .web_searcher import WebSearcher, WebSearcherConfig 7 | 8 | 9 | __all__ = [ 10 | "BaseSearcher", 11 | "BaseSearcherConfig", 12 | "KeywordSearcher", 13 | "KeywordSearcherConfig", 14 | "WebSearcher", 15 | "WebSearcherConfig", 16 | "DenseSearcher", 17 | "DenseSearcherConfig", 18 | "HybridSearcher", 19 | "HybridSearcherConfig", 20 | "HighLevalSearcher", 21 | "HighLevelSearcherConfig", 22 | "Searchers", 23 | "SearcherConfig", 24 | "load_searcher", 25 | ] 26 | -------------------------------------------------------------------------------- /searchers/dense_searcher.py: -------------------------------------------------------------------------------- 1 | import os 2 | from copy import deepcopy 3 | from dataclasses import dataclass 4 | from typing import Optional 5 | 6 | from flexrag.assistant import ASSISTANTS 7 | from flexrag.common_dataclass import RetrievedContext 8 | from flexrag.prompt import ChatTurn, ChatPrompt 9 | from flexrag.retriever import DenseRetriever, DenseRetrieverConfig, LocalRetriever 10 | from flexrag.utils import Choices, LOGGER_MANAGER 11 | 12 | from .searcher import BaseSearcher, BaseSearcherConfig 13 | 14 | 15 | logger = LOGGER_MANAGER.getLogger("levelrag.dense_searcher") 16 | 17 | 18 | @dataclass 19 | class DenseSearcherConfig(BaseSearcherConfig, DenseRetrieverConfig): 20 | rewrite_query: Choices(["never", "pseudo", "adaptive"]) = "never" # type: ignore 21 | max_rewrite_depth: int = 3 22 | hf_repo: Optional[str] = None 23 | 24 | 25 | @ASSISTANTS("dense", config_class=DenseSearcherConfig) 26 | class DenseSearcher(BaseSearcher): 27 | def __init__(self, cfg: DenseSearcherConfig) -> None: 28 | super().__init__(cfg) 29 | # setup Dense Searcher 30 | self.rewrite = cfg.rewrite_query 31 | self.rewrite_depth = cfg.max_rewrite_depth 32 | 33 | # load Dense Retrieve 34 | if cfg.hf_repo is not None: 35 | self.retriever = LocalRetriever.load_from_hub(cfg.hf_repo) 36 | else: 37 | self.retriever = DenseRetriever(cfg) 38 | 39 | # load prompts 40 | self.rewrite_with_ctx_prompt = ChatPrompt.from_json( 41 | os.path.join( 42 | os.path.dirname(__file__), 43 | "prompts", 44 | "rewrite_by_answer_with_context_prompt.json", 45 | ) 46 | ) 47 | self.rewrite_wo_ctx_prompt = ChatPrompt.from_json( 48 | os.path.join( 49 | os.path.dirname(__file__), 50 | "prompts", 51 | "rewrite_by_answer_without_context_prompt.json", 52 | ) 53 | ) 54 | self.verify_prompt = ChatPrompt.from_json( 55 | os.path.join( 56 | os.path.dirname(__file__), 57 | "prompts", 58 | "verify_prompt.json", 59 | ) 60 | ) 61 | return 62 | 63 | def search( 64 | self, question: str 65 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 66 | # rewrite the query 67 | if self.rewrite == "pseudo": 68 | query_to_search = self.rewrite_query(question) 69 | else: 70 | query_to_search = question 71 | 72 | # begin adaptive search 73 | ctxs = [] 74 | search_history = [] 75 | verification = False 76 | rewrite_depth = 0 77 | while (not verification) and (rewrite_depth < self.rewrite_depth): 78 | rewrite_depth += 1 79 | 80 | # search 81 | ctxs = self.retriever.search(query=[query_to_search])[0] 82 | search_history.append( 83 | { 84 | "query": query_to_search, 85 | "ctxs": ctxs, 86 | } 87 | ) 88 | 89 | # verify the contexts 90 | if self.rewrite == "adaptive": 91 | verification = self.verify_contexts(ctxs, question) 92 | else: 93 | verification = True 94 | 95 | # adaptive rewrite 96 | if (not verification) and (rewrite_depth < self.rewrite_depth): 97 | if rewrite_depth == 1: 98 | query_to_search = self.rewrite_query(question) 99 | else: 100 | query_to_search = self.rewrite_query(question, ctxs) 101 | return ctxs, search_history 102 | 103 | def rewrite_query( 104 | self, question: str, contexts: list[RetrievedContext] = [] 105 | ) -> str: 106 | # Rewrite the query to be more informative 107 | if len(contexts) == 0: 108 | prompt = deepcopy(self.rewrite_wo_ctx_prompt) 109 | user_prompt = f"Question: {question}" 110 | else: 111 | prompt = deepcopy(self.rewrite_with_ctx_prompt) 112 | user_prompt = "" 113 | for n, ctx in enumerate(contexts): 114 | user_prompt += f"Context {n}: {ctx.data['text']}\n\n" 115 | user_prompt += f"Question: {question}" 116 | prompt.update(ChatTurn(role="user", content=user_prompt)) 117 | query = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 118 | return f"{question} {query}" 119 | 120 | def verify_contexts( 121 | self, 122 | contexts: list[RetrievedContext], 123 | question: str, 124 | ) -> bool: 125 | prompt = deepcopy(self.verify_prompt) 126 | user_prompt = "" 127 | for n, ctx in enumerate(contexts): 128 | user_prompt += f"Context {n}: {ctx.data['text']}\n\n" 129 | user_prompt += f"Question: {question}" 130 | prompt.update(ChatTurn(role="user", content=user_prompt)) 131 | response = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 132 | return "yes" in response.lower() 133 | -------------------------------------------------------------------------------- /searchers/high_level_searcher.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from copy import deepcopy 4 | from dataclasses import dataclass 5 | 6 | from flexrag.assistant import ASSISTANTS 7 | from flexrag.common_dataclass import RetrievedContext 8 | from flexrag.prompt import ChatPrompt, ChatTurn 9 | 10 | from .hybrid_searcher import HybridSearcher, HybridSearcherConfig 11 | 12 | 13 | @dataclass 14 | class HighLevelSearcherConfig(HybridSearcherConfig): 15 | decompose: bool = False 16 | max_decompose_times: int = 100 17 | summarize_for_decompose: bool = False 18 | summarize_for_answer: bool = False 19 | 20 | 21 | @ASSISTANTS("highlevel", config_class=HighLevelSearcherConfig) 22 | class HighLevalSearcher(HybridSearcher): 23 | def __init__(self, cfg: HighLevelSearcherConfig) -> None: 24 | super().__init__(cfg) 25 | # set basic args 26 | if not cfg.decompose: 27 | self.max_decompose_times = 0 28 | else: 29 | self.max_decompose_times = cfg.max_decompose_times 30 | self.summarize_for_decompose = cfg.summarize_for_decompose 31 | self.summarize_for_answer = cfg.summarize_for_answer 32 | 33 | # load prompt 34 | self.decompose_prompt_w_ctx = ChatPrompt.from_json( 35 | os.path.join( 36 | os.path.dirname(__file__), 37 | "prompts", 38 | "decompose_with_context_prompt.json", 39 | ) 40 | ) 41 | self.decompose_prompt_wo_ctx = ChatPrompt.from_json( 42 | os.path.join( 43 | os.path.dirname(__file__), 44 | "prompts", 45 | "decompose_without_context_prompt.json", 46 | ) 47 | ) 48 | self.summarize_prompt = ChatPrompt.from_json( 49 | os.path.join( 50 | os.path.dirname(__file__), 51 | "prompts", 52 | "summarize_by_answer_prompt.json", 53 | ) 54 | ) 55 | return 56 | 57 | def decompose_question( 58 | self, 59 | question: str, 60 | search_history: list[dict[str, str | RetrievedContext]] = [], 61 | ) -> list[str]: 62 | # form prompt 63 | if len(search_history) > 0: 64 | prompt = deepcopy(self.decompose_prompt_w_ctx) 65 | ctx_str = self.compose_contexts(search_history) 66 | prompt.update( 67 | ChatTurn(role="user", content=f"Question: {question}\n\n{ctx_str}") 68 | ) 69 | else: 70 | prompt = deepcopy(self.decompose_prompt_wo_ctx) 71 | prompt.update(ChatTurn(role="user", content=f"Question: {question}")) 72 | 73 | # get response 74 | response = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 75 | if "No additional information is required" in response: 76 | return [] 77 | split_pattern = r"\[\d+\] ([^\[]+)" 78 | decompsed = re.findall(split_pattern, response) 79 | # If the question is not decomposed, fallback to original question 80 | if (len(decompsed) == 0) and (len(search_history) == 0): 81 | decompsed = [question] 82 | 83 | # deduplicate questions 84 | searched = set() 85 | for s in search_history: 86 | searched.add(s["question"]) 87 | decompsed = [i for i in decompsed if i not in searched] 88 | return decompsed 89 | 90 | def compose_contexts( 91 | self, search_history: list[dict[str, str | list[RetrievedContext]]] 92 | ) -> str: 93 | if self.summarize_for_decompose: 94 | summed_text = self.summarize_history(search_history) 95 | ctx_text = "" 96 | for n, text in enumerate(summed_text): 97 | ctx_text += f"Context {n + 1}: {text}\n\n" 98 | ctx_text = ctx_text[:-2] 99 | else: 100 | ctx_text = "" 101 | n = 1 102 | for item in search_history: 103 | for ctx in item["contexts"]: 104 | ctx_text += f"Context {n}: {ctx.data['text']}\n\n" 105 | n += 1 106 | ctx_text = ctx_text[:-2] 107 | return ctx_text 108 | 109 | def summarize_history( 110 | self, search_history: list[dict[str, str | list[RetrievedContext]]] 111 | ) -> list[str]: 112 | summed_text = [] 113 | for item in search_history: 114 | prompt = deepcopy(self.summarize_prompt) 115 | q = item["question"] 116 | usr_prompt = "" 117 | for n, ctx in enumerate(item["contexts"]): 118 | usr_prompt += f"Context {n + 1}: {ctx.data['text']}\n\n" 119 | usr_prompt += f"Question: {q}" 120 | prompt.update(ChatTurn(role="user", content=usr_prompt)) 121 | ans = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 122 | summed_text.append(ans) 123 | return summed_text 124 | 125 | def search( 126 | self, question: str 127 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 128 | contexts = [] 129 | search_history = [] 130 | decompose_times = self.max_decompose_times 131 | if decompose_times > 0: 132 | decomposed_questions = self.decompose_question(question) 133 | decompose_times -= 1 134 | else: 135 | decomposed_questions = [question] 136 | 137 | # search the decomposed_questions 138 | while len(decomposed_questions) > 0: 139 | q = decomposed_questions.pop(0) 140 | ctxs, _ = super().search(q) 141 | search_history.append({"question": q, "contexts": ctxs}) 142 | contexts.extend(ctxs) 143 | if (len(decomposed_questions) == 0) and (decompose_times > 0): 144 | decomposed_questions = self.decompose_question(question, search_history) 145 | decompose_times -= 1 146 | 147 | # post process 148 | if self.summarize_for_answer: 149 | summed_text = self.summarize_history(search_history) 150 | contexts = [ 151 | RetrievedContext( 152 | retriever="highlevel_searcher", 153 | query=j["question"], 154 | data={"text": i}, 155 | ) 156 | for i, j in zip(summed_text, search_history) 157 | ] 158 | return contexts, search_history 159 | -------------------------------------------------------------------------------- /searchers/hybrid_searcher.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | 3 | from flexrag.assistant import ASSISTANTS 4 | from flexrag.common_dataclass import RetrievedContext 5 | from flexrag.utils import Choices 6 | 7 | from .keyword_searcher import KeywordSearcher, KeywordSearcherConfig 8 | from .dense_searcher import DenseSearcher, DenseSearcherConfig 9 | from .searcher import BaseSearcher, BaseSearcherConfig 10 | from .web_searcher import WebSearcher, WebSearcherConfig 11 | 12 | 13 | @dataclass 14 | class HybridSearcherConfig(BaseSearcherConfig): 15 | searchers: list[Choices(["keyword", "web", "dense"])] = field(default_factory=list) # type: ignore 16 | keyword_config: KeywordSearcherConfig = field(default_factory=KeywordSearcherConfig) # fmt: skip 17 | web_config: WebSearcherConfig = field(default_factory=WebSearcherConfig) 18 | dense_config: DenseSearcherConfig = field(default_factory=DenseSearcherConfig) # fmt: skip 19 | 20 | 21 | @ASSISTANTS("hybrid", config_class=HybridSearcherConfig) 22 | class HybridSearcher(BaseSearcher): 23 | def __init__(self, cfg: HybridSearcherConfig) -> None: 24 | super().__init__(cfg) 25 | # load searchers 26 | self.searchers = self.load_searchers( 27 | searchers=cfg.searchers, 28 | bm25_cfg=cfg.keyword_config, 29 | web_cfg=cfg.web_config, 30 | dense_cfg=cfg.dense_config, 31 | ) 32 | return 33 | 34 | def load_searchers( 35 | self, 36 | searchers: list[str], 37 | bm25_cfg: KeywordSearcherConfig, 38 | web_cfg: WebSearcherConfig, 39 | dense_cfg: DenseSearcherConfig, 40 | ) -> dict[str, BaseSearcher]: 41 | searcher_list = {} 42 | for searcher in searchers: 43 | match searcher: 44 | case "keyword": 45 | searcher_list[searcher] = KeywordSearcher(bm25_cfg) 46 | case "web": 47 | searcher_list[searcher] = WebSearcher(web_cfg) 48 | case "dense": 49 | searcher_list[searcher] = DenseSearcher(dense_cfg) 50 | case _: 51 | raise ValueError(f"Searcher {searcher} not supported") 52 | return searcher_list 53 | 54 | def search( 55 | self, question: str 56 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 57 | # search the question using sub-searchers 58 | contexts = [] 59 | search_history = [] 60 | for name, searcher in self.searchers.items(): 61 | ctxs = searcher.search(question)[0] 62 | contexts.extend(ctxs) 63 | search_history.append( 64 | { 65 | "searcher": name, 66 | "context": ctxs, 67 | } 68 | ) 69 | return contexts, search_history 70 | -------------------------------------------------------------------------------- /searchers/keyword_searcher.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from copy import deepcopy 4 | from dataclasses import dataclass 5 | 6 | from flexrag.assistant import ASSISTANTS 7 | from flexrag.common_dataclass import RetrievedContext 8 | from flexrag.prompt import ChatTurn, ChatPrompt 9 | from flexrag.retriever import ElasticRetriever, ElasticRetrieverConfig 10 | from flexrag.utils import Choices, LOGGER_MANAGER 11 | 12 | from .searcher import BaseSearcher, BaseSearcherConfig 13 | 14 | 15 | logger = LOGGER_MANAGER.getLogger("levelrag.keyword_searcher") 16 | 17 | 18 | @dataclass 19 | class KeywordSearcherConfig(BaseSearcherConfig, ElasticRetrieverConfig): 20 | rewrite_query: Choices(["always", "never", "adaptive"]) = "never" # type: ignore 21 | feedback_depth: int = 1 22 | 23 | 24 | @ASSISTANTS("keyword", config_class=KeywordSearcherConfig) 25 | class KeywordSearcher(BaseSearcher): 26 | def __init__(self, cfg: KeywordSearcherConfig) -> None: 27 | super().__init__(cfg) 28 | # setup Keyword Searcher 29 | self.rewrite = cfg.rewrite_query 30 | self.feedback_depth = cfg.feedback_depth 31 | 32 | # load ElasticSearch Retriever 33 | self.retriever = ElasticRetriever(cfg) 34 | 35 | # load prompts 36 | self.rewrite_prompt = ChatPrompt.from_json( 37 | os.path.join( 38 | os.path.dirname(__file__), 39 | "prompts", 40 | "bm25_rewrite_prompt.json", 41 | ) 42 | ) 43 | self.verify_prompt = ChatPrompt.from_json( 44 | os.path.join( 45 | os.path.dirname(__file__), 46 | "prompts", 47 | "verify_prompt.json", 48 | ) 49 | ) 50 | self.refine_prompts = { 51 | "extend": ChatPrompt.from_json( 52 | os.path.join( 53 | os.path.dirname(__file__), 54 | "prompts", 55 | "bm25_refine_extend_prompt.json", 56 | ) 57 | ), 58 | "filter": ChatPrompt.from_json( 59 | os.path.join( 60 | os.path.dirname(__file__), 61 | "prompts", 62 | "bm25_refine_filter_prompt.json", 63 | ) 64 | ), 65 | "emphasize": ChatPrompt.from_json( 66 | os.path.join( 67 | os.path.dirname(__file__), 68 | "prompts", 69 | "bm25_refine_emphasize_prompt.json", 70 | ) 71 | ), 72 | } 73 | return 74 | 75 | def search( 76 | self, question: str 77 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 78 | retrieval_history = [] 79 | 80 | # rewrite the query 81 | match self.rewrite: 82 | case "always": 83 | query_to_search = self.rewrite_query(question) 84 | case "never": 85 | query_to_search = question 86 | case "adaptive": 87 | ctxs = self.retriever.search(query=[question])[0] 88 | verification = self.verify_contexts(ctxs, question) 89 | retrieval_history.append( 90 | { 91 | "query": question, 92 | "ctxs": ctxs, 93 | } 94 | ) 95 | if verification: 96 | return ctxs, retrieval_history 97 | query_to_search = self.rewrite_query(question) 98 | 99 | # begin BFS search 100 | search_stack = [(query_to_search, 1)] 101 | total_depth = self.feedback_depth + 1 102 | while len(search_stack) > 0: 103 | # search 104 | query_to_search, depth = search_stack.pop(0) 105 | ctxs = self.retriever.search(query=[query_to_search])[0] 106 | retrieval_history.append( 107 | { 108 | "query": query_to_search, 109 | "ctxs": ctxs, 110 | } 111 | ) 112 | 113 | # verify contexts 114 | if total_depth > 1: 115 | verification = self.verify_contexts(ctxs, question) 116 | else: 117 | verification = True 118 | if verification: 119 | break 120 | 121 | # if depth is already at the maximum, stop expanding 122 | if depth >= total_depth: 123 | continue 124 | 125 | # expand the search stack 126 | refined = self.refine_query( 127 | contexts=ctxs, 128 | base_query=question, 129 | current_query=query_to_search, 130 | ) 131 | search_stack.extend([(rq, depth + 1) for rq in refined]) 132 | return ctxs, retrieval_history 133 | 134 | def refine_query( 135 | self, 136 | contexts: list[RetrievedContext], 137 | base_query: str, 138 | current_query: str, 139 | ) -> list[str]: 140 | refined_queries = [] 141 | for prompt_type in self.refine_prompts: 142 | # prepare prompt 143 | prompt = deepcopy(self.refine_prompts[prompt_type]) 144 | ctx_str = "" 145 | for n, ctx in enumerate(contexts): 146 | ctx_str += f"Context {n}: {ctx.data['text']}\n\n" 147 | prompt.history[-1].content = ( 148 | f"{ctx_str}{prompt.history[-1].content}\n\n" 149 | f"Current query: {current_query}\n\n" 150 | f"The information you are looking for: {base_query}" 151 | ) 152 | response = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 153 | 154 | # prepare re patterns 155 | response_ = re.escape(response) 156 | pattern = f'("{response_}"(\^\d)?)|({response_})' 157 | 158 | # append refined query 159 | if prompt_type == "extend": 160 | refined_queries.append(f"{current_query} {response}") 161 | elif prompt_type == "filter": 162 | if re.search(pattern, current_query): 163 | refined_queries.append(re.sub(pattern, "", current_query)) 164 | else: 165 | refined_queries.append(f'{current_query} -"{response}"') 166 | elif prompt_type == "emphasize": 167 | if re.search(pattern, current_query): 168 | try: 169 | current_weight = re.search(pattern, current_query).group(2) 170 | current_weight = int(current_weight[1:]) 171 | except: 172 | current_weight = 1 173 | repl = re.escape(f'"{response}"^{current_weight + 1}') 174 | new_query = re.sub(pattern, repl, current_query) 175 | refined_queries.append(new_query) 176 | else: 177 | refined_queries.append(f'"{response}" {current_query}') 178 | return refined_queries 179 | 180 | def rewrite_query(self, info: str) -> str: 181 | # Rewrite the query to be more informative 182 | prompt = deepcopy(self.rewrite_prompt) 183 | prompt.update(ChatTurn(role="user", content=info)) 184 | query = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 185 | return query 186 | 187 | def verify_contexts( 188 | self, 189 | contexts: list[RetrievedContext], 190 | question: str, 191 | ) -> bool: 192 | prompt = deepcopy(self.verify_prompt) 193 | user_prompt = "" 194 | for n, ctx in enumerate(contexts): 195 | user_prompt += f"Context {n}: {ctx.data['text']}\n\n" 196 | user_prompt += f"Topic: {question}" 197 | prompt.update(ChatTurn(role="user", content=user_prompt)) 198 | response = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 199 | return "yes" in response.lower() 200 | -------------------------------------------------------------------------------- /searchers/prompts/bm25_refine_emphasize_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "You are an AI assistant tasked with helping users refine their search queries. When given a context retrieved from an initial query, your goal is to suggest the keywords that should be emphasized to help obtain more relevant and specific information.\n\nFollow these guidelines:\n1. Analyze the Given Context: Carefully read the provided context to understand what information has already been retrieved.\n2. Identify Missing Information: Determine what specific details or aspects might be missing or could be more relevant to the user's needs.\n3. Indicate Important Keywords: Provide keywords or phrases that should be emphasized to refine the search and yield more precise results." 5 | }, 6 | "history": [ 7 | { 8 | "role": "user", 9 | "content": "The above context was retrieved using the given query. However, the information may not fully address user needs. To refine the search and obtain more relevant results, please suggest **one** keyword or phrase that should be emphasized to narrow down the search. Please only reply your keyword and do not output any other words." 10 | } 11 | ], 12 | "demonstrations": [] 13 | } -------------------------------------------------------------------------------- /searchers/prompts/bm25_refine_extend_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "You are an AI assistant tasked with helping users refine their search queries. When given a context retrieved from an initial query, your goal is to suggest additional keywords or phrases that could help obtain more relevant and specific information.\n\nFollow these guidelines:\n1. Analyze the Given Context: Carefully read the provided context to understand what information has already been retrieved.\n2. Identify Missing Information: Determine what specific details or aspects might be missing or could be more relevant to the user's needs.\n3. Suggest Additional Keywords: Provide additional keywords or phrases that could help refine the search and yield more precise results. Ensure these keywords are directly related to the topic and can help narrow down the search to more relevant information." 5 | }, 6 | "history": [ 7 | { 8 | "role": "user", 9 | "content": "The above context was retrieved using the given query. However, the information may not fully address user needs. To refine the search and obtain more relevant results, please suggest **one** additional keyword or phrases that could help narrow down the search. Please only reply your keyword and do not output any other words." 10 | } 11 | ], 12 | "demonstrations": [] 13 | } -------------------------------------------------------------------------------- /searchers/prompts/bm25_refine_filter_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "The following context is retrieved using the given query. However, the context may not contain the information you are looking for. Please provide one keyword to filter out irrelevant information.\n\nFollow these guidelines:\n1. Analyze the Given Context: Carefully read the provided context to understand what information has already been retrieved.\n2. Identify Missing Information: Determine what specific details or aspects might be missing or could be more relevant to the user's needs.\n 3. Indicate Filter Keyword: Provide a keyword that can help filter out irrelevant information and narrow down the search to more relevant results." 5 | }, 6 | "history": [ 7 | { 8 | "role": "user", 9 | "content": "The above context was retrieved using the given query. However, the information may not fully address user needs. To refine the search and obtain more relevant results, please suggest **one** filter word that could help narrow down the search. Please only reply your filter word and do not output any other words." 10 | } 11 | ], 12 | "demonstrations": [] 13 | } -------------------------------------------------------------------------------- /searchers/prompts/bm25_rewrite_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Suggestions for Writing Queries for BM25 Search Engine\n1. Use Descriptive Keywords: Ensure your query includes all relevant keywords that describe what you are searching for.\n2. Incorporate Rare Terms: If you know any specific or rare terms related to your search, include them.\n3. Avoid Stop Words: Common words like \"the\", \"is\", and \"and\" may dilute the effectiveness of the query.\n4. Synonyms and Related Terms: Use synonyms and related terms to cover variations in how different documents might reference the same concept.\n5. Entity Searches: When searching for specific named entity, enclose them in double quotes.\nPlease optimize the following query for the BM25 Search Engine.\nPlease only reply your query and do not output any other words." 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "What is John Mayne's occupation?" 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "\"John Mayne\" occupation job career" 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "how many oar athletes are in the olympics" 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "\"oar athletes\" olympics number count participants" 26 | } 27 | ], 28 | [ 29 | { 30 | "role": "user", 31 | "content": "who introduced the system of civil services in india" 32 | }, 33 | { 34 | "role": "assistant", 35 | "content": "india civil services introduced foundation" 36 | } 37 | ], 38 | [ 39 | { 40 | "role": "user", 41 | "content": "Which leader did Hitler meet in the Brenner Pass in WWII?" 42 | }, 43 | { 44 | "role": "assistant", 45 | "content": "Hitler \"Brenner Pass\" WWII leader meeting" 46 | } 47 | ], 48 | [ 49 | { 50 | "role": "user", 51 | "content": "Which country does the airline Garuda come from?" 52 | }, 53 | { 54 | "role": "assistant", 55 | "content": "Garuda airline country origin" 56 | } 57 | ] 58 | ] 59 | } -------------------------------------------------------------------------------- /searchers/prompts/decompose_with_context_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Please first indicate the additional knowledge needed to answer the following question. If the question can be answered without external knowledge, answer \"No additional information is required\".\n\n" 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "Question: Which magazine was started first Arthur's Magazine or First for Women?\n\nContext 1: Arthur's Magazine was started in the 19th century.\n\nContext 2: First for Women was started in 1989." 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "No additional information is required." 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "Question: What nationality was Henry Valentine Miller's wife?\n\nContext 1: Henry Valentine Miller's wife is June Miller." 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "[1] What nationality was June Miller?" 26 | } 27 | ], 28 | [ 29 | { 30 | "role": "user", 31 | "content": "Question: Are director of film Move (1970 Film) and director of film Méditerranée (1963 Film) from the same country?\n\nContext 1: Move (1970 Film) was directed by Stuart Rosenberg.\n\nContext 2: Méditerranée (1963 Film) was directed by Jean-Daniel Pollet." 32 | }, 33 | { 34 | "role": "assistant", 35 | "content": "[1] What nationality was Stuart Rosenberg?\n[2] What nationality was Jean-Daniel Pollet?" 36 | } 37 | ], 38 | [ 39 | { 40 | "role": "user", 41 | "content": "Question: Who is Rhescuporis I (Odrysian)'s paternal grandfather?\n\nContext 1: Rhescuporis I (Odrysian)'s father is Cotys III." 42 | }, 43 | { 44 | "role": "assistant", 45 | "content": "[1] Who is Cotys III's father?" 46 | } 47 | ], 48 | [ 49 | { 50 | "role": "user", 51 | "content": "Question: Which film came out first, The Love Route or Engal Aasan?\n\nContext 1: The Love Route came out in 1915.\n\nContext 2: Engal Aasan came out in 2009." 52 | }, 53 | { 54 | "role": "assistant", 55 | "content": "No additional information is required." 56 | } 57 | ] 58 | ], 59 | "demonstrations_source": [ 60 | "hotpotqa-train", 61 | "hotpotqa-train", 62 | "hotpotqa-train" 63 | ] 64 | } -------------------------------------------------------------------------------- /searchers/prompts/decompose_without_context_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Please identify the external knowledge necessary to answer the following question. If multiple concurrent pieces of external knowledge are needed, list them all. If the required external knowledge is interdependent, list only the initial piece of external knowledge needed." 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "Question: Which magazine was started first Arthur's Magazine or First for Women?" 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "[1] When Arthur's Magazine was founded.\n[2] When First for Women was founded." 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "Question: What nationality was Henry Valentine Miller's wife?" 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "[1] Who was Henry Valentine Miller's wife?" 26 | } 27 | ], 28 | [ 29 | { 30 | "role": "user", 31 | "content": "Question: Are director of film Move (1970 Film) and director of film Méditerranée (1963 Film) from the same country?" 32 | }, 33 | { 34 | "role": "assistant", 35 | "content": "[1] Who direct the film Move (1970 Film)?\n[2] Who direct the film Méditerranée (1963 Film)?" 36 | } 37 | ], 38 | [ 39 | { 40 | "role": "user", 41 | "content": "Question: Who is Rhescuporis I (Odrysian)'s paternal grandfather?" 42 | }, 43 | { 44 | "role": "assistant", 45 | "content": "[1] Who is Rhescuporis I (Odrysian)'s father?" 46 | } 47 | ], 48 | [ 49 | { 50 | "role": "user", 51 | "content": "Question: Which film came out first, The Love Route or Engal Aasan?" 52 | }, 53 | { 54 | "role": "assistant", 55 | "content": "[1] When did The Love Route come out?\n[2] When did Engal Aasan come out?" 56 | } 57 | ] 58 | ], 59 | "demonstrations_source": [ 60 | "hotpotqa-train", 61 | "hotpotqa-train", 62 | "2wikimultihopqa-train", 63 | "2wikimultihopqa-train" 64 | ] 65 | } -------------------------------------------------------------------------------- /searchers/prompts/dense_rewrite_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Suggestions for Writing Queries for Dense Retrieval Search Engine\n1. Use Natural Language: Dense retrieval models are designed to understand and process natural language. Formulate queries in complete sentences or phrases as you would ask a question in a conversation.\n2. Incorporate Context: Provide context to your query to help the model understand the specific aspect of the topic you are interested in. Contextual information improves the accuracy of the retrieval.\n3. Ask Specific Questions: Dense retrieval models perform well with specific queries. Instead of a single keyword, use detailed questions or statements to convey your information need.\n4. Avoid Overly Technical Language: While dense retrieval can handle a variety of terms, overly technical or jargon-heavy language might not be necessary. Aim for clarity and simplicity.\nPlease optimize the following query for the Dense Retrieval Search Engine.\nPlease only reply your query and do not output any other words." 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "What is John Mayne's occupation?" 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "What is John Mayne's occupation?" 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "The year my pet monster come out" 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "When did My Pet Monster come out?" 26 | } 27 | ] 28 | ] 29 | } -------------------------------------------------------------------------------- /searchers/prompts/lucene_rewrite_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Suggestions for Writing Queries for BM25 Search Engine\n1. Use Descriptive Keywords: Ensure your query includes all relevant keywords that describe what you are searching for.\n2. Incorporate Rare Terms: If you know any specific or rare terms related to your search, include them.\n3. Avoid Stop Words: Common words like \"the\", \"is\", and \"and\" may dilute the effectiveness of the query.\n4. Synonyms and Related Terms: Use synonyms and related terms to cover variations in how different documents might reference the same concept.\n5. Phrase Searches: When searching for specific named entity, enclose them in double quotes.\n6. Use Boolean Operators: Use \"+\" for terms that must contains in the documents, \"-\" for terms that must not contains in the documents, and \"OR\" for terms that are optional.\nPlease optimize the following query for the BM25 Search Engine.\nPlease only reply your query and do not output any other words." 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "What is John Mayne's occupation?" 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "+\"John Mayne\" (occupation OR job OR career OR profession)" 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "what year did my pet monster come out" 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "+\"my pet monster\" (year OR release OR debut OR came)" 26 | } 27 | ], 28 | [ 29 | { 30 | "role": "user", 31 | "content": "Which prince is Queen Elizabeth II's youngest son?" 32 | }, 33 | { 34 | "role": "assistant", 35 | "content": "+\"Queen Elizabeth II\" +youngest +(prince OR son OR child)" 36 | } 37 | ] 38 | ] 39 | } -------------------------------------------------------------------------------- /searchers/prompts/rewrite_by_answer_with_context_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Please answer the following question using the provided contexts if relevant. If the provided contexts are not relevant, base your answer on your own knowledge." 5 | }, 6 | "history": [], 7 | "demonstrations": [] 8 | } -------------------------------------------------------------------------------- /searchers/prompts/rewrite_by_answer_without_context_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Please write a passage to answer the question." 5 | }, 6 | "history": [], 7 | "demonstrations": [] 8 | } -------------------------------------------------------------------------------- /searchers/prompts/summarize_by_answer_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Answer the following question in a single sentence using the provided contexts if relevant; if not, answer directly without additional words." 5 | }, 6 | "history": [], 7 | "demonstrations": [] 8 | } -------------------------------------------------------------------------------- /searchers/prompts/verify_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Your task is to verify whether any of the following contexts contains enough information to answer the following question. Please only reply 'yes' or 'no' and do not output any other words." 5 | }, 6 | "history": [], 7 | "demonstrations": [] 8 | } -------------------------------------------------------------------------------- /searchers/prompts/web_rewrite_prompt.json: -------------------------------------------------------------------------------- 1 | { 2 | "system": { 3 | "role": "system", 4 | "content": "Please optimize the following query for the Web search engine.\nSuggestions for Writing Queries for Web Search Engines:\n1. Use Specific Keywords: Identify and use the most relevant and specific keywords related to your search topic. This helps narrow down the results to the most pertinent web pages.\n2. Phrase Searches: Enclose exact phrases in quotation marks to search for those exact words in that exact order. This is useful for finding specific quotes, names, or titles.\n3. Ask Questions: Formulate your query as a question to get direct answers. For example, \"How to cook pasta?\" is likely to return step-by-step instructions.\n4. Synonyms and Variants: Include synonyms or different variations of a word to broaden your search. For instance, \"smartphone\" and \"mobile phone\" can yield different results.\nPlease only reply your query and do not output any other words." 5 | }, 6 | "history": [], 7 | "demonstrations": [ 8 | [ 9 | { 10 | "role": "user", 11 | "content": "What is John Mayne's occupation?" 12 | }, 13 | { 14 | "role": "assistant", 15 | "content": "What is \"John Mayne\"'s occupation?" 16 | } 17 | ], 18 | [ 19 | { 20 | "role": "user", 21 | "content": "The year my pet monster come out" 22 | }, 23 | { 24 | "role": "assistant", 25 | "content": "When did \"My Pet Monster\" come out?" 26 | } 27 | ] 28 | ] 29 | } -------------------------------------------------------------------------------- /searchers/searcher.py: -------------------------------------------------------------------------------- 1 | from abc import abstractmethod 2 | from copy import deepcopy 3 | from dataclasses import dataclass, field 4 | 5 | from flexrag.assistant import PREDEFINED_PROMPTS, AssistantBase 6 | from flexrag.common_dataclass import RetrievedContext 7 | from flexrag.models import GENERATORS, GenerationConfig, GeneratorConfig 8 | from flexrag.prompt import ChatPrompt, ChatTurn 9 | from flexrag.utils import LOGGER_MANAGER, Choices 10 | 11 | logger = LOGGER_MANAGER.getLogger("levelrag.searcher") 12 | 13 | 14 | @dataclass 15 | class BaseSearcherConfig(GeneratorConfig): 16 | gen_cfg: GenerationConfig = field(default_factory=GenerationConfig) 17 | response_type: Choices(["short", "long", "original"]) = "short" # type: ignore 18 | 19 | 20 | class BaseSearcher(AssistantBase): 21 | def __init__(self, cfg: BaseSearcherConfig) -> None: 22 | self.agent = GENERATORS.load(cfg) 23 | self.gen_cfg = cfg.gen_cfg 24 | if self.gen_cfg.sample_num > 1: 25 | logger.warning("Sample num > 1 is not supported for Searcher") 26 | self.gen_cfg.sample_num = 1 27 | 28 | # load assistant prompt 29 | match cfg.response_type: 30 | case "short": 31 | self.prompt_with_ctx = PREDEFINED_PROMPTS["shortform_with_context"] 32 | self.prompt_wo_ctx = PREDEFINED_PROMPTS["shortform_without_context"] 33 | case "long": 34 | self.prompt_with_ctx = PREDEFINED_PROMPTS["longform_with_context"] 35 | self.prompt_wo_ctx = PREDEFINED_PROMPTS["longform_without_context"] 36 | case "original": 37 | self.prompt_with_ctx = ChatPrompt() 38 | self.prompt_wo_ctx = ChatPrompt() 39 | case _: 40 | raise ValueError(f"Invalid response type: {cfg.response_type}") 41 | return 42 | 43 | @abstractmethod 44 | def search( 45 | self, question: str 46 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 47 | return 48 | 49 | def answer(self, question: str) -> tuple[str, list[RetrievedContext], dict]: 50 | ctxs, history = self.search(question) 51 | response, prompt = self.answer_with_contexts(question, ctxs) 52 | return response, ctxs, {"prompt": prompt, "search_histories": history} 53 | 54 | def answer_with_contexts( 55 | self, question: str, contexts: list[RetrievedContext] = [] 56 | ) -> tuple[str, ChatPrompt]: 57 | """Answer question with given contexts 58 | 59 | Args: 60 | question (str): The question to answer. 61 | contexts (list): The contexts searched by the searcher. 62 | 63 | Returns: 64 | response (str): response to the question 65 | prompt (ChatPrompt): prompt used. 66 | """ 67 | # prepare system prompt 68 | if len(contexts) > 0: 69 | prompt = deepcopy(self.prompt_with_ctx) 70 | else: 71 | prompt = deepcopy(self.prompt_wo_ctx) 72 | 73 | # prepare user prompt 74 | usr_prompt = "" 75 | for n, context in enumerate(contexts): 76 | ctx = context.data.get("text") 77 | usr_prompt += f"Context {n + 1}: {ctx}\n\n" 78 | usr_prompt += f"Question: {question}" 79 | 80 | # generate response 81 | prompt.update(ChatTurn(role="user", content=usr_prompt)) 82 | response = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 83 | return response, prompt 84 | -------------------------------------------------------------------------------- /searchers/web_searcher.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from copy import deepcopy 4 | from dataclasses import dataclass 5 | 6 | from flexrag.assistant import ASSISTANTS 7 | from flexrag.common_dataclass import RetrievedContext 8 | from flexrag.prompt import ChatPrompt, ChatTurn 9 | from flexrag.retriever import SimpleWebRetriever, SimpleWebRetrieverConfig 10 | 11 | from .searcher import BaseSearcher, BaseSearcherConfig 12 | 13 | logger = logging.getLogger("WebSearcher") 14 | 15 | 16 | @dataclass 17 | class WebSearcherConfig(BaseSearcherConfig, SimpleWebRetrieverConfig): 18 | rewrite_query: bool = False 19 | 20 | 21 | @ASSISTANTS("web", config_class=WebSearcherConfig) 22 | class WebSearcher(BaseSearcher): 23 | def __init__(self, cfg: WebSearcherConfig) -> None: 24 | super().__init__(cfg) 25 | # setup Web Searcher 26 | self.rewrite = cfg.rewrite_query 27 | 28 | # load Web Retrieve 29 | self.retriever = SimpleWebRetriever(cfg) 30 | 31 | # load prompt 32 | self.rewrite_prompt = ChatPrompt.from_json( 33 | os.path.join( 34 | os.path.dirname(__file__), 35 | "prompts", 36 | "web_rewrite_prompt.json", 37 | ) 38 | ) 39 | return 40 | 41 | def search( 42 | self, question: str 43 | ) -> tuple[list[RetrievedContext], list[dict[str, object]]]: 44 | # initialize search stack 45 | if self.rewrite: 46 | query_to_search = self.rewrite_query(question) 47 | else: 48 | query_to_search = question 49 | ctxs = self.retriever.search(query=[query_to_search])[0] 50 | for ctx in ctxs: 51 | ctx.data["text"] = ctx.data["snippet"] 52 | return ctxs, [] 53 | 54 | def rewrite_query(self, info: str) -> str: 55 | # Rewrite the query to be more informative 56 | user_prompt = f"Query: {info}" 57 | prompt = deepcopy(self.rewrite_prompt) 58 | prompt.update(ChatTurn(role="user", content=user_prompt)) 59 | query = self.agent.chat([prompt], generation_config=self.gen_cfg)[0][0] 60 | return query 61 | --------------------------------------------------------------------------------