An automated multi-agent framework for early-stage water treatment train design with retrieval augmentation and critic-based refinement

Hanzhang Liu^a,f,1, Zhaorui Jiang^a,b,d,1,*, Huiling Zhong^c, Jinshuo Li^b,e, Wei Pang^b, Yingfang Yuan^b,*

¹ Equal contribution | ^* Corresponding authors: Zhaorui Jiang, Yingfang Yuan

Paper Release GitHub stars License: MIT Python Hugging Face

FastAPI Streamlit RAG Multi-Agent WContBench OpenRouter

Paper · Release · Dataset · Installation · Quick Start · Evaluation · Citation

Figure 1. OpenAqua framework: a five-agent pipeline with retrieval augmentation and critic-based refinement.

Abstract

News

Highlights

Framework

Dataset and Benchmark

Download

# Option A — Hugging Face CLI
pip install -U "huggingface_hub[cli]"
huggingface-cli download zhaorui-bi/OpenAqua --repo-type dataset --local-dir ./hf_openaqua
# Option B — snapshot in Python
python - <<'PY'
from huggingface_hub import snapshot_download
snapshot_download(
 repo_id="zhaorui-bi/OpenAqua",
 repo_type="dataset",
 local_dir="./hf_openaqua",
)
PY

mkdir -p water_treatment_agent/data
unzip hf_openaqua/data.zip -d water_treatment_agent/data
cd water_treatment_agent
python scripts/build_indexes.py

Installation

git clone https://github.com/zhaorui-bi/OpenAqua.git
cd OpenAqua/water_treatment_agent
python -m venv .venv
source .venv/bin/activate # Windows: .venv\Scripts\activate
pip install -r requirements.txt
cp .env.example .env

OPENROUTER_API_KEY=sk-or-v1-your-key-here

Quick Start

1. Data and indexes

cd water_treatment_agent
python scripts/build_indexes.py

2. API

uvicorn app.api.main:app --reload --host 0.0.0.0 --port 8000

curl -X POST http://localhost:8000/recommend \
 -H "Content-Type: application/json" \
 -d '{
 "query": {
 "raw_query": "Groundwater with arsenic around 150 ug/L, low budget, no brine disposal",
 "source_water": "groundwater",
 "contaminants": ["arsenic"],
 "treatment_targets": {
 "arsenic_ug_L": 10,
 "compliance_standard": "WHO"
 },
 "constraints": {
 "budget": "low",
 "brine_disposal": false
 }
 },
 "top_k": 3
 }'

3. GUI

cd water_treatment_agent
streamlit run gui/app.py

4. Offline demo

cd water_treatment_agent
python scripts/run_full_demo.py

Evaluation

# Treatment-train evaluation
python test_openaqua.py \
 --benchmark-dir hf_openaqua/WContBench/WContBench_Easy \
 --predictions predictions.json \
 --output eval_easy.json
# Retrieval evaluation
python test_retrieval.py \
 --benchmark-dir hf_openaqua/WContBench/WContBench_Easy \
 --predictions retrieval_preds.json

python test_openaqua.py --write-example-predictions example_preds.json

cd water_treatment_agent
pytest tests -v

Repository

OpenAqua/
├── assets/ # logo + framework figure (paper-resolution + web preview)
├── CHANGELOG.md
├── test_openaqua.py # WContBench train-level metrics
├── test_retrieval.py # WContBench retrieval metrics
└── water_treatment_agent/
 ├── app/
 │ ├── agents/ # parser, retrieval, planner, critic, explainer
 │ ├── api/ # FastAPI entrypoint and routes
 │ ├── core/ # schemas, config, taxonomy, rules
 │ ├── rag/ # corpus builder, hybrid retriever, reranker
 │ ├── utils/ # scoring, evidence binding
 │ └── workflows/ # end-to-end orchestration
 ├── gui/ # Streamlit recommend / health / ingest
 ├── scripts/ # build_indexes, parse_pdf_cases, run_full_demo
 ├── tests/
 ├── requirements.txt
 └── .env.example

water_treatment_agent/data/
├── unit-level/tdb/ # from data.zip
├── unit-level/taxonomy.json
├── case-level/kb_cases.json
└── processed/indexes/ # created by scripts/build_indexes.py

Citation

@article{liu2026openaqua,
 title = {OpenAqua: An automated multi-agent framework for
 early-stage water treatment train design with retrieval
 augmentation and critic-based refinement},
 author = {Liu, Hanzhang and Jiang, Zhaorui and Zhong, Huiling
 and Li, Jinshuo and Pang, Wei and Yuan, Yingfang},
 journal = {Water Research},
 pages = {126761},
 year = {2026},
 issn = {0043-1354},
 doi = {10.1016/j.watres.2026.126761},
 url = {https://doi.org/10.1016/j.watres.2026.126761},
 note = {Equal contribution: Hanzhang Liu and Zhaorui Jiang.
 Corresponding authors: Zhaorui Jiang and Yingfang Yuan.
 Code: https://github.com/zhaorui-bi/OpenAqua.
 Dataset: https://huggingface.co/datasets/zhaorui-bi/OpenAqua}
}

License

Acknowledgements

Dataset Code

_{OpenAqua · multi-agent water treatment train design}