This repository contains the tools to reconstruct the dataset, access the traces, and perform behavioral analysis.
This repository provides a modular framework for generating, storing, and analszing agentic search traces. Below is the organizational structure of the source code:
ASQ/
βββ README.md
βββ requirements.txt
βββ setup.sh
βββ notebooks/
β βββ quickstart_trace_generation.ipynb # trace generation walkthrough
β βββ quickstart_data_access.ipynb # artifact loading & analysis walkthrough
βββ src/ # π main development root
βββ artifacts/ # artifact schemas and loaders
βββ data_utils/ # dataset and retrieval helpers
βββ trace_generation/ # agent setup, execution, extraction
Ensure you have Java installed to support the retrieval components. Use the following configuration to set up your local environment:
# create conda environment
conda create -n asq python=3.10 numpy==1.26.4 pandas -y
conda activate asq
# install dependencies
pip install -r requirements.txtUse these notebooks to get started quickly with trace generation and artifact loading and analysis.
notebooks/quickstart_trace_generation.ipynb: Generate traces.notebooks/quickstart_data_access.ipynb: Load trace artifacts and run basic analytics.
We now describe how to construct an ASQ-like dataset. For a practical walkthrough, see notebooks/quickstart_trace_generation.ipynb.
To reproduce the construction of ASQ, use the tutorial_main.py as shown in the aforementioned walkthrough about trace generation.
To instantiate one of the supported agents use the build_agent()method as follows.
from trace_generation.rag_agent.agent import build_agent
agent = build_agent(
model_type="search_r1", # agent family
model_id="Qwen-7B", # generator
retrieval_pipeline=retrieval_pipeline, # PyTerrier retrieval pipeline
agent_k=3, # number of retrieved documents
)In the following, we describe how to set the agent family, the generator and the retrieval pipeline
At the moment we support the following configurations:
model_type: agent family (search_r1orautorefine).model_id: generator checkpoint to use (see supported lists below).
The list of supported model_type agent families and of their model_id supported models can be seen using the trace_generation.help_supported_agents() and trace_generation.help_supported_generators() functions.
Use either an existing pyterrier.Transformer retrieval pipeline, or create your custom retriever by implement a PyTerrier transformer that:
- Accepts input rows with
qidandquery. - Returns rows with
qid,docno,rank,score, andtext. - Returns
titlewhen available.
See notebooks/quickstart_trace_generation.ipynb for some practical examples.
The next step is to prepare the input query set.
- If the input query set is an IRDS dataset:
- You can use our default
load_dataset()method passing the dataset's IRDS name without theirds:prefix.
- You can use our default
- If it is a custom dataset:
- build a queries DataFrame with
qidandquerycolumns.
- build a queries DataFrame with
See notebooks/quickstart_trace_generation.ipynb for some practical examples.
Use the following command to start generation:
res = agent(queries_df)Note: This process is computationally intensive. Pre-computed traces can be downloaded from Hugging Face Datasets: ASQ dataset.
By default traces are returned by the agent as an in-memory DataFrame and are not automatically saved to disk.
The raw agent output DataFrame contains qid, query, answer, all_queries, all_docids (and optionally all_thoughts) columns.
Use the DataFrameExtractor class to converts these into TSV artifacts and appends to existing files and persist artifacts in ASQ format.
extraction_order = [ExtractionType.ANSWERS, # define the extraction order
ExtractionType.ITER_QUERIES,
ExtractionType.RETRIEVED_DOCS,
ExtractionType.THOUGHTS]
extractor = DataFrameExtractor(extraction_order=extraction_order) # instantiate the extractor
results_df = agent(queries_df) # in-memory trace output
extractor.extract_all_and_save(results_df, output_paths) # persistTraces are stored in a hierarchical directory structure for easy filtering.
Directory traces/{dataset}/{retriever_config}/{agent_family}/{model}/ stores collections of traces for the following configuration:
dataset: queryset internal identifier or irds identifier (e.g., "hotpotqa-test")retriever_config: retrieval pipeline configuration (e.g., "BM25_k1000_electra_k3")agent_family: one of the supported agent families (e.g., "search_r1")model: one of the supported generators ofagent_family(e.g., "Qwen-7B")
To simplify path resolution and loading, use the artifacts.TraceCollectionConfig and OutputPaths classes, and the config.build_output_paths() methods.
See notebooks/quickstart_data_access.ipynb for a step-by-step example on how to use them.
To load full trace collection:
from artifacts import TraceCollection
trace = TraceCollection(config)
trace.load_data(load_docs=True)Artifacts can be also loaded in memory individually from a filepath.
Instead of manually building paths, you can use build_output_paths():
from trace_generation.output_paths import OutputPaths
output_paths: OutputPaths = config.build_output_paths()Then, you can load each artifact directly using their load function.
from artifacts.answers import load_answers
from artifacts.iter_queries import load_iter_queries
from artifacts.retrieved_docs import load_retrieved_docs
answers = load_answers(output_paths.answers)
iter_queries = load_iter_queries(output_paths.iter_queries)
retrieved_docs = load_retrieved_docs(output_paths.retrieved_docs)You can also pass any custom file path instead of output_paths.*.
answers = load_answers(f"{base_dir}/answers.tsv")
iter_queries = load_iter_queries(f"{base_dir}/iter_queries.tsv")
retrieved_docs = load_retrieved_docs(f"{base_dir}/retrieved_docs.tsv")After loading artifacts or traces, you can use several methods to quickly analyse them.
E.g.,:
AnswersArtifact.nunique_qids(): count answered queries.SyntheticQueriesArtifact.iteration_counts(): see how many queries required each iteration.SyntheticQueriesArtifact.query_trace_lengths(): map eachqidto its max iteration.RetrievedDocsArtifact.docids_for_qid(): inspect the docids retrieved documents for aqid.RetrievedDocsArtifact.format_runs_into_trec_format(): create a TREC run and access it withget_it_run().
Each artifact class also provides a .helper() method that prints the available methods and their purpose.
See src/quickstart_data_access.ipynb for more details and examples.
At the moment we support the following models:
| Agent Family | Supported Generators | Note |
|---|---|---|
search_r1 |
Qwen-3B, Qwen-7B, Qwen-14B | original paper |
autorefine |
Qwen-3B | original paper |
To add a new agent type, follow these steps:
- Implement the agent package (with
SUPPORTED_MODELSand a builder function) as a subpackage oftrace_generation/rag_agent/. - Register the new agent type using
trace_generation.rag_agent.register_agent().
The ASQ dataset is released under the MIT License. Individual source datasets may have their own licenses.
ASQ is derived from publicly available datasets and is intended solely for research on agentic search behaviour. The authors do not endorse or assume responsibility for the content or any biases present in the traces. The contents of these traces should not be interpreted as representing the views of the researchers or their institutions. Users are advised to apply safety and content filters when using them.
If you find our work useful, please cite it as follows:
@misc{fpezzuti2026asq,
title={A Picture of Agentic Search},
author={Pezzuti, Francesca and Frieder, Ophir and Silvestri, Fabrizio and MacAvaney, Sean and Tonellotto, Nicola},
year={2026},
eprint={2602.17518},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2602.17518},
}