PlanRAG: Plan-then-Retrieval Augmented Generation Framework

The building scenario requires Victoria 3 game savefiles (provided in /data/building/raw/). To generate simulated_questions.json in /data/building/questions/, run the simulator and example generator sequentially. This process takes approximately one hour.

export PYTHONPATH=.
python ./src/data_parsers/building/queries_gen/simulator.py
python ./src/data_parsers/building/example_gen/main.py

PlanRAG uses Neo4j to access Labeled Property Graph (LPG) data. On Ubuntu, install Neo4j Enterprise and the Cypher shell using the following steps:

1. Install dependencies and add the Neo4j repository.
2. Install neo4j-enterprise and cypher-shell.
3. Log in with the default credentials (neo4j/neo4j) and change the password immediately.

# Ubuntu
sudo apt update
sudo apt install apt-transport-https ca-certificates curl software-properties-common
sudo curl -fsSL https://debian.neo4j.com/neotechnology.gpg.key | sudo apt-key add -
sudo add-apt-repository "deb https://debian.neo4j.com stable 4.1"
sudo apt install neo4j-enterprise
sudo apt install cypher-shell

# Login with the default user/password 'neo4j'/'neo4j'.
cypher-shell -u neo4j
# Then, change to a new password and enter ':exit' to return to the terminal.

For relational data in the DQA benchmark, install MySQL on Ubuntu. You must create a database named DQA_rdb.

Note: To allow the Python environment to access MySQL without sudo privileges, you should update the root user to use mysql_native_password.

# Ubuntu
wget https://dev.mysql.com/get/mysql-apt-config_0.8.14-1_all.deb
sudo dpkg -i mysql-apt-config_0.8.14-1_all.deb
# Then, select 'OK' to proceed to the next step.
sudo apt update
sudo apt install -y mysql-server

# Login with root user to create the default database 'DQA_rdb'.
sudo mysql -u root
mysql> CREATE DATABASE DQA_rdb;
# (Optional) Make MySQL accessible without sudo privilege in the Python environment.
mysql> USE mysql;
mysql> update user set plugin='mysql_native_password' where user='root';
mysql> flush privileges;
# Then, enter ':exit' to return to the terminal.

Create a Conda environment with Python 3.8 and install the required dependencies.

conda create -n planRAG -y python=3.8 && conda activate planRAG
pip install -r requirements.txt

To use open-source models, it is highly recommended to use vllm as the model method.

1. Install vllm via pip.
2. Deploy the model as an OpenAI-compatible server using vllm.entrypoints.openai.api_server.
3. Run src/main.py with --model_method vllm.

pip install vllm

# Deploy model server
python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-70b-chat-hf

# Run experiment
python src/main.py \
  --technique IterRAG \
  --scenario locating \
  --database relational \
  --question_num 1 \
  --model meta-llama/Llama-2-70b-chat-hf \
  --model_method vllm

The locating scenario requires Europa Universalis IV game savefiles (provided in /data/locating/raw/). To generate simulated_question.json in /data/locating/questions/, run the simulator and example generator sequentially. This process takes approximately 30 minutes.

export PYTHONPATH=.
python ./src/data_parsers/locating/queries_gen/simulator.py
python ./src/data_parsers/locating/example_gen/main.py

Fill in the config.json file with your API keys and database connection details.

Important Recommendations:

• Use OpenAI's GPT-4.0 for best performance.
• Do not use the default password neo4j for the Neo4j configuration.

{
    "OPENAI_API_KEY": "Your OpenAI API key",
    "HUGGINGFACEHUB_API_TOKEN": "Your Hugging Face Hub API key", 
    "NEO4J": {
        "HOST": "Neo4j host address (default: 'bolt://localhost:7687')", 
        "USER": "User id (default: 'neo4j')",
        "PASSWORD": "User password (Do not fill with the default password 'neo4j')"
    },
    "MYSQL": {
        "HOST": "MySQL host address (default: 'localhost')",
        "USER": "User id (default: 'root')"
    }
}

PlanRAG supports different reasoning and retrieval strategies through its TaskManager abstraction. Depending on the --technique flag, the application initializes different logic:

• SingleRAG: Uses TaskManagerSingleRAG.
• IterRAG: Uses TaskManagerChain from the IterRAG module.
• PlanRAG: Uses TaskManagerChain from the PlanRAG module.
• PlanRAG_woReplan: Uses TaskManagerChain from the PlanRAG_woReplan module.

You can run the LangChain server locally by executing the main() function in src/langchain/server.py. This script uses docker-compose to pull the necessary images and start the services defined in the local docker-compose.yaml file. It relies on the langsmith.cli.main.get_docker_compose_command utility to construct the correct command.

To run the server, execute the script directly from your terminal.

python src/langchain/server.py

The json_example_generator function creates a JSON file containing structured examples for locating-related tasks. It processes raw CSV data (containing country-to-answer mappings) and uses parsers to extract geographical context (countries, provinces, and trade provinces) to build complex prompt components including business rules, questions, and goals.

Parameters:

• file_dir (str): The destination path for the generated JSON file. Defaults to ./data/locating/questions/questions.json.
• raw_dirs (list of str): A list of paths to raw CSV files. Each file is expected to contain a Country column and an Answer column. Defaults to `[