inputs/AgentSelect/other/PartIII_Compositional_Agents_Process/README.md

# Part III: Compositional Agents Process

This folder contains the data synthesis pipeline used to construct **Part III compositional agents**. The pipeline converts existing Part I and Part II supervision into new agent configurations that combine backbone LLMs and tools.

The goal is to generate query-conditioned compositional agents in the form:

```json
{
  "M": {"name": "Backbone LLM Name"},
  "T": {"tools": ["tool_name_1", "tool_name_2"]}
}
```

Each generated sample contains the original question, retrieved evidence from Part I and Part II, relevant tool descriptions, the final prompt sent to the LLM, and the generated ranked agent configurations.

## What this pipeline does

The process has three steps.

1. **Build retrieval indexes**
   - Part I questions are embedded into a Chroma database.
   - Part II questions are embedded into a Chroma database.
   - Tool descriptions are embedded into a Chroma database.

2. **Retrieve cross-part supervision**
   - For a Part I query, the pipeline retrieves similar Part II queries and their toolkit-oriented agents.
   - For a Part II query, the pipeline retrieves similar Part I queries and their backbone-LLM-oriented agents.
   - The pipeline also retrieves tool descriptions using LLM-generated tool search queries.

3. **Synthesize compositional agents**
   - The retrieved backbone agents, toolkit agents, and tools are merged into a structured suggestion prompt.
   - An LLM generates five ranked compositional agents using only the model and tool names appearing in the suggestion.

## Expected dataset layout

By default, the scripts assume the dataset is stored at `../dataset`:

```text
../dataset/
  PartI/
    agents/merge.json
    questions/merge.json
    rankings/merge.json
  PartII/
    agents/merge.json
    questions/merge.json
    rankings/merge.json
    tools/merge.json
    tools/merged_tools_fill.json   # optional; used first if available
  PartI_vector_db/
  PartII_vector_db/
  Tool_vector_db/
```

The vector database folders are generated by the indexing script.

## Installation

Create a clean environment and install the dependencies:

```bash
pip install -r requirements.txt
```

The LLM synthesis step requires an OpenAI-compatible API key:

```bash
export OPENAI_API_KEY="your_api_key"
```

If you use another OpenAI-compatible endpoint, adjust `partiii_compositional_agents/llm_utils.py` accordingly.

## Usage

### 1. Build all indexes

```bash
python scripts/build_indexes.py \
  --dataset-root ../dataset \
  --device cuda:0 \
  --recreate
```

To build only one index:

```bash
python scripts/build_indexes.py --which part_i_questions
python scripts/build_indexes.py --which part_ii_questions
python scripts/build_indexes.py --which tools
```

### 2. Generate samples from Part I queries

```bash
python scripts/synthesize_part_i.py \
  --dataset-root ../dataset \
  --output-root ./outputs \
  --device cuda:0 \
  --sample-size 5
```

Output:

```text
outputs/simulated_results_from_PartI_rebuttal/main.jsonl
```

### 3. Generate samples from Part II queries

```bash
python scripts/synthesize_part_ii.py \
  --dataset-root ../dataset \
  --output-root ./outputs \
  --device cuda:0 \
  --sample-size 200
```

Output:

```text
outputs/simulated_results_from_PartII_rebuttal/main.jsonl
```

## Output format

Each line in the output JSONL file is one synthesized training example:

```json
{
  "dirname": "PartII",
  "num": 0,
  "key": "PartII_question_0",
  "question": "user question",
  "suggestion": "retrieved Part I/Part II/tool evidence",
  "prompt": "final generation prompt",
  "recommendation_agents": [
    {"M": {"name": "..."}, "T": {"tools": ["..."]}, "C": {}}
  ]
}
```

## Code structure

```text
partiii_compositional_agents/
  config.py          # central paths and runtime configuration
  embeddings.py      # OpenAI, MiniLM, and BGE-M3 embedding wrappers
  knowledge_base.py  # document loading and Chroma index construction
  search.py          # similarity-search wrapper
  llm_utils.py       # tool-query generation and agent-generation prompts
  synthesis.py       # Part I -> Part III and Part II -> Part III synthesis logic
scripts/
  build_indexes.py
  synthesize_part_i.py
  synthesize_part_ii.py
```

## Notes for reproducibility

- The default dense encoder is `BAAI/bge-m3`.
- Tool and question retrieval use Chroma similarity search with a default threshold of `0.5` and `top-k = 5`.
- The generated agents are constrained to use only backbone LLM names and tool names found in the retrieved suggestion.
- The configuration can be changed in `ProjectConfig` or through script arguments.