Skip to content
FractalAIResearchLabsPublic

About

Repository for DiagAgent paper.

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

3 Commits

Folders and files

Repository files navigation

VAIDYA

VAIDYA is an open benchmark suite, runner, and evaluator for simulated medical conversations. A simulated patient (LLM) describes symptoms and answers follow-up questions, while a doctor model (LLM) conducts a multi-turn consultation. VAIDYA then evaluates the doctor's behaviour against gold cases using two ported pipelines:

  • DiagBench: did the doctor reach the expected diagnosis? (0/1 judge).
  • MedConvBench: did the doctor ask the rubric questions? (per-rubric coverage, weighted by rubric score).

The package is intentionally focused on the DoctorSimulator workflow only. Any proprietary system-2 / agentic doctor pipelines are out of scope.


Features

  • Multi-turn DoctorSimulator vs. PatientSimulator consultations.
  • Strict, single input mode: a JSON list file in the published DiagBench.json / MedConvBench.json schema.
  • Pluggable LLM backend during conversation generation through any OpenAI-compatible chat-completion API (OpenAI, xAI/Grok, Together/DeepSeek, locally hosted Qwen/MedGemma, etc.) with token usage and per-conversation cost calculation from a configurable price table.
  • Evaluation pipelines ported verbatim from the official dataset_eval/eval_{diagbench,medconvbench}.py scripts. Uses the OpenAI Responses API (gpt-4.1 judge by default) and reports eval-side token usage.
  • Config-driven multi-model batch execution (vaidya run-many) with per-model output directories and JSONL run logs.

Installation

VAIDYA targets Python 3.10+.

git clone https://github.com/FractalAIResearchLabs/VAIDYA.git
cd VAIDYA
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

Copy the environment template and fill in credentials for the providers you plan to use:

cp .env.example .env
# edit .env and set OPENAI_API_KEY, XAI_API_KEY, TOGETHER_API_KEY, etc.

The repo ships the gold case files at data/DiagBench.json (1000 cases) and data/MedConvBench.json (500 cases with expected_questions rubrics).


Quickstart

# 1. Generate conversations against the gold cases.
vaidya run \
  --input data/DiagBench.json \
  --output-dir runs/gpt4o \
  --doctor-model gpt-4o

# 2. Score the conversations against the diagnosis (0/1 judge).
vaidya evaluate \
  --benchmark diagbench \
  --conversation-dir runs/gpt4o \
  --case-file data/DiagBench.json \
  --output-file reports/gpt4o_diagbench.json

# 3. Or score them against the question-coverage rubric.
vaidya evaluate \
  --benchmark medconvbench \
  --conversation-dir runs/gpt4o \
  --case-file data/MedConvBench.json \
  --output-file reports/gpt4o_medconvbench.json

# 4. Drive multiple models in one shot.
vaidya run-many --config configs/runs.example.yaml

Each batch run writes one conversation JSON per case (patient_<id>.json) plus a manifest.json debug artifact.


CLI Reference

vaidya run

Run one DoctorSimulator batch over a JSON list of cases.

Flag Purpose
--input PATH JSON list file (DiagBench.json / MedConvBench.json shape). Required.
--output-dir DIR Directory to write conversation JSONs into. Required.
--doctor-model NAME Doctor model identifier. Required.
--doctor-provider NAME Override provider for the doctor model. Inferred from the model name when omitted.
--patient-model NAME Patient model. Defaults to gpt-4o.
--patient-provider NAME Provider for the patient model. Defaults to openai.
--max-rounds N Maximum conversation rounds. Defaults to 8.
--pricing-file PATH JSON file with per-model token prices. When omitted, cost is reported as 0 and a sample table is shipped at configs/pricing.json.
--no-skip-existing Re-generate even when an output file already exists.

vaidya evaluate

Score generated conversations using the official DiagBench / MedConvBench prompt and pipeline. The judge runs against the OpenAI Responses API; the API key is read from OPENAI_API_KEY (loaded from .env at CLI entry).

Flag Purpose
--benchmark {diagbench,medconvbench} Which evaluation pipeline to run. Required.
--conversation-dir DIR Directory of generated conversation JSONs (patient_<id>.json). Required.
--case-file PATH Path to DiagBench.json or MedConvBench.json. Required.
--output-file PATH Where to write the evaluation JSON. Required.
--evaluator-model NAME Judge model for the Responses API. Defaults to $VAIDYA_EVALUATOR_MODEL, $DIAGBENCH_EVALUATOR_MODEL, or gpt-4.1.
--extraction-model NAME Extraction model for --extract-diagnosis-flag. Defaults to gpt-4.1-mini.
--num-threads N Number of judge threads. Defaults to 128.
--extract-diagnosis-flag DiagBench only. When set, run an extra LLM extraction step instead of taking the last Doctor turn directly. Off by default, mirrors the official script.

vaidya run-many

Drive multiple model runs from one YAML file.

vaidya run-many --config configs/runs.example.yaml

See configs/runs.example.yaml for the schema. Each entry under doctor_models: becomes its own subdirectory under output_root and its own JSONL log under log_dir. The YAML input: field must point at a JSON list file; directories and non-.json files are rejected at config load.


Input format

VAIDYA accepts exactly one input shape: a top-level JSON list of case objects in the published DiagBench / MedConvBench schema. Each entry must carry an integer patient_id, a case_study body, and (for DiagBench scoring) expected_diagnosis. MedConvBench cases additionally carry an expected_questions list of {question, score} rubric items.

