Files
protocol-bicorder/analysis/scripts/bicorder_query.py
T
Protocolbot fb3bebcea0 fix: make --resume in batch pipeline actually skip completed work
Previously --resume re-queried every gradient in every row, overwriting
existing values — an interrupted run could not be resumed cheaply.

- bicorder_query.py: add --resume flag; skip gradients whose cells already
  have values, and report how many were skipped
- bicorder_batch.py: pass --resume through to query; skip fully-complete
  rows before invoking the query script; report partial rows
- bicorder_batch.py: import row/config helpers from bicorder_query instead
  of calling undefined names (would have crashed on --resume)
2026-09-23 07:58:04 -06:00

257 lines
8.4 KiB
Python

#!/usr/bin/env python3
"""
Query LLM for individual gradient values and update CSV.
This script generates prompts for each gradient, queries the LLM conversation,
and updates the CSV with the returned values.
"""
import csv
import json
import sys
import argparse
import subprocess
import re
from pathlib import Path
def load_bicorder_config(bicorder_path):
"""Load and parse the bicorder.json configuration file."""
with open(bicorder_path, 'r') as f:
return json.load(f)
def extract_gradients(bicorder_data):
"""Extract all gradients from the diagnostic sets."""
gradients = []
for diagnostic_set in bicorder_data['diagnostic']:
set_name = diagnostic_set['set_name']
for gradient in diagnostic_set['gradients']:
col_name = f"{set_name}_{gradient['term_left']}_vs_{gradient['term_right']}"
gradients.append({
'column_name': col_name,
'set_name': set_name,
'term_left': gradient['term_left'],
'term_left_description': gradient['term_left_description'],
'term_right': gradient['term_right'],
'term_right_description': gradient['term_right_description']
})
return gradients
def get_protocol_by_row(csv_path, row_number):
"""Get protocol data from CSV by row number (1-indexed)."""
with open(csv_path, 'r', encoding='utf-8') as f:
reader = csv.DictReader(f)
for i, row in enumerate(reader, start=1):
if i == row_number:
return {
'descriptor': row.get('Descriptor', '').strip(),
'description': row.get('Description', '').strip()
}
return None
def get_row_values(csv_path, row_number):
"""Get all existing values for a row (1-indexed) as a dict of column -> value."""
with open(csv_path, 'r', encoding='utf-8') as f:
reader = csv.DictReader(f)
for i, row in enumerate(reader, start=1):
if i == row_number:
return row
return None
def generate_gradient_prompt(protocol_descriptor, protocol_description, gradient):
"""Generate a prompt for a single gradient evaluation."""
return f"""Analyze this protocol: "{protocol_descriptor}"
Description: {protocol_description}
Evaluate the protocol on this gradient:
**{gradient['term_left']}** (1) vs **{gradient['term_right']}** (9)
- **{gradient['term_left']}**: {gradient['term_left_description']}
- **{gradient['term_right']}**: {gradient['term_right_description']}
Provide a rating from 1 to 9, where:
- 1 = strongly {gradient['term_left']}
- 5 = neutral/balanced/not applicable
- 9 = strongly {gradient['term_right']}
Respond with ONLY the number (1-9), optionally followed by a brief explanation.
"""
def query_llm(prompt, model=None):
"""Send prompt to llm CLI and get response."""
cmd = ['llm']
if model:
cmd.extend(['-m', model])
try:
result = subprocess.run(
cmd,
input=prompt,
text=True,
capture_output=True,
check=True
)
return result.stdout.strip()
except subprocess.CalledProcessError as e:
print(f" Error calling llm: {e.stderr}", file=sys.stderr)
return None
def extract_value(llm_response):
"""Extract numeric value (1-9) from LLM response."""
# Look for a number 1-9 at the start of the response
match = re.search(r'^(\d)', llm_response.strip())
if match:
value = int(match.group(1))
if 1 <= value <= 9:
return value
return None
def update_csv_cell(csv_path, row_number, column_name, value):
"""Update a specific cell in the CSV."""
# Read all rows
rows = []
with open(csv_path, 'r', encoding='utf-8') as f:
reader = csv.DictReader(f)
fieldnames = reader.fieldnames
for row in reader:
rows.append(row)
# Update the specific cell
if row_number <= len(rows):
rows[row_number - 1][column_name] = str(value)
# Write back
with open(csv_path, 'w', newline='', encoding='utf-8') as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
return True
return False
def main():
parser = argparse.ArgumentParser(
description='Query LLM for gradient values and update CSV',
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Example usage:
# Query all gradients for protocol in row 1
python3 bicorder_query.py analysis_output.csv 1
# Query specific model
python3 bicorder_query.py analysis_output.csv 1 -m mistral
# Dry run (show prompts without calling LLM)
python3 bicorder_query.py analysis_output.csv 1 --dry-run
"""
)
parser.add_argument('csv_path', help='CSV file to update')
parser.add_argument('row_number', type=int, help='Row number to analyze (1-indexed)')
parser.add_argument('-b', '--bicorder',
default='../bicorder.json',
help='Path to bicorder.json (default: ../bicorder.json)')
parser.add_argument('-m', '--model', help='LLM model to use')
parser.add_argument('--resume', action='store_true',
help='Skip gradients that already have values in the CSV')
parser.add_argument('--dry-run', action='store_true',
help='Show prompts without calling LLM or updating CSV')
args = parser.parse_args()
# Validate files exist
if not Path(args.csv_path).exists():
print(f"Error: CSV file '{args.csv_path}' not found", file=sys.stderr)
sys.exit(1)
if not Path(args.bicorder).exists():
print(f"Error: Bicorder config '{args.bicorder}' not found", file=sys.stderr)
sys.exit(1)
# Load protocol data
protocol = get_protocol_by_row(args.csv_path, args.row_number)
if protocol is None:
print(f"Error: Row {args.row_number} not found in CSV", file=sys.stderr)
sys.exit(1)
# Load bicorder config
bicorder_data = load_bicorder_config(args.bicorder)
gradients = extract_gradients(bicorder_data)
# Load existing values for this row (for resume mode)
existing_values = get_row_values(args.csv_path, args.row_number) or {}
# Count existing values for reporting
if args.resume:
already_filled = sum(1 for g in gradients if existing_values.get(g['column_name'], '').strip())
if already_filled:
print(f"Resume: {already_filled}/{len(gradients)} gradients already have values")
if args.dry_run:
print(f"DRY RUN: Row {args.row_number}, {len(gradients)} gradients")
print(f"Protocol: {protocol['descriptor']}\n")
else:
print(f"Protocol: {protocol['descriptor']}")
print(f"Loaded {len(gradients)} gradients, starting queries...")
# Process each gradient
for i, gradient in enumerate(gradients, 1):
gradient_short = gradient['column_name'].replace('_', ' ')
# In resume mode, skip gradients that already have a value
if args.resume and existing_values.get(gradient['column_name'], '').strip():
print(f"[{i}/{len(gradients)}] {gradient_short}: SKIP (already has value)")
continue
if not args.dry_run:
print(f"[{i}/{len(gradients)}] Querying: {gradient_short}...", flush=True)
# Generate prompt (including protocol context)
prompt = generate_gradient_prompt(
protocol['descriptor'],
protocol['description'],
gradient
)
if args.dry_run:
print(f"[{i}/{len(gradients)}] {gradient_short}")
print(f"Prompt:\n{prompt}\n")
continue
# Query LLM (new chat each time)
response = query_llm(prompt, args.model)
if response is None:
print(f"[{i}/{len(gradients)}] {gradient_short}: FAILED")
continue
# Extract value
value = extract_value(response)
if value is None:
print(f"[{i}/{len(gradients)}] {gradient_short}: WARNING - no valid value")
continue
# Update CSV
if update_csv_cell(args.csv_path, args.row_number, gradient['column_name'], value):
print(f"[{i}/{len(gradients)}] {gradient_short}: {value}")
else:
print(f"[{i}/{len(gradients)}] {gradient_short}: ERROR updating CSV")
if not args.dry_run:
print(f"\n✓ CSV updated: {args.csv_path}")
if __name__ == '__main__':
main()