[
  {
    "patient_id": 0,
    "age": null,
    "gender": "M",
    "chief_complaint": "Sore throat",
    "case_study": "Patient presents with Sore throat ...",
    "expected_diagnosis": "Strep throat",
    "expected_questions": [
      {"question": "Has there been a fever?", "score": 1}
    ],
    "source": "MTB"
  }
]

Directory inputs and single-case JSON object files are not supported; they are rejected with a CaseLoadError / ConfigError.


Output layout

Conversations (vaidya run)

One file per case, named patient_<patient_id>.json:

{
  "agent": "DoctorSimulator (gpt-4o)",
  "conversation": [
    {"Patient": "...", "Doctor": "..."}
  ],
  "candidate_diseases": ["..."],
  "disease_scores": {"Diagnosis A": 80.0},
  "token_usage": {
    "prompt_tokens": 5101,
    "completion_tokens": 584,
    "total_tokens": 6981
  },
  "total_cost": 0.0229,
  "metadata": {
    "case_id": "0",
    "doctor_model": "gpt-4o",
    "doctor_provider": "openai",
    "patient_model": "gpt-4o",
    "patient_provider": "openai",
    "created_at": "2026-05-07T06:30:00+00:00",
    "source_path": "data/DiagBench.json"
  }
}

A manifest.json is written alongside the conversations as a debug artifact (it is not consumed by the evaluator).

Evaluation (vaidya evaluate)

A single JSON file at --output-file with per-case results and an aggregate summary that includes eval-side token usage:

{
  "results": [
    {"patient_id": 0, "score": 1, "reasoning": "Diagnosis matches."}
  ],
  "summary": {
    "evaluated": 1000,
    "score_total": 742.0,
    "score_mean": 0.742,
    "eval_token_usage": {
      "prompt_tokens": 412034,
      "completion_tokens": 31021,
      "total_tokens": 443055,
      "prompt_cached_tokens": 0,
      "completion_reasoning_tokens": 0
    }
  }
}

For MedConvBench each per-case score is case_score / total_score over the rubric, where case_score sums the score values of rubric questions the doctor asked. The summary.score_mean divides by len(evaluated_cases) (no hardcoded denominator).


Providers and environment variables

VAIDYA treats every backend as an OpenAI-compatible chat completions API during conversation generation and configures it from environment variables (loaded from .env if present). Evaluation always goes through the OpenAI Responses API and uses OPENAI_API_KEY.

Provider name Env vars Default base URL
openai OPENAI_API_KEY, optional OPENAI_BASE_URL OpenAI default
xai XAI_API_KEY, optional XAI_BASE_URL https://api.x.ai/v1
together TOGETHER_API_KEY, optional TOGETHER_BASE_URL https://api.together.xyz/v1
local_openai optional LOCAL_OPENAI_API_KEY (defaults to not-required), optional LOCAL_OPENAI_BASE_URL http://localhost:8000/v1

When --doctor-provider is omitted, the provider is inferred from the model name (e.g. grok-* → xai, deepseek-* → together, qwen-* / medgemma-* / ii-medical-* → local_openai, otherwise openai).

To wire in a custom backend during conversation generation, register your own LLMProvider on an LLMProviderFactory and pass it to BenchmarkRunner.


Pricing

Token costs (for vaidya run) are computed from a JSON file with the shape:

{
  "models": {
    "gpt-4o": {
      "provider": "openai",
      "input_cost_per_million_tokens": 2.5,
      "output_cost_per_million_tokens": 10.0
    }
  }
}

Pass it via --pricing-file or pricing_file: in runs.example.yaml. When a model is not found in the table the per-conversation cost is reported as 0 without failing the run.


Project layout

VAIDYA/
  README.md
  pyproject.toml
  .env.example
  data/
    DiagBench.json           Gold cases: 1000 entries with expected_diagnosis.
    MedConvBench.json        Gold cases: 500 entries with expected_questions.
  configs/
    models.example.yaml      Model catalogue example.
    runs.example.yaml        Multi-model batch config example.
    pricing.json             Default per-model pricing.
  scripts/
    smoke_test.sh            Generate conversations on a small slice.
    smoke_eval.sh            Score the conversations from smoke_test.sh.
    run_all_evals.sh         Thin wrapper around `vaidya run-many`.
  src/diagbench/
    cli.py                   CLI entrypoints.
    config.py                YAML config loaders.
    schemas.py               Pydantic models for cases, turns, usage.
    data_loaders.py          Strict JSON-list loader.
    llm_providers.py         Provider abstraction and factory.
    pricing.py               Token cost calculation.
    patient_simulator.py     Simulated patient behavior.
    doctor_simulator.py      DoctorSimulator behavior and parsing.
    runner.py                Conversation runner and manifest writer.
    evaluator.py             Ported DiagBench / MedConvBench scoring pipelines.
    orchestration.py         Multi-run executor with logging.
  tests/                     Pytest suite (no network calls).

Development

source .venv/bin/activate
pip install -e ".[dev]"
python -m pytest
ruff check .

Tests use mocked providers and a FakeEvaluator so the suite runs without API keys.


Citation

@misc{vaidya2026,
  title={VAIDYA: Validated Agents for Intelligent Diagnosis and Yielded Analysis},
  author={Kalash Shah and Gautam Bhutani and Rohitaswa Sarbhangia and Snehan J},
  howpublished={\url{https://github.com/FractalAIResearchLabs/VAIDYA}},
  note={Accepted at ACL 2026 and AIME 2026. GitHub repository.},
  year={2026}
}

License

MIT. See LICENSE.

About

Repository for DiagAgent paper.

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages