{"version":1,"extractor":"496c9d83e6eb3668","files":{"download_datasets.py":{"size":4817,"mtimeMs":1789362123397.8032,"hash":"306c3ac9035cf9e7a512d4d9db84b76b305e763d75e5545743d2e2c9d77736f8","nodes":[{"id":"download_datasets.py","name":"download_datasets.py","kind":"file","path":"download_datasets.py","span":"L1-L132","signature":null,"exported":true,"origin":"ast","body_hash":"306c3ac9035cf9e7a512d4d9db84b76b305e763d75e5545743d2e2c9d77736f8","chars":4817,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"This file is used to download necessary raw datasets for the project.\" import os import argparse import shutil from datasets import load_dataset from src.paths import RAW_DATA_DIR DATASET_NAMES = [\"math500\", \"mmlu_pro\", \"truthfulqa\"] if __name__ == \"__main__\": parser = argparse.ArgumentParser() parser.add_argument(\"--dataset_name\", type=str, required=True, choices=DATASET_NAMES) parser.add_argument(\"--dataset_dir\", type=str, default=RAW_DATA_DIR) args = parser.parse_args() dataset_target_dir = args.dataset_dir if not os.path.exists(dataset_target_dir): os.makedirs(dataset_target_dir) if args.dataset_name == \"math500\": if not os.path.exists(os.path.join(dataset_target_dir, \"MATH500\")): print(f\"Downloading MATH500 dataset to {dataset_target_dir}...\") math500(dataset_target_dir) else: print(f\"MATH500 dataset already exists in {dataset_target_dir}\") elif args.dataset_name == \"mmlu_pro\": if not os.path.exists(os.path.join(dataset_target_dir, \"MMLU-PRO\")): print(f\"Downloading MMLU-Pro dataset to {dataset_target_dir}...\") mmlu_pro(dataset_target_dir) else: print(f\"MMLU-Pro dataset already exists in {dataset_target_dir}\") elif args.dataset_name == \"truthfulqa\": if not os.path.exists(os.path.join(dataset_target_dir, \"TruthfulQA\")): print(f\"Downloading TruthfulQA dataset to {dataset_target_dir}...\") truthfulqa(dataset_target_dir) else: print(f\"TruthfulQA dataset already exists in {dataset_target_dir}\")"},{"id":"download_datasets.py#math500","name":"math500","kind":"function","path":"download_datasets.py","span":"L13-L39","signature":"def math500(target_dir:str)","exported":true,"origin":"ast","body_hash":"0928e0f1fa994fdc6a121ce658b407f8133824f792419b7b63554082a471ba9f","body_text":"def math500(target_dir:str): \"\"\"Download MATH500 dataset to the target directory. Args: target_dir: str, the directory to save the dataset. Returns: None \"\"\" dataset_id = \"HuggingFaceH4/MATH-500\" target_dataset_dir = os.path.join(target_dir, \"MATH500\") if not os.path.exists(target_dataset_dir): os.makedirs(target_dataset_dir) dataset = load_dataset(\"HuggingFaceH4/MATH-500\", \"default\", cache_dir=target_dir) test_set = dataset[\"test\"] print(f\"DATASET INFO: {len(test_set)} test samples\") test_set.to_json(os.path.join(target_dataset_dir, \"test.json\")) tmp_dir_name = f\"{dataset_id.split('/')[0]}___{dataset_id.split('/')[1].lower()}\" shutil.rmtree(os.path.join(target_dir, tmp_dir_name)) lockfiles = [file for file in os.listdir(target_dir) if file.endswith(\".lock\")] for lockfile in lockfiles: os.remove(os.path.join(target_dir, lockfile)) print(f\"MATH500 dataset downloaded to {target_dataset_dir}\")","summary_state":"pending","summary":null,"crux":null},{"id":"download_datasets.py#mmlu_pro","name":"mmlu_pro","kind":"function","path":"download_datasets.py","span":"L42-L68","signature":"def mmlu_pro(target_dir:str)","exported":true,"origin":"ast","body_hash":"bedecbdd4c712fac16276620c63fd14ca8c2d7e3615f3a52099179bab29fd88e","body_text":"def mmlu_pro(target_dir:str): \"\"\"Download MMLU-Pro dataset to the target directory. Args: target_dir: str, the directory to save the dataset. Returns: None \"\"\" dataset_id = \"TIGER-Lab/MMLU-Pro\" target_dataset_dir = os.path.join(target_dir, \"MMLU-PRO\") if not os.path.exists(target_dataset_dir): os.makedirs(target_dataset_dir) dataset = load_dataset(\"TIGER-Lab/MMLU-Pro\", \"default\", cache_dir=target_dir) test_set = dataset[\"test\"] print(f\"DATASET INFO: {len(test_set)} test samples\") test_set.to_json(os.path.join(target_dataset_dir, \"mmlu_pro_test.json\")) tmp_dir_name = f\"{dataset_id.split('/')[0]}___{dataset_id.split('/')[1].lower()}\" shutil.rmtree(os.path.join(target_dir, tmp_dir_name)) lockfiles = [file for file in os.listdir(target_dir) if file.endswith(\".lock\")] for lockfile in lockfiles: os.remove(os.path.join(target_dir, lockfile)) print(f\"MMLU-Pro dataset downloaded to {target_dataset_dir}\")","summary_state":"pending","summary":null,"crux":null},{"id":"download_datasets.py#truthfulqa","name":"truthfulqa","kind":"function","path":"download_datasets.py","span":"L71-L98","signature":"def truthfulqa(target_dir:str)","exported":true,"origin":"ast","body_hash":"cb51d834968692313d45bff3a9ecb912411ea42aa1031168df99f89ca1fe2257","body_text":"def truthfulqa(target_dir:str): \"\"\"Download TruthfulQA dataset to the target directory. Args: target_dir: str, the directory to save the dataset. Returns: None \"\"\" dataset_id = \"truthfulqa/truthful_qa\" target_dataset_dir = os.path.join(target_dir, \"TruthfulQA\") if not os.path.exists(target_dataset_dir): os.makedirs(target_dataset_dir) for config in [\"generation\", \"multiple_choice\"]: dataset = load_dataset(dataset_id, config, cache_dir=target_dir) validation_set = dataset[\"validation\"] print(f\"DATASET ({config}) INFO: {len(validation_set)} validation samples\") validation_set.to_json(os.path.join(target_dataset_dir, f\"{config}_validation.json\")) tmp_dir_name = f\"{dataset_id.split('/')[0]}___{dataset_id.split('/')[1].lower()}\" shutil.rmtree(os.path.join(target_dir, tmp_dir_name)) lockfiles = [file for file in os.listdir(target_dir) if file.endswith(\".lock\")] for lockfile in lockfiles: os.remove(os.path.join(target_dir, lockfile)) print(f\"TruthfulQA dataset downloaded to {target_dataset_dir}\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"download_datasets.py","relation":"imports","specifier":"os","file":"download_datasets.py"},{"source":"download_datasets.py","relation":"imports","specifier":"argparse","file":"download_datasets.py"},{"source":"download_datasets.py","relation":"imports","specifier":"shutil","file":"download_datasets.py"},{"source":"download_datasets.py","relation":"imports","specifier":"datasets","file":"download_datasets.py"},{"source":"download_datasets.py","relation":"imports","specifier":"src.paths","file":"download_datasets.py"},{"source":"download_datasets.py","relation":"contains","targetId":"download_datasets.py#math500","file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"makedirs","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"load_dataset","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"len","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"to_json","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"lower","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"rmtree","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"listdir","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"endswith","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"remove","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#math500","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"contains","targetId":"download_datasets.py#mmlu_pro","file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"makedirs","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"load_dataset","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"len","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"to_json","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"lower","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"rmtree","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"listdir","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"endswith","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"remove","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#mmlu_pro","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"contains","targetId":"download_datasets.py#truthfulqa","file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"makedirs","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"load_dataset","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"len","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"to_json","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"lower","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"split","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"rmtree","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"listdir","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"endswith","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"remove","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py#truthfulqa","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"add_argument","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"add_argument","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"parse_args","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"makedirs","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"math500","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"mmlu_pro","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"exists","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"join","viaMember":true,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"truthfulqa","viaMember":false,"file":"download_datasets.py"},{"source":"download_datasets.py","relation":"calls","name":"print","viaMember":false,"file":"download_datasets.py"}]},"prepare_data.py":{"size":5232,"mtimeMs":1789362123397.8032,"hash":"8df30903e2d8211fde021a440b977374ed8cfbbe49c087c598bd9afc27f37ea2","nodes":[{"id":"prepare_data.py","name":"prepare_data.py","kind":"file","path":"prepare_data.py","span":"L1-L145","signature":null,"exported":true,"origin":"ast","body_hash":"8df30903e2d8211fde021a440b977374ed8cfbbe49c087c598bd9afc27f37ea2","chars":5232,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import json import argparse import random from src.paths import RAW_DATA_DIR, PROCESSED_DATA_DIR DATASET_NAMES = [\"math500\", \"mmlu_pro_engineering\", \"mmlu_pro_economics\", \"truthfulqa\"] if __name__ == \"__main__\": parser = argparse.ArgumentParser() parser.add_argument(\"--dataset_name\", type=str, required=True, choices=DATASET_NAMES) parser.add_argument(\"--seed\", type=int, default=42) args = parser.parse_args() if args.dataset_name == \"math500\": math500_processor() print(f\"Processed {args.dataset_name} dataset\") elif args.dataset_name == \"mmlu_pro_engineering\": mmlu_pro_engineering_processor(seed=args.seed) print(f\"Processed {args.dataset_name} dataset\") elif args.dataset_name == \"mmlu_pro_economics\": mmlu_pro_economics_processor(seed=args.seed) print(f\"Processed {args.dataset_name} dataset\") elif args.dataset_name == \"truthfulqa\": truthfulqa_processor(seed=args.seed) print(f\"Processed {args.dataset_name} dataset\")"},{"id":"prepare_data.py#load_data_from_json","name":"load_data_from_json","kind":"function","path":"prepare_data.py","span":"L12-L14","signature":"def load_data_from_json(file_path:str)","exported":true,"origin":"ast","body_hash":"3abc58a21df308e7d2758c03c887ab09c37ca017eccbb9035dad375e61e7cfd5","body_text":"def load_data_from_json(file_path:str): with open(file_path, \"r\", encoding=\"utf-8\") as f: return [json.loads(line) for line in f]","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#dump_data_to_json","name":"dump_data_to_json","kind":"function","path":"prepare_data.py","span":"L17-L21","signature":"def dump_data_to_json(data:list, file_path:str)","exported":true,"origin":"ast","body_hash":"32de26f3d487b7e3aa3037eacb0e2314a09d914ea4dbf10879f606fd3fb1e628","body_text":"def dump_data_to_json(data:list, file_path:str): with open(file_path, \"w\", encoding=\"utf-8\") as f: for item in data: json_line = json.dumps(item, ensure_ascii=False) f.write(json_line + \"\\n\")","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#math500_processor","name":"math500_processor","kind":"function","path":"prepare_data.py","span":"L24-L46","signature":"def math500_processor()","exported":true,"origin":"ast","body_hash":"309bc860ddab866a47670cee86be7afaae712af9e20c59dcf4151d7c02d905e4","body_text":"def math500_processor(): if not os.path.exists(os.path.join(PROCESSED_DATA_DIR, \"math500\")): os.makedirs(os.path.join(PROCESSED_DATA_DIR, \"math500\")) source_path = os.path.join(RAW_DATA_DIR, \"MATH500\", \"test.json\") train_data, test_data = [], [] with open(source_path, \"r\") as f: for line in f: entry = json.loads(line) sample = { \"question\": entry[\"problem\"], \"answer\": entry[\"answer\"], \"solution\": entry[\"solution\"], \"category\": entry[\"subject\"], } if entry.get(\"level\") == 5: test_data.append(sample) else: train_data.append(sample) dump_data_to_json(train_data, os.path.join(PROCESSED_DATA_DIR, \"math500\", \"train.jsonl\")) dump_data_to_json(test_data, os.path.join(PROCESSED_DATA_DIR, \"math500\", \"test.jsonl\"))","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#_mmlu_pro_subset_processor","name":"_mmlu_pro_subset_processor","kind":"function","path":"prepare_data.py","span":"L49-L77","signature":"def _mmlu_pro_subset_processor(category: str, seed: int = 42)","exported":false,"origin":"ast","body_hash":"f2187cb7893d0b4bacb5f230af0990f209f9d4d06a0b1abcc3fbbdb3ecd02766","body_text":"def _mmlu_pro_subset_processor(category: str, seed: int = 42): dataset_path = os.path.join(RAW_DATA_DIR, \"MMLU-PRO\", \"mmlu_pro_test.json\") subset_data = [] with open(dataset_path, \"r\") as f: for line in f: entry = json.loads(line) if entry.get(\"category\", \"\").lower() != category.lower(): continue options = entry[\"options\"] option_keys = [chr(ord(\"A\") + i) for i in range(len(options))] sample = { \"question\": entry[\"question\"], \"options\": dict(zip(option_keys, options)), \"answer\": entry[\"answer\"], \"category\": entry[\"category\"], } subset_data.append(sample) random.seed(seed) random.shuffle(subset_data) split = int(len(subset_data) * 0.75) train_data, test_data = subset_data[:split], subset_data[split:] out_dir = os.path.join(PROCESSED_DATA_DIR, \"mmlu_pro\", category.lower()) if not os.path.exists(out_dir): os.makedirs(out_dir) dump_data_to_json(train_data, os.path.join(out_dir, \"train.jsonl\")) dump_data_to_json(test_data, os.path.join(out_dir, \"test.jsonl\"))","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#mmlu_pro_engineering_processor","name":"mmlu_pro_engineering_processor","kind":"function","path":"prepare_data.py","span":"L80-L81","signature":"def mmlu_pro_engineering_processor(seed: int = 42)","exported":true,"origin":"ast","body_hash":"70f802dc9568998504d5617612c547929fdc9f4d09fa83e5245c4e5a7c1c4d02","body_text":"def mmlu_pro_engineering_processor(seed: int = 42): _mmlu_pro_subset_processor(\"engineering\", seed=seed)","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#mmlu_pro_economics_processor","name":"mmlu_pro_economics_processor","kind":"function","path":"prepare_data.py","span":"L84-L85","signature":"def mmlu_pro_economics_processor(seed: int = 42)","exported":true,"origin":"ast","body_hash":"ab1ac7fed3101180e542d20e4112a547c034bdcb1796749465b1db7fe03950a1","body_text":"def mmlu_pro_economics_processor(seed: int = 42): _mmlu_pro_subset_processor(\"economics\", seed=seed)","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_data.py#truthfulqa_processor","name":"truthfulqa_processor","kind":"function","path":"prepare_data.py","span":"L88-L121","signature":"def truthfulqa_processor(seed: int = 42)","exported":true,"origin":"ast","body_hash":"ef329197af5f1cf2c64b105f97a8b4d2e866dee93ea50a6709a88b58ae7c0add","body_text":"def truthfulqa_processor(seed: int = 42): if not os.path.exists(os.path.join(PROCESSED_DATA_DIR, \"truthfulqa\")): os.makedirs(os.path.join(PROCESSED_DATA_DIR, \"truthfulqa\")) source_path = os.path.join(RAW_DATA_DIR, \"TruthfulQA\", \"multiple_choice_validation.json\") random.seed(seed) data = [] with open(source_path, \"r\", encoding=\"utf-8\") as f: for line in f: entry = json.loads(line) choices = entry[\"mc1_targets\"][\"choices\"] labels = entry[\"mc1_targets\"][\"labels\"] if not (4 <= len(choices) <= 9): continue correct_choice = choices[labels.index(1)] shuffled_choices = choices[:] random.shuffle(shuffled_choices) option_keys = [chr(ord(\"A\") + i) for i in range(len(shuffled_choices))] answer = option_keys[shuffled_choices.index(correct_choice)] sample = { \"question\": entry[\"question\"], \"options\": dict(zip(option_keys, shuffled_choices)), \"answer\": answer, \"category\": \"truthfulqa\", } data.append(sample) random.shuffle(data) split = int(len(data) * 0.75) train_data, test_data = data[:split], data[split:] dump_data_to_json(train_data, os.path.join(PROCESSED_DATA_DIR, \"truthfulqa\", \"train.jsonl\")) dump_data_to_json(test_data, os.path.join(PROCESSED_DATA_DIR, \"truthfulqa\", \"test.jsonl\"))","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"prepare_data.py","relation":"imports","specifier":"os","file":"prepare_data.py"},{"source":"prepare_data.py","relation":"imports","specifier":"json","file":"prepare_data.py"},{"source":"prepare_data.py","relation":"imports","specifier":"argparse","file":"prepare_data.py"},{"source":"prepare_data.py","relation":"imports","specifier":"random","file":"prepare_data.py"},{"source":"prepare_data.py","relation":"imports","specifier":"src.paths","file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#load_data_from_json","file":"prepare_data.py"},{"source":"prepare_data.py#load_data_from_json","relation":"calls","name":"open","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#load_data_from_json","relation":"calls","name":"loads","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#dump_data_to_json","file":"prepare_data.py"},{"source":"prepare_data.py#dump_data_to_json","relation":"calls","name":"open","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#dump_data_to_json","relation":"calls","name":"dumps","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#dump_data_to_json","relation":"calls","name":"write","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#math500_processor","file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"exists","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"makedirs","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"open","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"loads","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"get","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"append","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"append","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#math500_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#_mmlu_pro_subset_processor","file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"open","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"loads","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"lower","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"get","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"lower","viaMember":true,"file":"prepare_data.py","recvType":"str"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"chr","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"ord","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"range","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"len","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"dict","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"zip","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"append","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"seed","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"shuffle","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"int","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"len","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"lower","viaMember":true,"file":"prepare_data.py","recvType":"str"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"exists","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"makedirs","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#_mmlu_pro_subset_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#mmlu_pro_engineering_processor","file":"prepare_data.py"},{"source":"prepare_data.py#mmlu_pro_engineering_processor","relation":"calls","name":"_mmlu_pro_subset_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#mmlu_pro_economics_processor","file":"prepare_data.py"},{"source":"prepare_data.py#mmlu_pro_economics_processor","relation":"calls","name":"_mmlu_pro_subset_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"contains","targetId":"prepare_data.py#truthfulqa_processor","file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"exists","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"makedirs","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"seed","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"open","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"loads","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"len","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"index","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"shuffle","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"chr","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"ord","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"range","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"len","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"index","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"dict","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"zip","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"append","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"shuffle","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"int","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"len","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"dump_data_to_json","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py#truthfulqa_processor","relation":"calls","name":"join","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"parse_args","viaMember":true,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"math500_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"print","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"mmlu_pro_engineering_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"print","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"mmlu_pro_economics_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"print","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"truthfulqa_processor","viaMember":false,"file":"prepare_data.py"},{"source":"prepare_data.py","relation":"calls","name":"print","viaMember":false,"file":"prepare_data.py"}]},"prepare_memory.py":{"size":5541,"mtimeMs":1789362123397.8032,"hash":"ff2bdde66b2e406ccca383d406a692c1358e363d1964d8660bc63a1df7e5b137","nodes":[{"id":"prepare_memory.py","name":"prepare_memory.py","kind":"file","path":"prepare_memory.py","span":"L1-L141","signature":null,"exported":true,"origin":"ast","body_hash":"ff2bdde66b2e406ccca383d406a692c1358e363d1964d8660bc63a1df7e5b137","chars":5541,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import argparse import chromadb import numpy as np from src.memory import MemoryVectorDB from src.models import EmbeddingModel from src.paths import MEMORY_DATA_DIR, PERSISTENT_DIR from src.utils import ( EmbeddingConfig, load_configs_from_yaml, read_memory_data, sample_memories, construct_memory, construct_response_memory, ) if __name__ == \"__main__\": main()"},{"id":"prepare_memory.py#parse_args","name":"parse_args","kind":"function","path":"prepare_memory.py","span":"L19-L36","signature":"def parse_args()","exported":true,"origin":"ast","body_hash":"89c95963e15c13612fafd87e62bf9c4ac3f5db78c98a21240b72b8dd9da07e2d","body_text":"def parse_args(): parser = argparse.ArgumentParser(description=\"Prepare Memory DB\") # data settings parser.add_argument(\"--task\", type=str, required=True, help=\"Task subdirectory under MEMORY_DATA_DIR\") parser.add_argument(\"--model_name\", type=str, required=True, help=\"Model subdirectory under MEMORY_DATA_DIR/<task>\") parser.add_argument(\"--collection_name\", type=str, default=\"\", help=\"Optional suffix for the ChromaDB collection name\") parser.add_argument(\"--batch_size\", type=int, default=256, help=\"Number of records per batch when adding to ChromaDB\") # embedding settings parser.add_argument(\"--embedding_model\", type=str, default=\"bge-m3\", help=\"SentenceTransformer model name or path\") # sample settings parser.add_argument(\"--sample_size\", type=int, default=-1, help=\"Number of unique questions to sample (-1 = all)\") parser.add_argument(\"--seed\", type=int, default=42, help=\"Random seed for sampling\") args = parser.parse_args() return args","summary_state":"pending","summary":null,"crux":null},{"id":"prepare_memory.py#main","name":"main","kind":"function","path":"prepare_memory.py","span":"L39-L136","signature":"def main()","exported":true,"origin":"ast","body_hash":"b9c61b40d51f9db7216d57a290507a4218f44ea82833f91c324639da0cdc8fea","body_text":"def main(): args = parse_args() # load embedding configurations and model configs = load_configs_from_yaml(\"configs.yaml\") embedding_config = EmbeddingConfig(configs[\"embedding_configs\"][args.embedding_model]) embedding_model = EmbeddingModel(embedding_config) if args.collection_name: collection_name = f\"{args.task}_{args.model_name}_{args.collection_name}\" else: collection_name = f\"{args.task}_{args.model_name}\" print(f\"Main Collection name: {collection_name}\") # Reset collections if they already exist client = chromadb.PersistentClient(path=str(PERSISTENT_DIR)) collection_names = client.list_collections() for name in [collection_name, f\"{collection_name}_responses\"]: if name in collection_names: client.delete_collection(name) print(f\"Deleted existing collection '{name}'\") # initialize memory database memory_db = MemoryVectorDB( persistent_dir=str(PERSISTENT_DIR), collection_name=collection_name, embedding_model=embedding_model, verbose=True, ) # load memory data from JSONL files data_dir = os.path.join(MEMORY_DATA_DIR, args.task, args.model_name) records = read_memory_data(data_dir) if not records: print(\"No memory data found. Exiting.\") return print(f\"Loaded {len(records)} memory cases from {data_dir}\") # ------------------------------------------------------------------ # 1. Create memory entries for debate state retrieval # ------------------------------------------------------------------ records_sample = sample_memories(records, sample_size=args.sample_size, seed=args.seed) print(f\"Sampled {len(records_sample)} memory cases.\") all_db_ids = [] all_value_texts = [] batch_size = args.batch_size for batch_start in range(0, len(records_sample), batch_size): batch = records_sample[batch_start: batch_start + batch_size] db_ids, embedding_texts, value_texts, meta_datas = zip(*[construct_memory(row) for row in batch]) print(f\"Adding batch {batch_start // batch_size + 1} ({len(batch)} records)...\") memory_db.add_batch_memories( db_ids=list(db_ids), embedding_texts=list(embedding_texts), document_texts=list(value_texts), add_meta_datas=list(meta_datas), ) all_db_ids.extend(db_ids) all_value_texts.extend(value_texts) print(f\"\\nDone. Total memory entries in collection: {memory_db.get_memories_count()}\") # Pre-compute and save document embeddings for fast MMR re-ranking at inference time. print(\"\\nPre-computing document embeddings...\") doc_embeddings = embedding_model.encode(all_value_texts, show_progress_bar=True) doc_embeddings = np.array(doc_embeddings, dtype=np.float16) npz_path = os.path.join(PERSISTENT_DIR, f\"{collection_name}_doc_embeddings.npz\") np.savez(npz_path, db_ids=np.array(all_db_ids), embeddings=doc_embeddings) print(f\"Saved document embeddings to {npz_path} (shape: {doc_embeddings.shape})\") # ------------------------------------------------------------------ # 2. Create response entries for confidence estimation # ------------------------------------------------------------------ seen_ids = set() resp_ids, resp_texts, resp_metas = [], [], [] for row in records: if not row.get(\"current_response\"): continue db_id, text, meta = construct_response_memory(row) if db_id not in seen_ids: resp_ids.append(db_id) resp_texts.append(text) resp_metas.append(meta) seen_ids.add(db_id) for batch_start in range(0, len(resp_ids), batch_size): memory_db.add_batch_responses( db_ids=resp_ids[batch_start: batch_start + batch_size], texts=resp_texts[batch_start: batch_start + batch_size], meta_datas=resp_metas[batch_start: batch_start + batch_size], ) print(f\"\\nDone. Total response entries in collection: {len(resp_ids)}\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"prepare_memory.py","relation":"imports","specifier":"os","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"argparse","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"chromadb","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"src.memory","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"src.models","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"src.paths","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"imports","specifier":"src.utils","file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"contains","targetId":"prepare_memory.py#parse_args","file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#parse_args","relation":"calls","name":"parse_args","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"contains","targetId":"prepare_memory.py#main","file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"parse_args","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"load_configs_from_yaml","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"EmbeddingConfig","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"EmbeddingModel","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"PersistentClient","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"str","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"list_collections","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"delete_collection","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"MemoryVectorDB","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"str","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"join","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"read_memory_data","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"sample_memories","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"range","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"zip","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"construct_memory","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"add_batch_memories","viaMember":true,"file":"prepare_memory.py","recvType":"MemoryVectorDB"},{"source":"prepare_memory.py#main","relation":"calls","name":"list","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"list","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"list","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"list","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"extend","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"extend","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"get_memories_count","viaMember":true,"file":"prepare_memory.py","recvType":"MemoryVectorDB"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"encode","viaMember":true,"file":"prepare_memory.py","recvType":"EmbeddingModel"},{"source":"prepare_memory.py#main","relation":"calls","name":"array","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"join","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"savez","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"array","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"set","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"get","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"construct_response_memory","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"append","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"append","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"append","viaMember":true,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"add","viaMember":true,"file":"prepare_memory.py","recvType":"set"},{"source":"prepare_memory.py#main","relation":"calls","name":"range","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"add_batch_responses","viaMember":true,"file":"prepare_memory.py","recvType":"MemoryVectorDB"},{"source":"prepare_memory.py#main","relation":"calls","name":"print","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py#main","relation":"calls","name":"len","viaMember":false,"file":"prepare_memory.py"},{"source":"prepare_memory.py","relation":"calls","name":"main","viaMember":false,"file":"prepare_memory.py"}]},"run_cot.py":{"size":9525,"mtimeMs":1789362123397.8032,"hash":"9a6571beaf492b10ca10c5923f4a4eb290f860d28b96f3cc6e1128694e7dccc4","nodes":[{"id":"run_cot.py","name":"run_cot.py","kind":"file","path":"run_cot.py","span":"L1-L220","signature":null,"exported":true,"origin":"ast","body_hash":"9a6571beaf492b10ca10c5923f4a4eb290f860d28b96f3cc6e1128694e7dccc4","chars":9525,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import json import argparse from datetime import datetime from typing import Dict, List os.environ[\"VLLM_WORKER_MULTIPROC_METHOD\"] = \"spawn\" from src.agent import Agent, AgentId, Query from src.memory import MemoryVectorDB from src.models import build_language_model, EmbeddingModel from src.chain_of_thought import ChainOfThought from src.paths import PERSISTENT_DIR, LOG_DIR from src.utils import LLMConfig, EmbeddingConfig, load_configs_from_yaml, load_env_file, DataLoader if __name__ == \"__main__\": main()"},{"id":"run_cot.py#parse_args","name":"parse_args","kind":"function","path":"run_cot.py","span":"L17-L52","signature":"def parse_args()","exported":true,"origin":"ast","body_hash":"45a049ac44e528576ea82282ff8c5cf68024f2b77a361a58c8e297e8b7e69307","body_text":"def parse_args(): parser = argparse.ArgumentParser(description=\"Run Chain-of-Thought / Self-Consistency evaluation\") # self-consistency settings parser.add_argument(\"--task\", type=str, default=\"MATH500\", help=\"Task / dataset type\") parser.add_argument(\"--self_consistency\", action=\"store_true\", help=\"Enable self-consistency with multi-path sampling\") parser.add_argument(\"--num_reasoning_paths\", type=int, default=9, help=\"Number of reasoning paths for self-consistency\") # agent settings parser.add_argument(\"--model_name\", type=str, default=\"qwen3-8b\", help=\"The name of adopted LLM\") parser.add_argument(\"--if_use_memory\", action=\"store_true\", help=\"Retrieve past cases as in-context examples\") parser.add_argument(\"--memory_model\", type=str, default=\"\", help=\"Model name whose memory collection to use (default: same as --model_name)\") parser.add_argument(\"--collection_name\", type=str, default=\"\", help=\"Optional suffix for the ChromaDB collection name\") parser.add_argument(\"--embedding_model\", type=str, default=\"bge-m3\", help=\"The name of adopted embedding model\") parser.add_argument(\"--n_retrieval\", type=int, default=1, help=\"Number of retrieved examples\") parser.add_argument(\"--retrieval_policy\", type=str, default=\"similarity\", choices=[\"similarity\", \"random\", \"diversity\", \"positive_only\"], help=\"Memory retrieval policy (only active with --if_use_memory; 'similarity' is the natural default for CoT)\") parser.add_argument(\"--seed\", type=int, default=None, help=\"Base seed for stochastic retrieval policies (random/diversity)\") parser.add_argument(\"--memory_agent_id\", type=str, default=\"Debater-0\", help=\"agent_id whose memory to retrieve from (memory is stored per debate agent)\") # general settings parser.add_argument(\"--exp_name\", type=str, default=\"test\", help=\"Optional prefix for the output filename\") parser.add_argument(\"--run_suffix\", type=str, default=\"\", help=\"Optional suffix appended before .json\") parser.add_argument(\"--save_log\", action=\"store_true\", help=\"Whether to save evaluation log to disk\") parser.add_argument(\"--train\", action=\"store_true\", help=\"Use training split instead of test split\") parser.add_argument(\"--batch_size\", type=int, default=64, help=\"Number of queries per batch\") parser.add_argument(\"--limit\", type=int, default=None, help=\"Cap the number of questions (for quick tests)\") # gpu settings parser.add_argument(\"--gpu_id\", type=str, default=\"0\", help=\"The GPU ID.\") parser.add_argument(\"--embed_gpu_id\", type=str, default=None, help=\"Dedicated GPU for the embedding model. Default None = share --gpu_id.\") parser.add_argument(\"--parallel\", action=\"store_true\", default=False, help=\"Whether to use parallel mode.\") args = parser.parse_args() return args","summary_state":"pending","summary":null,"crux":null},{"id":"run_cot.py#build_agent","name":"build_agent","kind":"function","path":"run_cot.py","span":"L55-L70","signature":"def build_agent(args, memory_db=None) -> Agent","exported":true,"origin":"ast","body_hash":"58c1b7a1420e18a0da7e3b24c9648aeeae3a43778033d95418d3f62b3ae047ec","body_text":"def build_agent(args, memory_db=None) -> Agent: agent_id = AgentId(agent_type=\"CoT\", agent_key=\"0\") agent = Agent( agent_id=agent_id, task_type=args.task, if_use_memory=args.if_use_memory, memory_db=memory_db, n_retrieval=args.n_retrieval, retrieval_policy=args.retrieval_policy, seed=args.seed, # CoT has no debate rounds; retrieve on round 0 and reuse a debate # agent's memory (memory is stored per Debater-{i}). retrieve_first_round=True, memory_agent_id=args.memory_agent_id, ) return agent","summary_state":"pending","summary":null,"crux":null},{"id":"run_cot.py#load_questions","name":"load_questions","kind":"function","path":"run_cot.py","span":"L73-L92","signature":"def load_questions(task_type: str, limit: int | None, is_training: bool) -> List[Query]","exported":true,"origin":"ast","body_hash":"2d7d6158ecee70f4dfa08f46aea5026691c9422a73f806ef823c79b39db7a1f5","body_text":"def load_questions(task_type: str, limit: int | None, is_training: bool) -> List[Query]: if is_training: raw_data = DataLoader.load_dataset(task_type + \"_TRAIN\") else: raw_data = DataLoader.load_dataset(task_type) queries = [] for idx, item in enumerate(raw_data): if limit is not None and idx >= limit: break content, answer, solution, category = DataLoader.format_question(item, task_type) queries.append(Query( content=content, query_id=idx, answer=answer, solution=solution, category=category, )) return queries","summary_state":"pending","summary":null,"crux":null},{"id":"run_cot.py#compute_metrics","name":"compute_metrics","kind":"function","path":"run_cot.py","span":"L95-L99","signature":"def compute_metrics(results: List[Dict]) -> Dict","exported":true,"origin":"ast","body_hash":"9960de980eb147c12b4c08466b9d816ceee86619f0ae459193a470768e6dfbc1","body_text":"def compute_metrics(results: List[Dict]) -> Dict: total = len(results) correct = sum(1 for r in results if r.get(\"is_final_correct\", False)) accuracy = correct / total if total > 0 else 0.0 return {\"correct\": correct, \"total\": total, \"accuracy\": accuracy}","summary_state":"pending","summary":null,"crux":null},{"id":"run_cot.py#main","name":"main","kind":"function","path":"run_cot.py","span":"L102-L215","signature":"def main()","exported":true,"origin":"ast","body_hash":"6085edd89846996c42bcb8764ea82f4be79a8290f62acb739fa8f0301382e03f","body_text":"def main(): args = parse_args() # Load API keys etc. from .env (no-op if the file is absent) load_env_file() # GPU settings gpus = args.gpu_id.split(\",\") tensor_parallel_size = len(gpus) if args.parallel else 1 # Embedding device: shares the vLLM GPU(s) by default; --embed_gpu_id puts it # on a dedicated card, which is then added to CUDA_VISIBLE_DEVICES. visible_gpus = list(gpus) if args.parallel else [args.gpu_id] embed_device = None if args.embed_gpu_id is not None: if args.embed_gpu_id not in visible_gpus: visible_gpus.append(args.embed_gpu_id) embed_device = f\"cuda:{visible_gpus.index(args.embed_gpu_id)}\" os.environ[\"CUDA_VISIBLE_DEVICES\"] = \",\".join(visible_gpus) # load configurations configs = load_configs_from_yaml(\"configs.yaml\") llm_config = LLMConfig(configs[\"llm_configs\"][\"general_configs\"], configs[\"llm_configs\"][args.model_name], name=args.model_name) llm_config.tensor_parallel_size = tensor_parallel_size print(\"=\" * 80) reasoning_mode = \"Self-Consistency\" if args.self_consistency else \"CoT\" print(f\"Task: {args.task} | Reasoning: {reasoning_mode}\") if args.self_consistency: print(f\"Reasoning paths per query: {max(1, args.num_reasoning_paths)}\") print(f\"Model: {llm_config.model} | Backend: {llm_config.backend}\") print(\"=\" * 80) model = build_language_model(llm_config) memory_db = None if args.if_use_memory: # load embedding model embedding_config = EmbeddingConfig(configs[\"embedding_configs\"][args.embedding_model]) embedding_model = EmbeddingModel(embedding_config, device=embed_device) memory_model = args.memory_model if args.memory_model else args.model_name if args.collection_name: collection_name = f\"{args.task}_{memory_model}_{args.collection_name}\" else: collection_name = f\"{args.task}_{memory_model}\" print(f\"Collection name: {collection_name} (memory from: {memory_model}, agent: {args.memory_agent_id})\") # initialize memory database memory_db = MemoryVectorDB( persistent_dir=PERSISTENT_DIR, collection_name=collection_name, embedding_model=embedding_model, ) npz_path = os.path.join(PERSISTENT_DIR, f\"{collection_name}_doc_embeddings.npz\") if os.path.exists(npz_path): memory_db.load_doc_embeddings(npz_path) else: print(f\"Warning: doc embeddings not found at {npz_path}, MMR will fall back to on-the-fly encode.\") start_time = datetime.now() agent = build_agent(args, memory_db=memory_db) cot = ChainOfThought( agent=agent, model=model, task_type=args.task, num_reasoning_paths=args.num_reasoning_paths, self_consistency=args.self_consistency, ) queries = load_questions(args.task, args.limit, args.train) print(f\"Loaded {len(queries)} questions\") all_results = [] for batch_start in range(0, len(queries), args.batch_size): batch_idx = batch_start // args.batch_size + 1 batch = queries[batch_start: batch_start + args.batch_size] print(f\"\\n>>> Batch {batch_idx} \" f\"(questions {batch_start}-{batch_start + len(batch) - 1})\") all_results.extend(cot.run(batch)) metrics = compute_metrics(all_results) print(\"\\n\" + \"=\" * 80) print(f\"Accuracy: {metrics['correct']}/{metrics['total']} = {metrics['accuracy']:.4f}\") print(\"=\" * 80) token_summary = cot.get_token_usage_summary() print(f\"Token usage: {token_summary}\") elapsed = datetime.now() - start_time print(f\"\\nTotal time: {elapsed}\") if args.save_log: out_dir = os.path.join(LOG_DIR, args.task, args.model_name) os.makedirs(out_dir, exist_ok=True) timestamp = datetime.now().strftime(\"%Y%m%d%H%M\") prefix = f\"{args.exp_name}_\" if args.exp_name else \"\" suffix = f\"_{args.run_suffix}\" if args.run_suffix else \"\" out_path = os.path.join( out_dir, f\"{prefix}{timestamp}{suffix}.json\" ) with open(out_path, \"w\", encoding=\"utf-8\") as f: json.dump({ \"args\": vars(args), \"metrics\": metrics, \"time_consumption\": str(elapsed), \"token_usage\": token_summary, \"results\": all_results, }, f, ensure_ascii=False, indent=2) print(f\"Results saved to {out_path}\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"run_cot.py","relation":"imports","specifier":"os","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"json","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"argparse","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"datetime","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"typing","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.agent","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.memory","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.models","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.chain_of_thought","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.paths","file":"run_cot.py"},{"source":"run_cot.py","relation":"imports","specifier":"src.utils","file":"run_cot.py"},{"source":"run_cot.py","relation":"contains","targetId":"run_cot.py#parse_args","file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#parse_args","relation":"calls","name":"parse_args","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py","relation":"contains","targetId":"run_cot.py#build_agent","file":"run_cot.py"},{"source":"run_cot.py#build_agent","relation":"calls","name":"AgentId","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#build_agent","relation":"calls","name":"Agent","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py","relation":"contains","targetId":"run_cot.py#load_questions","file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"load_dataset","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"load_dataset","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"enumerate","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"format_question","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"append","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#load_questions","relation":"calls","name":"Query","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py","relation":"contains","targetId":"run_cot.py#compute_metrics","file":"run_cot.py"},{"source":"run_cot.py#compute_metrics","relation":"calls","name":"len","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#compute_metrics","relation":"calls","name":"sum","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#compute_metrics","relation":"calls","name":"get","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py","relation":"contains","targetId":"run_cot.py#main","file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"parse_args","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"load_env_file","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"split","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"list","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"append","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"index","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"load_configs_from_yaml","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"LLMConfig","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"max","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"build_language_model","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"EmbeddingConfig","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"EmbeddingModel","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"MemoryVectorDB","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"exists","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"load_doc_embeddings","viaMember":true,"file":"run_cot.py","recvType":"MemoryVectorDB"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"build_agent","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"ChainOfThought","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"load_questions","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"range","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"extend","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"run","viaMember":true,"file":"run_cot.py","recvType":"ChainOfThought"},{"source":"run_cot.py#main","relation":"calls","name":"compute_metrics","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"get_token_usage_summary","viaMember":true,"file":"run_cot.py","recvType":"ChainOfThought"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"makedirs","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"strftime","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"open","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"dump","viaMember":true,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"vars","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"str","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_cot.py"},{"source":"run_cot.py","relation":"calls","name":"main","viaMember":false,"file":"run_cot.py"}]},"run_mad.py":{"size":9887,"mtimeMs":1789362123398.8032,"hash":"e9a6e8b8705dcbf15d91ac7bed6827900850719d2b38cf6e84b7dc491a8b25bd","nodes":[{"id":"run_mad.py","name":"run_mad.py","kind":"file","path":"run_mad.py","span":"L1-L224","signature":null,"exported":true,"origin":"ast","body_hash":"e9a6e8b8705dcbf15d91ac7bed6827900850719d2b38cf6e84b7dc491a8b25bd","chars":9887,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import json import argparse from datetime import datetime from typing import Dict, List os.environ[\"VLLM_WORKER_MULTIPROC_METHOD\"] = \"spawn\" from src.agent import Agent, AgentId, Query from src.memory import MemoryVectorDB from src.models import build_language_model, EmbeddingModel from src.multi_agent_debate import MultiAgentDebate from src.paths import PERSISTENT_DIR, LOG_DIR from src.utils import LLMConfig, EmbeddingConfig, load_configs_from_yaml, load_env_file, DataLoader if __name__ == \"__main__\": main()"},{"id":"run_mad.py#parse_args","name":"parse_args","kind":"function","path":"run_mad.py","span":"L17-L55","signature":"def parse_args()","exported":true,"origin":"ast","body_hash":"50e63f75046398c2ed40e3c45cb91428ed0fe289cfd3475af25d50b429b80a24","body_text":"def parse_args(): parser = argparse.ArgumentParser(description=\"Run Multi-Agent Debate\") # debate settings parser.add_argument(\"--task\", type=str, default=\"MATH500\", help=\"Task / dataset type\") parser.add_argument(\"--n_agents\", type=int, default=3, help=\"Number of debate agents\") parser.add_argument(\"--max_round\", type=int, default=3, help=\"Max debate rounds\") parser.add_argument(\"--consensus_threshold\", type=float, default=0.8, help=\"Fraction of agents that must agree to declare consensus\") parser.add_argument(\"--if_summarize_state\", action=\"store_true\", help=\"Whether to summarize debate state for memory retrieval\") # agent settings parser.add_argument(\"--model_name\", type=str, default=\"qwen3-8b\", help=\"The name of adopted LLM\") parser.add_argument(\"--if_use_memory\", action=\"store_true\", help=\"Use memory or not\") parser.add_argument(\"--memory_model\", type=str, default=\"\", help=\"Model name whose memory collection to use (default: same as --model_name)\") parser.add_argument(\"--if_confidence_weight\", action=\"store_true\", help=\"Weight other agents' responses by memory-estimated confidence\") parser.add_argument(\"--collection_name\", type=str, default=\"\", help=\"Optional suffix for the ChromaDB collection name\") parser.add_argument(\"--embedding_model\", type=str, default=\"bge-m3\", help=\"The name of adopted embedding model\") parser.add_argument(\"--n_retrieval\", type=int, default=1, help=\"Number of retrieval\") parser.add_argument(\"--retrieval_policy\", type=str, default=\"debate_state_aware\", choices=[\"debate_state_aware\", \"similarity\", \"random\", \"diversity\", \"fixed_lambda\", \"positive_only\"], help=\"Memory retrieval policy (only active with --if_use_memory)\") parser.add_argument(\"--seed\", type=int, default=None, help=\"Base seed for stochastic retrieval policies (random/diversity)\") parser.add_argument(\"--ref_lambda\", type=float, default=0.7, help=\"Reference MMR lambda used with no consensus; the constant lambda for the fixed_lambda policy\") # general settings parser.add_argument(\"--exp_name\", type=str, default=\"main\", help=\"Optional prefix for the output filename\") parser.add_argument(\"--run_suffix\", type=str, default=\"\", help=\"Optional suffix appended before .json\") parser.add_argument(\"--save_log\", action=\"store_true\", help=\"Whether to save debate log to disk\") parser.add_argument(\"--train\", action=\"store_true\", help=\"Train or not\") parser.add_argument(\"--batch_size\", type=int, default=64, help=\"Number of queries per batch\") parser.add_argument(\"--limit\", type=int, default=None, help=\"Cap the number of questions (for quick tests)\") # gpu settings parser.add_argument(\"--gpu_id\", type=str, default=\"0\", help=\"The GPU ID.\") parser.add_argument(\"--embed_gpu_id\", type=str, default=None, help=\"Dedicated GPU for the embedding model. Default None = share --gpu_id.\") parser.add_argument(\"--parallel\", action=\"store_true\", default=False, help=\"Whether to use parallel mode.\") args = parser.parse_args() return args","summary_state":"pending","summary":null,"crux":null},{"id":"run_mad.py#build_agents","name":"build_agents","kind":"function","path":"run_mad.py","span":"L58-L74","signature":"def build_agents(args, memory_db=None) -> List[Agent]","exported":true,"origin":"ast","body_hash":"3c0492775460ac03396688e2de904488543ecde26eaa5fe69f81fba8f0b16a91","body_text":"def build_agents(args, memory_db=None) -> List[Agent]: agents = [] for i in range(args.n_agents): agent_id = AgentId(agent_type=\"Debater\", agent_key=str(i)) agent = Agent( agent_id=agent_id, task_type=args.task, if_use_memory=args.if_use_memory, memory_db=memory_db, n_retrieval=args.n_retrieval, if_confidence_weight=args.if_confidence_weight, retrieval_policy=args.retrieval_policy, seed=args.seed, ref_lambda=args.ref_lambda, ) agents.append(agent) return agents","summary_state":"pending","summary":null,"crux":null},{"id":"run_mad.py#load_questions","name":"load_questions","kind":"function","path":"run_mad.py","span":"L77-L96","signature":"def load_questions(task_type: str, limit: int | None, is_training: bool) -> List[Query]","exported":true,"origin":"ast","body_hash":"7c672cc3872f3106fd22913dd3d1269eb7e7e4a911fddd46f114837c459c0b26","body_text":"def load_questions(task_type: str, limit: int | None, is_training: bool) -> List[Query]: if is_training: raw_data = DataLoader.load_dataset(task_type + \"_TRAIN\") else: raw_data = DataLoader.load_dataset(task_type) queries = [] for idx, item in enumerate(raw_data): if limit is not None and idx >= limit: break content, answer, solution, category = DataLoader.format_question(item, task_type) queries.append(Query( content=content, query_id=idx, answer=answer, solution=solution, category=category, )) return queries","summary_state":"pending","summary":null,"crux":null},{"id":"run_mad.py#compute_metrics","name":"compute_metrics","kind":"function","path":"run_mad.py","span":"L99-L103","signature":"def compute_metrics(results: List[Dict]) -> Dict","exported":true,"origin":"ast","body_hash":"9960de980eb147c12b4c08466b9d816ceee86619f0ae459193a470768e6dfbc1","body_text":"def compute_metrics(results: List[Dict]) -> Dict: total = len(results) correct = sum(1 for r in results if r.get(\"is_final_correct\", False)) accuracy = correct / total if total > 0 else 0.0 return {\"correct\": correct, \"total\": total, \"accuracy\": accuracy}","summary_state":"pending","summary":null,"crux":null},{"id":"run_mad.py#main","name":"main","kind":"function","path":"run_mad.py","span":"L106-L219","signature":"def main()","exported":true,"origin":"ast","body_hash":"e5d960caf069e0a3fa26be74051f11b0fafe6496bb5189c9d4821ddccb1c8fa4","body_text":"def main(): args = parse_args() # Load API keys etc. from .env (no-op if the file is absent) load_env_file() # GPU settings gpus = args.gpu_id.split(\",\") tensor_parallel_size = len(gpus) if args.parallel else 1 # Embedding device: shares the vLLM GPU(s) by default; --embed_gpu_id puts it on a dedicated card. # The embed GPU is added to CUDA_VISIBLE_DEVICES so it is reachable. visible_gpus = list(gpus) embed_device = None if args.embed_gpu_id is not None: if args.embed_gpu_id not in visible_gpus: visible_gpus.append(args.embed_gpu_id) embed_device = f\"cuda:{visible_gpus.index(args.embed_gpu_id)}\" os.environ[\"CUDA_VISIBLE_DEVICES\"] = \",\".join(visible_gpus) # load configurations configs = load_configs_from_yaml(\"configs.yaml\") llm_config = LLMConfig(configs[\"llm_configs\"][\"general_configs\"], configs[\"llm_configs\"][args.model_name], name=args.model_name) llm_config.tensor_parallel_size = tensor_parallel_size print(\"=\" * 80) print(f\"Task: {args.task} | Agents: {args.n_agents} | \" f\"Max rounds: {args.max_round} | Consensus: {args.consensus_threshold}\") print(f\"Model: {llm_config.model} | Backend: {llm_config.backend}\") print(\"=\" * 80) model = build_language_model(llm_config) memory_db = None if args.if_use_memory or args.if_confidence_weight: # load embedding model embedding_config = EmbeddingConfig(configs[\"embedding_configs\"][args.embedding_model]) embedding_model = EmbeddingModel(embedding_config, device=embed_device) memory_model = args.memory_model if args.memory_model else args.model_name if args.collection_name: collection_name = f\"{args.task}_{memory_model}_{args.collection_name}\" else: collection_name = f\"{args.task}_{memory_model}\" print(f\"Collection name: {collection_name} (memory from: {memory_model})\") # initialize memory database memory_db = MemoryVectorDB( persistent_dir=PERSISTENT_DIR, collection_name=collection_name, embedding_model=embedding_model, ) npz_path = os.path.join(PERSISTENT_DIR, f\"{collection_name}_doc_embeddings.npz\") if os.path.exists(npz_path): memory_db.load_doc_embeddings(npz_path) else: print(f\"Warning: doc embeddings not found at {npz_path}, MMR will fall back to on-the-fly encode.\") start_time = datetime.now() agents = build_agents(args, memory_db=memory_db) debate = MultiAgentDebate( agents=agents, model=model, task_type=args.task, max_round=args.max_round, consensus_threshold=args.consensus_threshold, if_summarize_state=args.if_summarize_state, if_confidence_weight=args.if_confidence_weight, ) queries = load_questions(args.task, args.limit, args.train) print(f\"Loaded {len(queries)} questions\") all_results = [] for batch_start in range(0, len(queries), args.batch_size): batch_idx = batch_start // args.batch_size + 1 batch = queries[batch_start: batch_start + args.batch_size] print(f\"\\n>>> Batch {batch_idx} \" f\"(questions {batch_start}-{batch_start + len(batch) - 1})\") all_results.extend(debate.run(queries=batch, is_training=args.train, batch_idx=batch_idx)) metrics = compute_metrics(all_results) print(\"\\n\" + \"=\" * 80) print(f\"Accuracy: {metrics['correct']}/{metrics['total']} = {metrics['accuracy']:.4f}\") print(\"=\" * 80) token_summary = debate.get_token_usage_summary() print(f\"Token usage: {token_summary}\") elapsed = datetime.now() - start_time print(f\"\\nTotal time: {elapsed}\") if args.save_log: out_dir = os.path.join(LOG_DIR, args.task, args.model_name) os.makedirs(out_dir, exist_ok=True) timestamp = datetime.now().strftime(\"%Y%m%d%H%M\") prefix = f\"{args.exp_name}_\" if args.exp_name else \"\" suffix = f\"_{args.run_suffix}\" if args.run_suffix else \"\" out_path = os.path.join( out_dir, f\"{prefix}agent{args.n_agents}_round{args.max_round}_{timestamp}{suffix}.json\" ) with open(out_path, \"w\", encoding=\"utf-8\") as f: json.dump({ \"args\": vars(args), \"metrics\": metrics, \"time_consumption\": str(elapsed), \"token_usage\": token_summary, \"results\": all_results, }, f, ensure_ascii=False, indent=2) print(f\"Results saved to {out_path}\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"run_mad.py","relation":"imports","specifier":"os","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"json","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"argparse","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"datetime","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"typing","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.agent","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.memory","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.models","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.multi_agent_debate","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.paths","file":"run_mad.py"},{"source":"run_mad.py","relation":"imports","specifier":"src.utils","file":"run_mad.py"},{"source":"run_mad.py","relation":"contains","targetId":"run_mad.py#parse_args","file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"add_argument","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#parse_args","relation":"calls","name":"parse_args","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py","relation":"contains","targetId":"run_mad.py#build_agents","file":"run_mad.py"},{"source":"run_mad.py#build_agents","relation":"calls","name":"range","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#build_agents","relation":"calls","name":"AgentId","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#build_agents","relation":"calls","name":"str","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#build_agents","relation":"calls","name":"Agent","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#build_agents","relation":"calls","name":"append","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py","relation":"contains","targetId":"run_mad.py#load_questions","file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"load_dataset","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"load_dataset","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"enumerate","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"format_question","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"append","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#load_questions","relation":"calls","name":"Query","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py","relation":"contains","targetId":"run_mad.py#compute_metrics","file":"run_mad.py"},{"source":"run_mad.py#compute_metrics","relation":"calls","name":"len","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#compute_metrics","relation":"calls","name":"sum","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#compute_metrics","relation":"calls","name":"get","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py","relation":"contains","targetId":"run_mad.py#main","file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"parse_args","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"load_env_file","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"split","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"list","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"append","viaMember":true,"file":"run_mad.py","recvType":"list"},{"source":"run_mad.py#main","relation":"calls","name":"index","viaMember":true,"file":"run_mad.py","recvType":"list"},{"source":"run_mad.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"load_configs_from_yaml","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"LLMConfig","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"build_language_model","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"EmbeddingConfig","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"EmbeddingModel","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"MemoryVectorDB","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"exists","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"load_doc_embeddings","viaMember":true,"file":"run_mad.py","recvType":"MemoryVectorDB"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"build_agents","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"MultiAgentDebate","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"load_questions","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"range","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"len","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"extend","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"run","viaMember":true,"file":"run_mad.py","recvType":"MultiAgentDebate"},{"source":"run_mad.py#main","relation":"calls","name":"compute_metrics","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"get_token_usage_summary","viaMember":true,"file":"run_mad.py","recvType":"MultiAgentDebate"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"makedirs","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"strftime","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"now","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"join","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"open","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"dump","viaMember":true,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"vars","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"str","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py#main","relation":"calls","name":"print","viaMember":false,"file":"run_mad.py"},{"source":"run_mad.py","relation":"calls","name":"main","viaMember":false,"file":"run_mad.py"}]},"src/agent.py":{"size":18686,"mtimeMs":1789362123399.8032,"hash":"137992b7b83ee1a745b10e98c7f89c2a06864cbd0634d4ff8ddaf2136b1b8f14","nodes":[{"id":"src/agent.py","name":"agent.py","kind":"file","path":"src/agent.py","span":"L1-L428","signature":null,"exported":true,"origin":"ast","body_hash":"137992b7b83ee1a745b10e98c7f89c2a06864cbd0634d4ff8ddaf2136b1b8f14","chars":18676,"summary_state":"pending","summary":null,"crux":null,"body_text":"import numpy as np from dataclasses import dataclass from typing import List, Dict, Optional, Tuple from src.memory import MemoryVectorDB from src.retrieval_policy import build_retrieval_policy from src.prompts import MAD_SYS_PROMPTS, TASK_PROMPT, DEBATE_PROMPT @dataclass @dataclass(frozen=True) @dataclass"},{"id":"src/agent.py#Query","name":"Query","kind":"class","path":"src/agent.py","span":"L11-L23","signature":"class Query","exported":true,"origin":"ast","body_hash":"5b5b267ea593c8ea53ac212a90443017c6f7f79ee966d7d738863fc87add2c77","body_text":"class Query: \"\"\"Encapsulates a query and its expected answer.\"\"\" content: str query_id: int answer: str solution: str category: str def __str__(self): return f\"The Query {self.query_id} is :: {self.content}\" def __repr__(self): return self.__str__()","summary_state":"pending","summary":null,"crux":null},{"id":"src/agent.py#Query.__str__","name":"__str__","kind":"method","path":"src/agent.py","span":"L19-L20","signature":"def __str__(self)","exported":false,"origin":"ast","body_hash":"51eb38a8f94a49dbf7525c3ea9e91464e065f5a790fd3ed364829fe48433dfc0","body_text":"def __str__(self): return f\"The Query {self.query_id} is :: {self.content}\"","summary_state":"pending","summary":null,"crux":null,"owner":"Query"},{"id":"src/agent.py#Query.__repr__","name":"__repr__","kind":"method","path":"src/agent.py","span":"L22-L23","signature":"def __repr__(self)","exported":false,"origin":"ast","body_hash":"700390b46df22b46e518b9ace82ed2db94b4a504974e2309ca324c4fe5c9b44d","body_text":"def __repr__(self): return self.__str__()","summary_state":"pending","summary":null,"crux":null,"owner":"Query"},{"id":"src/agent.py#AgentId","name":"AgentId","kind":"class","path":"src/agent.py","span":"L27-L35","signature":"class AgentId","exported":true,"origin":"ast","body_hash":"f2465abd5f867bc8a10c1d26400b75f84b76c4d1cfd383be4537aa2000b9cae5","body_text":"class AgentId: agent_type: str agent_key: str def __str__(self) -> str: return f\"{self.agent_type}-{self.agent_key}\" def __repr__(self) -> str: return self.__str__()","summary_state":"pending","summary":null,"crux":null},{"id":"src/agent.py#AgentId.__str__","name":"__str__","kind":"method","path":"src/agent.py","span":"L31-L32","signature":"def __str__(self) -> str","exported":false,"origin":"ast","body_hash":"4f3bec120c18ecc0fc3cd09a1e3298e964ac0f4a6b03690e66b22edcef7fea82","body_text":"def __str__(self) -> str: return f\"{self.agent_type}-{self.agent_key}\"","summary_state":"pending","summary":null,"crux":null,"owner":"AgentId"},{"id":"src/agent.py#AgentId.__repr__","name":"__repr__","kind":"method","path":"src/agent.py","span":"L34-L35","signature":"def __repr__(self) -> str","exported":false,"origin":"ast","body_hash":"8a19aab5e355b09ec08b4c7b65f59cb0b5eac756392c13897570aeb9fe2fa0eb","body_text":"def __repr__(self) -> str: return self.__str__()","summary_state":"pending","summary":null,"crux":null,"owner":"AgentId"},{"id":"src/agent.py#RetrievedMemory","name":"RetrievedMemory","kind":"class","path":"src/agent.py","span":"L39-L50","signature":"class RetrievedMemory","exported":true,"origin":"ast","body_hash":"47e72389e38a52abb5a076ebc5c7e1e45aa89b7b999ac6a6abc803f255084bc8","body_text":"class RetrievedMemory: content: str query_id: int agent_id: AgentId is_correct: bool = False similarity: float = 0.0 def __str__(self): return f\"Memory of {self.agent_id} for Q{self.query_id}: {self.content}\" def __repr__(self): return self.__str__()","summary_state":"pending","summary":null,"crux":null},{"id":"src/agent.py#RetrievedMemory.__str__","name":"__str__","kind":"method","path":"src/agent.py","span":"L46-L47","signature":"def __str__(self)","exported":false,"origin":"ast","body_hash":"a9a8249e47c5c39f0a0878ef8b9f1d321a85115aee5fbb4ff814f33d1469d675","body_text":"def __str__(self): return f\"Memory of {self.agent_id} for Q{self.query_id}: {self.content}\"","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievedMemory"},{"id":"src/agent.py#RetrievedMemory.__repr__","name":"__repr__","kind":"method","path":"src/agent.py","span":"L49-L50","signature":"def __repr__(self)","exported":false,"origin":"ast","body_hash":"700390b46df22b46e518b9ace82ed2db94b4a504974e2309ca324c4fe5c9b44d","body_text":"def __repr__(self): return self.__str__()","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievedMemory"},{"id":"src/agent.py#Agent","name":"Agent","kind":"class","path":"src/agent.py","span":"L53-L427","signature":"class Agent","exported":true,"origin":"ast","body_hash":"7c2540ffb2f2333ebe49d70a54210df8be40f188856f9e25cc238fdfd95afe09","body_text":"class Agent: \"\"\" Manages prompt construction and memory retrieval for a single debate agent. Conversation histories are stored per query_id so that one Agent instance can participate in debates on multiple querys simultaneously (batch mode). Inference is delegated to the shared LanguageModel instance. \"\"\" def __init__( self, agent_id: AgentId, task_type: str, if_use_memory: bool = False, memory_db: Optional[MemoryVectorDB] = None, n_retrieval: int = 3, if_confidence_weight: bool = False, retrieval_policy: str = \"debate_state_aware\", seed: Optional[int] = None, ref_lambda: Optional[float] = None, retrieve_first_round: bool = False, memory_agent_id: Optional[str] = None, ) -> None: self.agent_id = agent_id self.task_type = task_type self.if_use_memory = if_use_memory self.memory_db = memory_db self.n_retrieval = n_retrieval self.if_confidence_weight = if_confidence_weight # Retrieve on round 0 too (used by CoT, which has no debate rounds). self._retrieve_first_round = retrieve_first_round # agent_id string used to filter the memory collection; defaults to this # agent's own id. CoT overrides it (e.g. \"Debater-0\") to reuse memory # collected under the debate agent ids. self._memory_agent_id = memory_agent_id or str(agent_id) # Offset the base seed by the agent index policy_seed = None if seed is None else seed + int(agent_id.agent_key) self._policy = build_retrieval_policy(retrieval_policy, seed=policy_seed, ref_lambda=ref_lambda) self._system_prompt = MAD_SYS_PROMPTS[task_type][int(agent_id.agent_key)] # Per-query conversation state # _histories[query_id] = list of {\"role\": ..., \"content\": ...} self._histories: Dict[int, List[Dict[str, str]]] = {} # _rounds[query_id] = number of completed rounds self._rounds: Dict[int, int] = {} # ------------------------------------------------------------------ # Conversation management # ------------------------------------------------------------------ def initialize_conversation(self, query: Query) -> None: \"\"\"Register a new query and reset its conversation state.\"\"\" self._histories[query.query_id] = [] self._rounds[query.query_id] = 0 def clear_conversation(self, query: Query) -> None: self._histories.pop(query.query_id, None) self._rounds.pop(query.query_id, None) def clear_all_conversations(self) -> None: self._histories.clear() self._rounds.clear() # ------------------------------------------------------------------ # Prompt building # ------------------------------------------------------------------ def _build_prompt( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]], retrieved_memories: Optional[List[RetrievedMemory]] = None, other_confidences: Optional[Dict[AgentId, float]] = None, ) -> str: task_prompt = \"\" if retrieved_memories: task_prompt += \"Here are some examples:\\n\\n<examples>\\n\" for idx, memory in enumerate(retrieved_memories): task_prompt += f'<example index=\"{idx}\">\\n{memory.content}\\n</example>\\n' task_prompt += \"</examples>\\n\\n\" if prev_responses is not None: other_agents = [ (aid, resp) for aid, (resp, _) in prev_responses.items() if aid != self.agent_id and resp ] if self.if_confidence_weight and other_confidences: task_prompt += \"These are the solutions to the problem from other agents with your confidences:\\n\\n<other_solutions>\\n\" else: task_prompt += \"These are the solutions to the problem from other agents:\\n\\n<other_solutions>\\n\" for idx, (aid, response) in enumerate(other_agents): confidence_str = \"\" if other_confidences and aid in other_confidences: if other_confidences[aid] > 0.55: confidence_str = f'<confidence>\\nHigh, historically reliable on relative problems.\\n</confidence>\\n' elif other_confidences[aid] < 0.45: confidence_str = f'<confidence>\\nLow, low historical accuracy on relative problems.\\n</confidence>\\n' task_prompt += f'<other_solution index=\"{idx}\">\\n{response}\\n{confidence_str}</other_solution>\\n' task_prompt += \"</other_solutions>\\n\\n\" task_prompt += DEBATE_PROMPT[self.task_type].format(QUESTION=query.content) else: task_prompt += TASK_PROMPT[self.task_type].format(QUESTION=query.content) return task_prompt def build_messages( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]] = None, prev_round_summary: Optional[str] = None, prev_consensus_ratio: Optional[float] = None, prefetched: Optional[Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]] = None, ) -> List[Dict[str, str]]: \"\"\" Assemble the full message list (system + history + new user turn). Memory retrieval is handled internally when if_use_memory is True. prev_responses=None → initial round (round 0) prev_responses={aid: (resp, emb)} → debate round; includes all agents (self + others); emb is a pre-computed unit-norm embedding or None prev_round_summary → summary of the previous round (for memory retrieval) prev_consensus_ratio → consensus ratio of the previous round (for memory re-ranking) prefetched → (retrieved_memories, other_confidences) pre-computed by batch_retrieve_","summary_state":"pending","summary":null,"crux":null},{"id":"src/agent.py#Agent.__init__","name":"__init__","kind":"method","path":"src/agent.py","span":"L61-L100","signature":"def __init__( self, agent_id: AgentId, task_type: str, if_use_memory: bool = False, memory_db: Optional[MemoryVectorDB] = None, n_retrieval: int = 3, if_confidence_weight: bool = False, retrieval_policy: str = \"debate_state_aware\", seed: Optional[int] = None, ref_lambda: Optional[float] = None, retrieve_first_round: bool = False, memory_agent_id: Optional[str] = None, ) -> None","exported":false,"origin":"ast","body_hash":"ba966bb91506bd17f9ddb8c66e616755221f8fa08c78e0e0f58eba67d5c11b42","body_text":"def __init__( self, agent_id: AgentId, task_type: str, if_use_memory: bool = False, memory_db: Optional[MemoryVectorDB] = None, n_retrieval: int = 3, if_confidence_weight: bool = False, retrieval_policy: str = \"debate_state_aware\", seed: Optional[int] = None, ref_lambda: Optional[float] = None, retrieve_first_round: bool = False, memory_agent_id: Optional[str] = None, ) -> None: self.agent_id = agent_id self.task_type = task_type self.if_use_memory = if_use_memory self.memory_db = memory_db self.n_retrieval = n_retrieval self.if_confidence_weight = if_confidence_weight # Retrieve on round 0 too (used by CoT, which has no debate rounds). self._retrieve_first_round = retrieve_first_round # agent_id string used to filter the memory collection; defaults to this # agent's own id. CoT overrides it (e.g. \"Debater-0\") to reuse memory # collected under the debate agent ids. self._memory_agent_id = memory_agent_id or str(agent_id) # Offset the base seed by the agent index policy_seed = None if seed is None else seed + int(agent_id.agent_key) self._policy = build_retrieval_policy(retrieval_policy, seed=policy_seed, ref_lambda=ref_lambda) self._system_prompt = MAD_SYS_PROMPTS[task_type][int(agent_id.agent_key)] # Per-query conversation state # _histories[query_id] = list of {\"role\": ..., \"content\": ...} self._histories: Dict[int, List[Dict[str, str]]] = {} # _rounds[query_id] = number of completed rounds self._rounds: Dict[int, int] = {}","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.initialize_conversation","name":"initialize_conversation","kind":"method","path":"src/agent.py","span":"L106-L109","signature":"def initialize_conversation(self, query: Query) -> None","exported":true,"origin":"ast","body_hash":"814a1a589fd5653fa88b5185a9b9cb1554d9cf1009a35e7a1e1d7d07d23e99be","body_text":"def initialize_conversation(self, query: Query) -> None: \"\"\"Register a new query and reset its conversation state.\"\"\" self._histories[query.query_id] = [] self._rounds[query.query_id] = 0","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.clear_conversation","name":"clear_conversation","kind":"method","path":"src/agent.py","span":"L111-L113","signature":"def clear_conversation(self, query: Query) -> None","exported":true,"origin":"ast","body_hash":"bafcd877d6901ada1e1acea988dea91532a2567939ea939f02b76dc91e83e19f","body_text":"def clear_conversation(self, query: Query) -> None: self._histories.pop(query.query_id, None) self._rounds.pop(query.query_id, None)","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.clear_all_conversations","name":"clear_all_conversations","kind":"method","path":"src/agent.py","span":"L115-L117","signature":"def clear_all_conversations(self) -> None","exported":true,"origin":"ast","body_hash":"081d917c26c4444c98139b74db2f3c0a080173ec40d6a3addb840cfc99afd48b","body_text":"def clear_all_conversations(self) -> None: self._histories.clear() self._rounds.clear()","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent._build_prompt","name":"_build_prompt","kind":"method","path":"src/agent.py","span":"L123-L164","signature":"def _build_prompt( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]], retrieved_memories: Optional[List[RetrievedMemory]] = None, other_confidences: Optional[Dict[AgentId, float]] = None, ) -> str","exported":false,"origin":"ast","body_hash":"5bca85e13af3766328ba79784183adc3097db0695a8ea541b4a8f7be1b26be1c","body_text":"def _build_prompt( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]], retrieved_memories: Optional[List[RetrievedMemory]] = None, other_confidences: Optional[Dict[AgentId, float]] = None, ) -> str: task_prompt = \"\" if retrieved_memories: task_prompt += \"Here are some examples:\\n\\n<examples>\\n\" for idx, memory in enumerate(retrieved_memories): task_prompt += f'<example index=\"{idx}\">\\n{memory.content}\\n</example>\\n' task_prompt += \"</examples>\\n\\n\" if prev_responses is not None: other_agents = [ (aid, resp) for aid, (resp, _) in prev_responses.items() if aid != self.agent_id and resp ] if self.if_confidence_weight and other_confidences: task_prompt += \"These are the solutions to the problem from other agents with your confidences:\\n\\n<other_solutions>\\n\" else: task_prompt += \"These are the solutions to the problem from other agents:\\n\\n<other_solutions>\\n\" for idx, (aid, response) in enumerate(other_agents): confidence_str = \"\" if other_confidences and aid in other_confidences: if other_confidences[aid] > 0.55: confidence_str = f'<confidence>\\nHigh, historically reliable on relative problems.\\n</confidence>\\n' elif other_confidences[aid] < 0.45: confidence_str = f'<confidence>\\nLow, low historical accuracy on relative problems.\\n</confidence>\\n' task_prompt += f'<other_solution index=\"{idx}\">\\n{response}\\n{confidence_str}</other_solution>\\n' task_prompt += \"</other_solutions>\\n\\n\" task_prompt += DEBATE_PROMPT[self.task_type].format(QUESTION=query.content) else: task_prompt += TASK_PROMPT[self.task_type].format(QUESTION=query.content) return task_prompt","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.build_messages","name":"build_messages","kind":"method","path":"src/agent.py","span":"L166-L211","signature":"def build_messages( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]] = None, prev_round_summary: Optional[str] = None, prev_consensus_ratio: Optional[float] = None, prefetched: Optional[Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]] = None, ) -> List[Dict[str, str]]","exported":true,"origin":"ast","body_hash":"b08378165339a15fae656f20688158921f45df3038c14e236b746a5745753d24","body_text":"def build_messages( self, query: Query, prev_responses: Optional[Dict[AgentId, Tuple[str, Optional[np.ndarray]]]] = None, prev_round_summary: Optional[str] = None, prev_consensus_ratio: Optional[float] = None, prefetched: Optional[Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]] = None, ) -> List[Dict[str, str]]: \"\"\" Assemble the full message list (system + history + new user turn). Memory retrieval is handled internally when if_use_memory is True. prev_responses=None → initial round (round 0) prev_responses={aid: (resp, emb)} → debate round; includes all agents (self + others); emb is a pre-computed unit-norm embedding or None prev_round_summary → summary of the previous round (for memory retrieval) prev_consensus_ratio → consensus ratio of the previous round (for memory re-ranking) prefetched → (retrieved_memories, other_confidences) pre-computed by batch_retrieve_all; skips internal DB queries when provided \"\"\" round_idx = self._rounds.get(query.query_id, 0) retrieved_memories = None other_confidences = None if (round_idx > 0 or self._retrieve_first_round) and (self.if_use_memory or self.if_confidence_weight) and self.memory_db: if prefetched is not None: all_retrieved, other_confidences = prefetched retrieved_memories = all_retrieved if self.if_use_memory else None else: # No prefetched: run the same retrieval policy for this single query result = self.batch_retrieve_all( queries=[query], prev_responses_map={query.query_id: prev_responses}, prev_round_summaries={query.query_id: prev_round_summary}, prev_consensus_ratios={query.query_id: prev_consensus_ratio}, ) all_retrieved, other_confidences = result.get(query.query_id, ([], None)) retrieved_memories = all_retrieved if self.if_use_memory else None task_prompt = self._build_prompt(query, prev_responses, retrieved_memories, other_confidences) self._histories[query.query_id].append({\"role\": \"user\", \"content\": task_prompt}) messages = [{\"role\": \"system\", \"content\": self._system_prompt}] messages.extend(self._histories[query.query_id]) return messages","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.record_response","name":"record_response","kind":"method","path":"src/agent.py","span":"L213-L216","signature":"def record_response(self, query: Query, response: str) -> None","exported":true,"origin":"ast","body_hash":"079d06937938af035f09dbfbb6a1c9eeae494af88be0c1fe9d7de9f493e373bf","body_text":"def record_response(self, query: Query, response: str) -> None: \"\"\"Append a completed turn to the conversation history for query_id.\"\"\" self._histories[query.query_id].append({\"role\": \"assistant\", \"content\": response}) self._rounds[query.query_id] = self._rounds.get(query.query_id, 0) + 1","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent.batch_retrieve_all","name":"batch_retrieve_all","kind":"method","path":"src/agent.py","span":"L222-L322","signature":"def batch_retrieve_all( self, queries: List[Query], prev_responses_map: Dict[int, Dict[AgentId, Tuple[str, Optional[np.ndarray]]]], prev_round_summaries: Dict[int, Optional[str]], prev_consensus_ratios: Dict[int, Optional[float]], ) -> Dict[int, Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]]","exported":true,"origin":"ast","body_hash":"4c3525feff97e1f342e1f5d6950c2dd3d611c801520472d3140e4c27d7cbee51","body_text":"def batch_retrieve_all( self, queries: List[Query], prev_responses_map: Dict[int, Dict[AgentId, Tuple[str, Optional[np.ndarray]]]], prev_round_summaries: Dict[int, Optional[str]], prev_consensus_ratios: Dict[int, Optional[float]], ) -> Dict[int, Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]]: \"\"\" Batch memory retrieval for a list of queries in a single round. For each query that needs memory (round > 0, if_use_memory, memory_db exists), builds all query texts and delegates the two-stage retrieval to the active retrieval policy, then estimates per-peer confidence from each query's candidate pool. Returns: {query_id: (retrieved_memories, other_confidences)} Only includes entries for queries where retrieval was performed. \"\"\" if not (self.if_use_memory or self.if_confidence_weight) or self.memory_db is None: return {} # Collect queries that actually need retrieval (round > 0) active: List[Tuple[int, str]] = [] # (query_id, embedding_text) for query in queries: round_idx = self._rounds.get(query.query_id, 0) if round_idx == 0 and not self._retrieve_first_round: continue prev_responses = prev_responses_map.get(query.query_id, {}) prev_agent_response, _ = prev_responses.get(self.agent_id, (None, None)) prev_round_summary = prev_round_summaries.get(query.query_id) embedding_text = self._build_query_text(query, prev_agent_response, prev_round_summary) active.append((query.query_id, embedding_text)) if not active: return {} # Candidate generation + selection. # When memory is on, the retrieval policy owns the whole two-stage retrieval # and returns (retrieved, candidates) per query. When only confidence # weighting is on there is no active memory policy, so fall back to the # plain semantic candidate pool (no retrieval) for confidence. if self.if_use_memory: policy_out = self._policy.batch_retrieve( memory_db=self.memory_db, agent_id=self._memory_agent_id, queries=active, n_retrieval=self.n_retrieval, prev_consensus_ratios=prev_consensus_ratios, ) else: batch_candidates = self.memory_db.batch_query_similar( query_texts=[emb_text for _, emb_text in active], n_results=self.n_retrieval * 3, filter_metadata={\"agent_id\": {\"$eq\": self._memory_agent_id}}, ) policy_out = {qid: ([], candidates) for (qid, _), candidates in zip(active, batch_candidates)} # If confidence_weight is on, collect all (query_id, round, agent_id) keys needed # for get_response_embeddings across all candidates in one batch call resp_emb_cache: Dict[Tuple, Optional[np.ndarray]] = {} if self.if_confidence_weight: other_agent_strs = [ str(aid) for aid in ( next(iter(prev_responses_map.values()), {}).keys() ) if aid != self.agent_id ] if other_agent_strs: keys_needed: List[Tuple] = [] for _, candidates in policy_out.values(): for case_meta, _, _, _ in candidates: hist_query_id = case_meta.get(\"question_id\") hist_round = int(case_meta.get(\"current_round\", 1)) - 1 if hist_round < 0: continue for aid_str in other_agent_strs: key = (hist_query_id, hist_round, aid_str) if key not in resp_emb_cache: keys_needed.append(key) resp_emb_cache[key] = None # placeholder if keys_needed: fetched = self.memory_db.batch_get_response_embeddings(keys_needed) resp_emb_cache.update(fetched) # Per-query: build RetrievedMemory from the retrieved cases + estimate confidences # from the candidates using cached response embeddings results: Dict[int, Tuple[List[RetrievedMemory], Optional[Dict[AgentId, float]]]] = {} for qid, (retrieved, candidates) in policy_out.items(): retrieved_memories = self._to_retrieved_memories(qid, retrieved) if self.if_use_memory else [] other_confidences = None if self.if_confidence_weight: prev_responses = prev_responses_map.get(qid, {}) other_confidences = self._estimate_confidences_cached( prev_responses, candidates, resp_emb_cache ) results[qid] = (retrieved_memories, other_confidences) return results","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent._estimate_confidences_cached","name":"_estimate_confidences_cached","kind":"method","path":"src/agent.py","span":"L324-L381","signature":"def _estimate_confidences_cached( self, prev_responses: Dict[AgentId, Tuple[str, Optional[np.ndarray]]], candidate_memories: List[Tuple[Dict, str, float, Optional[np.ndarray]]], resp_emb_cache: Dict[Tuple, Optional[np.ndarray]], ) -> Dict[AgentId, float]","exported":false,"origin":"ast","body_hash":"b46a2a14195cc31cc26d649b1accf5fce26121bf36713a1ab2b1847baf6aa322","body_text":"def _estimate_confidences_cached( self, prev_responses: Dict[AgentId, Tuple[str, Optional[np.ndarray]]], candidate_memories: List[Tuple[Dict, str, float, Optional[np.ndarray]]], resp_emb_cache: Dict[Tuple, Optional[np.ndarray]], ) -> Dict[AgentId, float]: \"\"\" Estimate a confidence score for each other agent's previous response from its similarity to historical responses in the candidate pool. For each other agent o and each candidate case c (which records what other agents answered at round c[\"current_round\"] - 1): sim(o, c) = cosine(current response of o, historical response of o in c) reward(c) = response_correct of c (correctness at that round) confidence(o) = sum_c [ sim(o, c) * reward(c) ] / (sum_c sim(o, c) + eps) Response embeddings are read from a pre-fetched cache (resp_emb_cache) rather than via individual DB calls; missing entries are skipped (sim treated as 0). \"\"\" if not candidate_memories or self.memory_db is None: return {} other_agents = [aid for aid in prev_responses if aid != self.agent_id] if not other_agents: return {} other_agent_strs = [str(aid) for aid in other_agents] precomputed = [prev_responses[aid][1] for aid in other_agents] if all(e is not None for e in precomputed): current_embs = np.stack(precomputed).astype(np.float32) else: embed_fn = self.memory_db.embedding_model.encode current_texts = [prev_responses[aid][0] for aid in other_agents] current_embs = np.array(embed_fn(current_texts), dtype=np.float32) weighted_sim_sum = np.zeros(len(other_agents)) sim_sum = np.zeros(len(other_agents)) for case_meta, _, _, _ in candidate_memories: hist_query_id = case_meta.get(\"question_id\") hist_round = int(case_meta.get(\"current_round\", 1)) - 1 if hist_round < 0: continue round_correct = 1.0 if case_meta.get(\"response_correct\", False) else 0.0 for i, aid_str in enumerate(other_agent_strs): hist_vec = resp_emb_cache.get((hist_query_id, hist_round, aid_str)) if hist_vec is None: continue sim = float(np.dot(current_embs[i], hist_vec)) sim = max(sim, 0.0) weighted_sim_sum[i] += sim * round_correct sim_sum[i] += sim eps = 1e-8 confidences = weighted_sim_sum / (sim_sum + eps) return {aid: float(confidences[i]) for i, aid in enumerate(other_agents)}","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent._build_query_text","name":"_build_query_text","kind":"method","path":"src/agent.py","span":"L383-L410","signature":"def _build_query_text( self, query: Query, prev_agent_response: Optional[str] = None, prev_round_summary: Optional[str] = None, ) -> str","exported":false,"origin":"ast","body_hash":"9ad1bd02ba32abf30c3f00c8e813584e632df1ebfd0a8ea40a96c6e12cc381ac","body_text":"def _build_query_text( self, query: Query, prev_agent_response: Optional[str] = None, prev_round_summary: Optional[str] = None, ) -> str: \"\"\" Mirror _build_embedding_text from memory_utils: constructs the debate-state string used as the retrieval key. 1. task (query content) 2. agent's previous response (from prev_responses, None at round 0) 3. previous round summary (None at round 0 or when summarization is off) \"\"\" embedding_text = \"\" # 1. task embedding_text += f\"{query.content}\\n\" # 2. agent previous response (None for round 0) if prev_agent_response is not None: embedding_text += f\"{prev_agent_response}\\n\" # 3. previous debate summary (None for round 0) if prev_round_summary is not None: embedding_text += f\"{prev_round_summary}\\n\" return embedding_text.strip()","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"},{"id":"src/agent.py#Agent._to_retrieved_memories","name":"_to_retrieved_memories","kind":"method","path":"src/agent.py","span":"L412-L427","signature":"def _to_retrieved_memories( self, query_id: int, retrieved: List[Tuple[Dict, str, float, Optional[np.ndarray]]], ) -> List[RetrievedMemory]","exported":false,"origin":"ast","body_hash":"cd944958b189782e12833531dc554478838aee07dcab7b1ab04426b3ea756999","body_text":"def _to_retrieved_memories( self, query_id: int, retrieved: List[Tuple[Dict, str, float, Optional[np.ndarray]]], ) -> List[RetrievedMemory]: \"\"\"Wrap the raw candidate tuples retrieved by a policy as RetrievedMemory.\"\"\" return [ RetrievedMemory( content=doc, query_id=query_id, agent_id=self.agent_id, is_correct=bool(meta.get(\"response_correct\", False)), similarity=float(max(0.0, 1.0 - distance)), ) for meta, doc, distance, _ in retrieved ]","summary_state":"pending","summary":null,"crux":null,"owner":"Agent"}],"rawEdges":[{"source":"src/agent.py","relation":"imports","specifier":"dataclasses","file":"src/agent.py"},{"source":"src/agent.py","relation":"imports","specifier":"typing","file":"src/agent.py"},{"source":"src/agent.py","relation":"imports","specifier":"src.memory","file":"src/agent.py"},{"source":"src/agent.py","relation":"imports","specifier":"src.retrieval_policy","file":"src/agent.py"},{"source":"src/agent.py","relation":"imports","specifier":"src.prompts","file":"src/agent.py"},{"source":"src/agent.py","relation":"contains","targetId":"src/agent.py#Query","file":"src/agent.py"},{"source":"src/agent.py#Query","relation":"contains","targetId":"src/agent.py#Query.__str__","file":"src/agent.py"},{"source":"src/agent.py#Query","relation":"contains","targetId":"src/agent.py#Query.__repr__","file":"src/agent.py"},{"source":"src/agent.py#Query.__repr__","relation":"calls","name":"__str__","viaMember":true,"file":"src/agent.py","recvType":"Query"},{"source":"src/agent.py","relation":"calls","name":"dataclass","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py","relation":"contains","targetId":"src/agent.py#AgentId","file":"src/agent.py"},{"source":"src/agent.py#AgentId","relation":"contains","targetId":"src/agent.py#AgentId.__str__","file":"src/agent.py"},{"source":"src/agent.py#AgentId","relation":"contains","targetId":"src/agent.py#AgentId.__repr__","file":"src/agent.py"},{"source":"src/agent.py#AgentId.__repr__","relation":"calls","name":"__str__","viaMember":true,"file":"src/agent.py","recvType":"AgentId"},{"source":"src/agent.py","relation":"contains","targetId":"src/agent.py#RetrievedMemory","file":"src/agent.py"},{"source":"src/agent.py#RetrievedMemory","relation":"contains","targetId":"src/agent.py#RetrievedMemory.__str__","file":"src/agent.py"},{"source":"src/agent.py#RetrievedMemory","relation":"contains","targetId":"src/agent.py#RetrievedMemory.__repr__","file":"src/agent.py"},{"source":"src/agent.py#RetrievedMemory.__repr__","relation":"calls","name":"__str__","viaMember":true,"file":"src/agent.py","recvType":"RetrievedMemory"},{"source":"src/agent.py","relation":"contains","targetId":"src/agent.py#Agent","file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.__init__","file":"src/agent.py"},{"source":"src/agent.py#Agent.__init__","relation":"calls","name":"str","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.__init__","relation":"calls","name":"int","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.__init__","relation":"calls","name":"build_retrieval_policy","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.__init__","relation":"calls","name":"int","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.initialize_conversation","file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.clear_conversation","file":"src/agent.py"},{"source":"src/agent.py#Agent.clear_conversation","relation":"calls","name":"pop","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.clear_conversation","relation":"calls","name":"pop","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.clear_all_conversations","file":"src/agent.py"},{"source":"src/agent.py#Agent.clear_all_conversations","relation":"calls","name":"clear","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.clear_all_conversations","relation":"calls","name":"clear","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent._build_prompt","file":"src/agent.py"},{"source":"src/agent.py#Agent._build_prompt","relation":"calls","name":"enumerate","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._build_prompt","relation":"calls","name":"items","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._build_prompt","relation":"calls","name":"enumerate","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._build_prompt","relation":"calls","name":"format","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._build_prompt","relation":"calls","name":"format","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.build_messages","file":"src/agent.py"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"batch_retrieve_all","viaMember":true,"file":"src/agent.py","recvType":"Agent"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"_build_prompt","viaMember":true,"file":"src/agent.py","recvType":"Agent"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"append","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.build_messages","relation":"calls","name":"extend","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.record_response","file":"src/agent.py"},{"source":"src/agent.py#Agent.record_response","relation":"calls","name":"append","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.record_response","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent.batch_retrieve_all","file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"_build_query_text","viaMember":true,"file":"src/agent.py","recvType":"Agent"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"append","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"batch_retrieve","viaMember":true,"file":"src/agent.py","recvType":"build_retrieval_policy"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"batch_query_similar","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"zip","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"str","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"keys","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"next","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"iter","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"values","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"values","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"int","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"append","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"batch_get_response_embeddings","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"update","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"items","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"_to_retrieved_memories","viaMember":true,"file":"src/agent.py","recvType":"Agent"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent.batch_retrieve_all","relation":"calls","name":"_estimate_confidences_cached","viaMember":true,"file":"src/agent.py","recvType":"Agent"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent._estimate_confidences_cached","file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"str","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"all","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"astype","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"stack","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"array","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"embed_fn","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"zeros","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"len","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"zeros","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"len","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"int","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"enumerate","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"float","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"dot","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"max","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"float","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._estimate_confidences_cached","relation":"calls","name":"enumerate","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent._build_query_text","file":"src/agent.py"},{"source":"src/agent.py#Agent._build_query_text","relation":"calls","name":"strip","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent","relation":"contains","targetId":"src/agent.py#Agent._to_retrieved_memories","file":"src/agent.py"},{"source":"src/agent.py#Agent._to_retrieved_memories","relation":"calls","name":"RetrievedMemory","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._to_retrieved_memories","relation":"calls","name":"bool","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._to_retrieved_memories","relation":"calls","name":"get","viaMember":true,"file":"src/agent.py"},{"source":"src/agent.py#Agent._to_retrieved_memories","relation":"calls","name":"float","viaMember":false,"file":"src/agent.py"},{"source":"src/agent.py#Agent._to_retrieved_memories","relation":"calls","name":"max","viaMember":false,"file":"src/agent.py"}]},"src/chain_of_thought.py":{"size":7247,"mtimeMs":1789362123400.8032,"hash":"1d7b5217b81cd9f001977935d45725028e20f46c4cc4719d28ee0b3714225e93","nodes":[{"id":"src/chain_of_thought.py","name":"chain_of_thought.py","kind":"file","path":"src/chain_of_thought.py","span":"L1-L174","signature":null,"exported":true,"origin":"ast","body_hash":"1d7b5217b81cd9f001977935d45725028e20f46c4cc4719d28ee0b3714225e93","chars":7247,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import re import json from collections import Counter from typing import Dict, List, Any from src.agent import Query, Agent from src.models import LanguageModel from src.checker import ANSWER_CHECKER"},{"id":"src/chain_of_thought.py#remove_think_tags","name":"remove_think_tags","kind":"function","path":"src/chain_of_thought.py","span":"L12-L13","signature":"def remove_think_tags(text: str) -> str","exported":true,"origin":"ast","body_hash":"f30d594c93fd7f89cc4a9f60dca524a516252b37114812bcd4767112b6116d3b","body_text":"def remove_think_tags(text: str) -> str: return re.sub(r\"<think>.*?</think>\", \"\", text, flags=re.DOTALL)","summary_state":"pending","summary":null,"crux":null},{"id":"src/chain_of_thought.py#ChainOfThought","name":"ChainOfThought","kind":"class","path":"src/chain_of_thought.py","span":"L16-L173","signature":"class ChainOfThought","exported":true,"origin":"ast","body_hash":"3380fdeb249ffddbcc5f0194b405cd794ec357dd8f1f8c58a8c4dd762ffb3e09","body_text":"class ChainOfThought: \"\"\" Chain-of-Thought / Self-Consistency inference wrapper. \"\"\" def __init__( self, agent: Agent, model: LanguageModel, task_type: str, num_reasoning_paths: int = 1, self_consistency: bool = False, ) -> None: self.agent = agent self.model = model self.num_reasoning_paths = num_reasoning_paths self.self_consistency = self_consistency self.task_type = task_type self._checker = ANSWER_CHECKER[task_type] self.cot_log = None def run(self, queries: List[Query]) -> List[Dict[str, Any]]: for query in queries: self.agent.initialize_conversation(query) # Optional similarity-based retrieval: fetch in-context examples for the # whole batch in one shot. CoT has no debate state, so prev_responses is # empty and retrieval falls back to a plain semantic query on the question # (Agent is built with retrieve_first_round=True so round 0 retrieves). prefetched_map: Dict[int, Any] = {} if self.agent.if_use_memory: prefetched_map = self.agent.batch_retrieve_all( queries=queries, prev_responses_map={query.query_id: {} for query in queries}, prev_round_summaries={query.query_id: None for query in queries}, prev_consensus_ratios={query.query_id: None for query in queries}, ) formated_messages = [ self._build_cot_messages(query, prefetched_map.get(query.query_id)) for query in queries ] final_results = [] if self.self_consistency and self.num_reasoning_paths > 1: flat_meta, sc_messages = [], [] for query, messages in zip(queries, formated_messages): sc_messages.extend([messages] * self.num_reasoning_paths) flat_meta.extend([(query, path_idx) for path_idx in range(self.num_reasoning_paths)]) model_output = self.model(sc_messages, task_name=\"self-consistency\") responses = [remove_think_tags(r).strip() for r in model_output[\"responses\"]] responses_by_qid = {query.query_id: [] for query in queries} for i, (query, path_idx) in enumerate(flat_meta): pred_answer = self._checker.parse_answer(responses[i]) is_correct = self._checker.check_answer(pred_answer, query.answer) responses_by_qid[query.query_id].append({ \"path_id\": path_idx, \"response\": responses[i], \"pred_answer\": pred_answer, \"is_correct\": is_correct, }) for query in queries: reasoning_paths = sorted(responses_by_qid[query.query_id], key=lambda x: x[\"path_id\"]) _, common_answer = self._check_consensus(reasoning_paths) is_final_correct = bool(self._checker.check_answer(common_answer, query.answer)) result = { \"query\": query.content, \"answer\": query.answer, \"reasoning_paths\": reasoning_paths, \"final_answer\": common_answer, \"is_final_correct\": is_final_correct, } self._attach_retrieved_memories(result, prefetched_map.get(query.query_id)) final_results.append(result) else: model_output = self.model(formated_messages, task_name=\"cot\") responses = [remove_think_tags(r).strip() for r in model_output[\"responses\"]] for query, response in zip(queries, responses): pred_answer = self._checker.parse_answer(response) is_final_correct = bool(self._checker.check_answer(pred_answer, query.answer)) result = { \"query\": query.content, \"answer\": query.answer, \"response\": response, \"final_answer\": pred_answer, \"is_final_correct\": is_final_correct, } self._attach_retrieved_memories(result, prefetched_map.get(query.query_id)) final_results.append(result) self.cot_log = final_results for query in queries: self.agent.clear_conversation(query) return final_results def _attach_retrieved_memories(self, result: Dict[str, Any], prefetched) -> None: \"\"\" Record the in-context memory cases retrieved for this query, mirroring the per-round `retrieved_memories` format in multi_agent_debate.py. Only added when memory is enabled; prefetched is (retrieved_memories, other_confidences). \"\"\" if not self.agent.if_use_memory: return retrieved, _ = prefetched if prefetched else (None, None) result[\"retrieved_memories\"] = [{ \"content\": rm.content, \"is_correct\": rm.is_correct, \"similarity\": rm.similarity, } for rm in (retrieved or [])] def _build_cot_messages(self, query: Query, prefetched=None) -> List[Dict[str, str]]: messages = self.agent.build_messages(query=query, prev_responses=None, prefetched=prefetched) # Keep the original task-specific answer format while nudging CoT style. messages[-1][\"content\"] += \"\\nPlease solve the problem step by step.\" return messages def _check_consensus(self, reasoning_paths: List[str], consensus_threshold: float = 0.8) -> tuple[bool, str]: \"\"\" Check if all agents have reached consensus on the same answer. Returns: is_consensus: bool common_answer: str (the answer that the majority agreed on, or None if no consensus) \"\"\" if not reasoning_paths: return False, \"\" normalized_answers = [output[\"pred_answer\"].lower().strip() for output in reasoning_paths] answer_correctness = [output[\"is_correct\"] for output in reasoning_paths] if sum(answer_correctness) == len(reasoning_paths): return True, normalized_answers[0] answer_counts = Counter(normalized_answers) common_answer, max_count = answer_counts.most_common(1)[0] consensus_ratio = max_count / len(reasoning_paths) is_consensu","summary_state":"pending","summary":null,"crux":null},{"id":"src/chain_of_thought.py#ChainOfThought.__init__","name":"__init__","kind":"method","path":"src/chain_of_thought.py","span":"L21-L37","signature":"def __init__( self, agent: Agent, model: LanguageModel, task_type: str, num_reasoning_paths: int = 1, self_consistency: bool = False, ) -> None","exported":false,"origin":"ast","body_hash":"e5c2d0706bdad2ba87998db08109871c8a55cbe0888456f461f598802dceba02","body_text":"def __init__( self, agent: Agent, model: LanguageModel, task_type: str, num_reasoning_paths: int = 1, self_consistency: bool = False, ) -> None: self.agent = agent self.model = model self.num_reasoning_paths = num_reasoning_paths self.self_consistency = self_consistency self.task_type = task_type self._checker = ANSWER_CHECKER[task_type] self.cot_log = None","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought.run","name":"run","kind":"method","path":"src/chain_of_thought.py","span":"L39-L120","signature":"def run(self, queries: List[Query]) -> List[Dict[str, Any]]","exported":true,"origin":"ast","body_hash":"bcfb75dbf11e18ff817023513e3ac145863c2cba53feccc312d8c361ed762242","body_text":"def run(self, queries: List[Query]) -> List[Dict[str, Any]]: for query in queries: self.agent.initialize_conversation(query) # Optional similarity-based retrieval: fetch in-context examples for the # whole batch in one shot. CoT has no debate state, so prev_responses is # empty and retrieval falls back to a plain semantic query on the question # (Agent is built with retrieve_first_round=True so round 0 retrieves). prefetched_map: Dict[int, Any] = {} if self.agent.if_use_memory: prefetched_map = self.agent.batch_retrieve_all( queries=queries, prev_responses_map={query.query_id: {} for query in queries}, prev_round_summaries={query.query_id: None for query in queries}, prev_consensus_ratios={query.query_id: None for query in queries}, ) formated_messages = [ self._build_cot_messages(query, prefetched_map.get(query.query_id)) for query in queries ] final_results = [] if self.self_consistency and self.num_reasoning_paths > 1: flat_meta, sc_messages = [], [] for query, messages in zip(queries, formated_messages): sc_messages.extend([messages] * self.num_reasoning_paths) flat_meta.extend([(query, path_idx) for path_idx in range(self.num_reasoning_paths)]) model_output = self.model(sc_messages, task_name=\"self-consistency\") responses = [remove_think_tags(r).strip() for r in model_output[\"responses\"]] responses_by_qid = {query.query_id: [] for query in queries} for i, (query, path_idx) in enumerate(flat_meta): pred_answer = self._checker.parse_answer(responses[i]) is_correct = self._checker.check_answer(pred_answer, query.answer) responses_by_qid[query.query_id].append({ \"path_id\": path_idx, \"response\": responses[i], \"pred_answer\": pred_answer, \"is_correct\": is_correct, }) for query in queries: reasoning_paths = sorted(responses_by_qid[query.query_id], key=lambda x: x[\"path_id\"]) _, common_answer = self._check_consensus(reasoning_paths) is_final_correct = bool(self._checker.check_answer(common_answer, query.answer)) result = { \"query\": query.content, \"answer\": query.answer, \"reasoning_paths\": reasoning_paths, \"final_answer\": common_answer, \"is_final_correct\": is_final_correct, } self._attach_retrieved_memories(result, prefetched_map.get(query.query_id)) final_results.append(result) else: model_output = self.model(formated_messages, task_name=\"cot\") responses = [remove_think_tags(r).strip() for r in model_output[\"responses\"]] for query, response in zip(queries, responses): pred_answer = self._checker.parse_answer(response) is_final_correct = bool(self._checker.check_answer(pred_answer, query.answer)) result = { \"query\": query.content, \"answer\": query.answer, \"response\": response, \"final_answer\": pred_answer, \"is_final_correct\": is_final_correct, } self._attach_retrieved_memories(result, prefetched_map.get(query.query_id)) final_results.append(result) self.cot_log = final_results for query in queries: self.agent.clear_conversation(query) return final_results","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought._attach_retrieved_memories","name":"_attach_retrieved_memories","kind":"method","path":"src/chain_of_thought.py","span":"L122-L135","signature":"def _attach_retrieved_memories(self, result: Dict[str, Any], prefetched) -> None","exported":false,"origin":"ast","body_hash":"1eed6781ad7a073a36a14119433637344f03fc868c750107b6c1f66b4872d9cc","body_text":"def _attach_retrieved_memories(self, result: Dict[str, Any], prefetched) -> None: \"\"\" Record the in-context memory cases retrieved for this query, mirroring the per-round `retrieved_memories` format in multi_agent_debate.py. Only added when memory is enabled; prefetched is (retrieved_memories, other_confidences). \"\"\" if not self.agent.if_use_memory: return retrieved, _ = prefetched if prefetched else (None, None) result[\"retrieved_memories\"] = [{ \"content\": rm.content, \"is_correct\": rm.is_correct, \"similarity\": rm.similarity, } for rm in (retrieved or [])]","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought._build_cot_messages","name":"_build_cot_messages","kind":"method","path":"src/chain_of_thought.py","span":"L137-L141","signature":"def _build_cot_messages(self, query: Query, prefetched=None) -> List[Dict[str, str]]","exported":false,"origin":"ast","body_hash":"73badf41b2c67bc010279a137f2b99d78fed18a9b467b4f5755f0918e2ce3f28","body_text":"def _build_cot_messages(self, query: Query, prefetched=None) -> List[Dict[str, str]]: messages = self.agent.build_messages(query=query, prev_responses=None, prefetched=prefetched) # Keep the original task-specific answer format while nudging CoT style. messages[-1][\"content\"] += \"\\nPlease solve the problem step by step.\" return messages","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought._check_consensus","name":"_check_consensus","kind":"method","path":"src/chain_of_thought.py","span":"L143-L165","signature":"def _check_consensus(self, reasoning_paths: List[str], consensus_threshold: float = 0.8) -> tuple[bool, str]","exported":false,"origin":"ast","body_hash":"cca7f2c7dff14a82a39cdd5e8f8c21fc37e804f5f3fa9f8da934824fa14b52ca","body_text":"def _check_consensus(self, reasoning_paths: List[str], consensus_threshold: float = 0.8) -> tuple[bool, str]: \"\"\" Check if all agents have reached consensus on the same answer. Returns: is_consensus: bool common_answer: str (the answer that the majority agreed on, or None if no consensus) \"\"\" if not reasoning_paths: return False, \"\" normalized_answers = [output[\"pred_answer\"].lower().strip() for output in reasoning_paths] answer_correctness = [output[\"is_correct\"] for output in reasoning_paths] if sum(answer_correctness) == len(reasoning_paths): return True, normalized_answers[0] answer_counts = Counter(normalized_answers) common_answer, max_count = answer_counts.most_common(1)[0] consensus_ratio = max_count / len(reasoning_paths) is_consensus = consensus_ratio >= consensus_threshold return is_consensus, common_answer","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought.save_cot_log","name":"save_cot_log","kind":"method","path":"src/chain_of_thought.py","span":"L167-L170","signature":"def save_cot_log(self, filename: str) -> None","exported":true,"origin":"ast","body_hash":"2df9d20e14037720acbc7094fa820882db72abf43d4c7edb5eeb50500464bb39","body_text":"def save_cot_log(self, filename: str) -> None: os.makedirs(os.path.dirname(filename), exist_ok=True) with open(filename, \"w\", encoding=\"utf-8\") as f: json.dump(self.cot_log, f, ensure_ascii=False, indent=2)","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"},{"id":"src/chain_of_thought.py#ChainOfThought.get_token_usage_summary","name":"get_token_usage_summary","kind":"method","path":"src/chain_of_thought.py","span":"L172-L173","signature":"def get_token_usage_summary(self)","exported":true,"origin":"ast","body_hash":"61cb8dd073485b509287cb1bf05553edd4cd4fde28c40e096d0616cd81dee054","body_text":"def get_token_usage_summary(self): return self.model.get_token_usage_summary()","summary_state":"pending","summary":null,"crux":null,"owner":"ChainOfThought"}],"rawEdges":[{"source":"src/chain_of_thought.py","relation":"imports","specifier":"os","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"re","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"json","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"collections","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"typing","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"src.agent","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"src.models","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"imports","specifier":"src.checker","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"contains","targetId":"src/chain_of_thought.py#remove_think_tags","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#remove_think_tags","relation":"calls","name":"sub","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought.__init__","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought.run","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"initialize_conversation","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"batch_retrieve_all","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"_build_cot_messages","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"get","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"zip","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"extend","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"extend","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"range","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"model","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"strip","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"remove_think_tags","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"enumerate","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"parse_answer","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"check_answer","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"append","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"sorted","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"_check_consensus","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"bool","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"check_answer","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"_attach_retrieved_memories","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"get","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"append","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"model","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"strip","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"remove_think_tags","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"zip","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"parse_answer","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"bool","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"check_answer","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"_attach_retrieved_memories","viaMember":true,"file":"src/chain_of_thought.py","recvType":"ChainOfThought"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"get","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"append","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.run","relation":"calls","name":"clear_conversation","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought._attach_retrieved_memories","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought._build_cot_messages","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._build_cot_messages","relation":"calls","name":"build_messages","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought._check_consensus","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"strip","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"lower","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"sum","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"len","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"Counter","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"most_common","viaMember":true,"file":"src/chain_of_thought.py","recvType":"Counter"},{"source":"src/chain_of_thought.py#ChainOfThought._check_consensus","relation":"calls","name":"len","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought.save_cot_log","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.save_cot_log","relation":"calls","name":"makedirs","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.save_cot_log","relation":"calls","name":"dirname","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.save_cot_log","relation":"calls","name":"open","viaMember":false,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.save_cot_log","relation":"calls","name":"dump","viaMember":true,"file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought","relation":"contains","targetId":"src/chain_of_thought.py#ChainOfThought.get_token_usage_summary","file":"src/chain_of_thought.py"},{"source":"src/chain_of_thought.py#ChainOfThought.get_token_usage_summary","relation":"calls","name":"get_token_usage_summary","viaMember":true,"file":"src/chain_of_thought.py"}]},"src/checker.py":{"size":9231,"mtimeMs":1789362123400.8032,"hash":"fa1d3b7f80055d94d509ad6e2cc312f8e4f63141caeb2bec37a4ab4e97a72f1a","nodes":[{"id":"src/checker.py","name":"checker.py","kind":"file","path":"src/checker.py","span":"L1-L283","signature":null,"exported":true,"origin":"ast","body_hash":"fa1d3b7f80055d94d509ad6e2cc312f8e4f63141caeb2bec37a4ab4e97a72f1a","chars":9220,"summary_state":"pending","summary":null,"crux":null,"body_text":"import re import regex from math import isclose from abc import ABC, abstractmethod try: from sympy import simplify, parse_expr from latex2sympy2 import latex2sympy LATEX2SYMPY_AVAILABLE = True except ImportError: LATEX2SYMPY_AVAILABLE = False # ============================================ # Utility functions # ============================================ ANSWER_CHECKER = { \"MATH500\": MATHChecker(), \"MMLUPro_Engineering\": MMLUProChecker(), \"MMLUPro_Economics\": MMLUProChecker(), \"TruthfulQA\": MMLUProChecker(), }"},{"id":"src/checker.py#BaseChecker","name":"BaseChecker","kind":"class","path":"src/checker.py","span":"L14-L21","signature":"class BaseChecker(ABC)","exported":true,"origin":"ast","body_hash":"0d3129c9856b932154ea2c30345a587057db487f0b28473c5bc018673464919d","body_text":"class BaseChecker(ABC): @abstractmethod def check_answer(self, pred_answer: str, expected_answer: str) -> float: pass @abstractmethod def parse_answer(self, response: str) -> str: pass","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#BaseChecker.check_answer","name":"check_answer","kind":"method","path":"src/checker.py","span":"L16-L17","signature":"def check_answer(self, pred_answer: str, expected_answer: str) -> float","exported":true,"origin":"ast","body_hash":"7af2c6a84df3e9f626bdcc28820f7dd1c3807b6565a7f954ab99a90f6005d3dc","body_text":"def check_answer(self, pred_answer: str, expected_answer: str) -> float: pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseChecker"},{"id":"src/checker.py#BaseChecker.parse_answer","name":"parse_answer","kind":"method","path":"src/checker.py","span":"L20-L21","signature":"def parse_answer(self, response: str) -> str","exported":true,"origin":"ast","body_hash":"ce4f06136d8a022965eee9fccd91d84373945ec06674601f462325b0131c28ec","body_text":"def parse_answer(self, response: str) -> str: pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseChecker"},{"id":"src/checker.py#MATHChecker","name":"MATHChecker","kind":"class","path":"src/checker.py","span":"L24-L120","signature":"class MATHChecker(BaseChecker)","exported":true,"origin":"ast","body_hash":"00da200cd4f0956f68bc1cd6854f8545de3570ad1e375c689f59099c23d0c5d3","body_text":"class MATHChecker(BaseChecker): def check_answer(self, pred_answer: str, expected_answer: str) -> bool: try: is_correct = self._math_equal(pred_answer, expected_answer) return 1.0 if is_correct else 0.0 except Exception as e: print(f\"Error in evaluation: {e}\") print(f\"Prediction: {pred_answer}\") print(f\"Ground Truth: {expected_answer}\") def parse_answer(self, response: str) -> str: def _find_box(string:str) -> str: res = regex.findall(r\"\\\\boxed\\{(.*)\\}\", string) if not res: res = regex.findall(r\"\\\\fbox\\{(.*)\\}\", string) if not res: return None return res[-1] if response is None: return \"\" box_format_ans = _find_box(response) if box_format_ans: return box_format_ans gsm8k_format_ans = re.findall(r\"####\\s*(-?[0-9\\.\\,]+)\", response) if gsm8k_format_ans: return gsm8k_format_ans[-1].replace(\",\", \"\").replace(\"$\", \"\") # Then try to find \"The answer is X\" format answer_pattern = re.findall(r'(?:[Tt]he(?:\\s+final)?(?:\\s+answer)?(?:\\s+is)?:?)\\s*([^\\n]+)', response) if answer_pattern: return answer_pattern[-1] ans_line = response.split(\"\\n\")[-1] return ans_line def _math_equal(self, pred_answer: str, expected_answer: str) -> bool: if pred_answer is None or expected_answer is None: return False pred_raw = str(pred_answer) reference_raw = str(expected_answer) pm_pattern = r\"(.+?)\\s*\\\\pm\\s*(.+)\" def _check_pm_match(pm_str, list_str): pm_match = re.match(pm_pattern, pm_str.strip()) if not pm_match: return False parts = re.split(r'[,;]', list_str) if len(parts) != 2: return False A = pm_match.group(1).strip() B = pm_match.group(2).strip() pm_set = {strip_string(f\"{A}+{B}\"), strip_string(f\"{A}-{B}\")} list_set = {strip_string(parts[0]), strip_string(parts[1])} return pm_set == list_set if r\"\\pm\" in reference_raw and _check_pm_match(reference_raw, pred_raw): return True if r\"\\pm\" in pred_raw and _check_pm_match(pred_raw, reference_raw): return True # Apply enhanced normalization to both prediction and reference prediction_str = strip_string(pred_raw) reference_str = strip_string(reference_raw) # Path 0: Direct string equality after normalization if prediction_str == reference_str: return True # Path 1: Numerical comparison if is_digit(prediction_str) and is_digit(reference_str): pred_float = parse_digits(prediction_str) ref_float = parse_digits(reference_str) if pred_float is not None and ref_float is not None and isclose(pred_float, ref_float, rel_tol=1e-4): return True # Path 2: Unordered Tuple/Vector comparison pred_parts = [p.strip() for p in re.split(r'[,;]', prediction_str.strip('()[]{}')) if p.strip()] ref_parts = [r.strip() for r in re.split(r'[,;]', reference_str.strip('()[]{}')) if r.strip()] if len(pred_parts) > 1 and len(pred_parts) == len(ref_parts): # Sort the parts before comparing to handle unordered lists if sorted(pred_parts) == sorted(ref_parts): return True # Path 3: Symbolic comparison if symbolic_equal(prediction_str, reference_str): return True return False","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#MATHChecker.check_answer","name":"check_answer","kind":"method","path":"src/checker.py","span":"L25-L32","signature":"def check_answer(self, pred_answer: str, expected_answer: str) -> bool","exported":true,"origin":"ast","body_hash":"f2d2545269cab6de50ab50d3b0eeb39157608596463bd99cd9a307c65b98b2d8","body_text":"def check_answer(self, pred_answer: str, expected_answer: str) -> bool: try: is_correct = self._math_equal(pred_answer, expected_answer) return 1.0 if is_correct else 0.0 except Exception as e: print(f\"Error in evaluation: {e}\") print(f\"Prediction: {pred_answer}\") print(f\"Ground Truth: {expected_answer}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MATHChecker"},{"id":"src/checker.py#MATHChecker.parse_answer","name":"parse_answer","kind":"method","path":"src/checker.py","span":"L34-L60","signature":"def parse_answer(self, response: str) -> str","exported":true,"origin":"ast","body_hash":"c2ce512cfe39074a72ca4a507b66ef9d8593afa11e010e0b4862fffe49df17ab","body_text":"def parse_answer(self, response: str) -> str: def _find_box(string:str) -> str: res = regex.findall(r\"\\\\boxed\\{(.*)\\}\", string) if not res: res = regex.findall(r\"\\\\fbox\\{(.*)\\}\", string) if not res: return None return res[-1] if response is None: return \"\" box_format_ans = _find_box(response) if box_format_ans: return box_format_ans gsm8k_format_ans = re.findall(r\"####\\s*(-?[0-9\\.\\,]+)\", response) if gsm8k_format_ans: return gsm8k_format_ans[-1].replace(\",\", \"\").replace(\"$\", \"\") # Then try to find \"The answer is X\" format answer_pattern = re.findall(r'(?:[Tt]he(?:\\s+final)?(?:\\s+answer)?(?:\\s+is)?:?)\\s*([^\\n]+)', response) if answer_pattern: return answer_pattern[-1] ans_line = response.split(\"\\n\")[-1] return ans_line","summary_state":"pending","summary":null,"crux":null,"owner":"MATHChecker"},{"id":"src/checker.py#MATHChecker.parse_answer._find_box","name":"_find_box","kind":"function","path":"src/checker.py","span":"L35-L41","signature":"def _find_box(string:str) -> str","exported":false,"origin":"ast","body_hash":"29f4927a7cfe184a6ccec608326836e33af50af98c79584ef55da337d32da637","body_text":"def _find_box(string:str) -> str: res = regex.findall(r\"\\\\boxed\\{(.*)\\}\", string) if not res: res = regex.findall(r\"\\\\fbox\\{(.*)\\}\", string) if not res: return None return res[-1]","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#MATHChecker._math_equal","name":"_math_equal","kind":"method","path":"src/checker.py","span":"L62-L120","signature":"def _math_equal(self, pred_answer: str, expected_answer: str) -> bool","exported":false,"origin":"ast","body_hash":"ee187ac2608b4296311d2aa8ba123b0cd421120020158ba5bf1c161f54cbe264","body_text":"def _math_equal(self, pred_answer: str, expected_answer: str) -> bool: if pred_answer is None or expected_answer is None: return False pred_raw = str(pred_answer) reference_raw = str(expected_answer) pm_pattern = r\"(.+?)\\s*\\\\pm\\s*(.+)\" def _check_pm_match(pm_str, list_str): pm_match = re.match(pm_pattern, pm_str.strip()) if not pm_match: return False parts = re.split(r'[,;]', list_str) if len(parts) != 2: return False A = pm_match.group(1).strip() B = pm_match.group(2).strip() pm_set = {strip_string(f\"{A}+{B}\"), strip_string(f\"{A}-{B}\")} list_set = {strip_string(parts[0]), strip_string(parts[1])} return pm_set == list_set if r\"\\pm\" in reference_raw and _check_pm_match(reference_raw, pred_raw): return True if r\"\\pm\" in pred_raw and _check_pm_match(pred_raw, reference_raw): return True # Apply enhanced normalization to both prediction and reference prediction_str = strip_string(pred_raw) reference_str = strip_string(reference_raw) # Path 0: Direct string equality after normalization if prediction_str == reference_str: return True # Path 1: Numerical comparison if is_digit(prediction_str) and is_digit(reference_str): pred_float = parse_digits(prediction_str) ref_float = parse_digits(reference_str) if pred_float is not None and ref_float is not None and isclose(pred_float, ref_float, rel_tol=1e-4): return True # Path 2: Unordered Tuple/Vector comparison pred_parts = [p.strip() for p in re.split(r'[,;]', prediction_str.strip('()[]{}')) if p.strip()] ref_parts = [r.strip() for r in re.split(r'[,;]', reference_str.strip('()[]{}')) if r.strip()] if len(pred_parts) > 1 and len(pred_parts) == len(ref_parts): # Sort the parts before comparing to handle unordered lists if sorted(pred_parts) == sorted(ref_parts): return True # Path 3: Symbolic comparison if symbolic_equal(prediction_str, reference_str): return True return False","summary_state":"pending","summary":null,"crux":null,"owner":"MATHChecker"},{"id":"src/checker.py#MATHChecker._math_equal._check_pm_match","name":"_check_pm_match","kind":"function","path":"src/checker.py","span":"L71-L86","signature":"def _check_pm_match(pm_str, list_str)","exported":false,"origin":"ast","body_hash":"a2b76e16bcd7605b16f5e13e3b082973b66a78828274e69aab5f0111428c076e","body_text":"def _check_pm_match(pm_str, list_str): pm_match = re.match(pm_pattern, pm_str.strip()) if not pm_match: return False parts = re.split(r'[,;]', list_str) if len(parts) != 2: return False A = pm_match.group(1).strip() B = pm_match.group(2).strip() pm_set = {strip_string(f\"{A}+{B}\"), strip_string(f\"{A}-{B}\")} list_set = {strip_string(parts[0]), strip_string(parts[1])} return pm_set == list_set","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#MMLUProChecker","name":"MMLUProChecker","kind":"class","path":"src/checker.py","span":"L123-L146","signature":"class MMLUProChecker(BaseChecker)","exported":true,"origin":"ast","body_hash":"ee10fd3dad68c09aa6a756b7c3356f7c6d8092491abaabe56ddd2dfb6d43f732","body_text":"class MMLUProChecker(BaseChecker): def check_answer(self, pred_answer: str, expected_answer: str) -> bool: pred_char = self._normalize_char(pred_answer) correct_char = self._normalize_char(expected_answer) return 1.0 if pred_char == correct_char else 0.0 def parse_answer(self, response: str) -> str: matches = re.findall(r\"\\(\\(([Ａ-ＪA-Ja-j])\\)\\)\", response) if len(matches) == 0: matches = re.findall(r\"\\(([Ａ-ＪA-Ja-j])\\)\", response) if len(matches) == 0: return \"\" return matches[-1].lower() def _normalize_char(self, char: str) -> str: if len(char) == 1 and 0xFF01 <= ord(char) <= 0xFF5E: char = chr(ord(char) - 0xFEE0) char = char.replace(\"：\", \":\") return char.strip().lower()[0] if char else \"\"","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#MMLUProChecker.check_answer","name":"check_answer","kind":"method","path":"src/checker.py","span":"L124-L128","signature":"def check_answer(self, pred_answer: str, expected_answer: str) -> bool","exported":true,"origin":"ast","body_hash":"4847e46a99cc37f6909153d6a742e076fc3b7c48c091da062c4219734ff96a14","body_text":"def check_answer(self, pred_answer: str, expected_answer: str) -> bool: pred_char = self._normalize_char(pred_answer) correct_char = self._normalize_char(expected_answer) return 1.0 if pred_char == correct_char else 0.0","summary_state":"pending","summary":null,"crux":null,"owner":"MMLUProChecker"},{"id":"src/checker.py#MMLUProChecker.parse_answer","name":"parse_answer","kind":"method","path":"src/checker.py","span":"L130-L138","signature":"def parse_answer(self, response: str) -> str","exported":true,"origin":"ast","body_hash":"2cdda48399b1cac5ca37bd23ab6e825d374c332017e3ce751cf4168c98c23755","body_text":"def parse_answer(self, response: str) -> str: matches = re.findall(r\"\\(\\(([Ａ-ＪA-Ja-j])\\)\\)\", response) if len(matches) == 0: matches = re.findall(r\"\\(([Ａ-ＪA-Ja-j])\\)\", response) if len(matches) == 0: return \"\" return matches[-1].lower()","summary_state":"pending","summary":null,"crux":null,"owner":"MMLUProChecker"},{"id":"src/checker.py#MMLUProChecker._normalize_char","name":"_normalize_char","kind":"method","path":"src/checker.py","span":"L140-L146","signature":"def _normalize_char(self, char: str) -> str","exported":false,"origin":"ast","body_hash":"78d9d90fcbd04685477d7dcb842b2f00c49abb7141f65c5b12a47cfa46347f63","body_text":"def _normalize_char(self, char: str) -> str: if len(char) == 1 and 0xFF01 <= ord(char) <= 0xFF5E: char = chr(ord(char) - 0xFEE0) char = char.replace(\"：\", \":\") return char.strip().lower()[0] if char else \"\"","summary_state":"pending","summary":null,"crux":null,"owner":"MMLUProChecker"},{"id":"src/checker.py#strip_string","name":"strip_string","kind":"function","path":"src/checker.py","span":"L162-L205","signature":"def strip_string(string: str) -> str","exported":true,"origin":"ast","body_hash":"9c50bb0ccdc2fdfe3365c49533fd32873b5d0fcf8d48a5acae1569d8df34f296","body_text":"def strip_string(string: str) -> str: \"\"\"Enhanced string normalization to handle units, LaTeX commands, and other inconsistencies.\"\"\" string = str(string).strip() # Remove LaTeX wrappers for text/units string = re.sub(r\"\\\\text\\{(.*?)\\}\", r\"\\1\", string) string = re.sub(r\"\\\\mbox\\{(.*?)\\}\", r\"\\1\", string) # Normalize and remove common symbols and units string = string.replace(\"^{\\\\circ}\", \"\") string = string.replace(\"\\\\circ\", \"\") string = string.replace(\"°\", \"\") string = string.replace(\"\\\\%\", \"\") string = string.replace(\"%\", \"\") string = re.sub(r\"\\s*degrees\", \"\", string, flags=re.IGNORECASE) # Remove common textual units (can be expanded) string = re.sub(r\"\\s*inches\\^2\", \"\", string, flags=re.IGNORECASE) string = re.sub(r\"\\s*square inches\", \"\", string, flags=re.IGNORECASE) # Standard replacements string = string.replace(\"\\n\", \"\") if string.endswith(\".\"): string = string[:-1] string = string.replace(\"\\\\!\", \"\") string = string.replace(\"\\\\ \", \" \") string = string.replace(\"\\\\\\\\\", \"\\\\\") string = string.replace(\"\\\\\\n\", \"\\\\\") string = string.replace(\"tfrac\", \"frac\") string = string.replace(\"dfrac\", \"frac\") string = string.replace(\"\\\\left\", \"\") string = string.replace(\"\\\\right\", \"\") # Remove dollar signs string = string.replace(\"$\", \"\") # Remove spaces and special characters last for clean comparison string = re.sub(r\" \", \"\", string) string = re.sub(\"\\u200b\", \"\", string) # Zero-width space string = _fix_fracs(string) string = _fix_sqrt(string) return string.strip()","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#_fix_fracs","name":"_fix_fracs","kind":"function","path":"src/checker.py","span":"L208-L237","signature":"def _fix_fracs(string)","exported":false,"origin":"ast","body_hash":"b05d1742f77b907f1bc8b2a89e4689e2933e14a22eff957215c80fa71f442301","body_text":"def _fix_fracs(string): \"\"\"Fix fraction formatting in LaTeX strings.\"\"\" substrs = string.split(\"\\\\frac\") new_str = substrs[0] if len(substrs) > 1: substrs = substrs[1:] for substr in substrs: new_str += \"\\\\frac\" if len(substr) > 0 and substr[0] == \"{\": new_str += substr else: try: assert len(substr) >= 2 except: return string a = substr[0] b = substr[1] if b != \"{\": if len(substr) > 2: post_substr = substr[2:] new_str += \"{\" + a + \"}{\" + b + \"}\" + post_substr else: new_str += \"{\" + a + \"}{\" + b + \"}\" else: if len(substr) > 2: post_substr = substr[2:] new_str += \"{\" + a + \"}\" + b + post_substr else: new_str += \"{\" + a + \"}\" + b return new_str","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#_fix_sqrt","name":"_fix_sqrt","kind":"function","path":"src/checker.py","span":"L240-L242","signature":"def _fix_sqrt(string)","exported":false,"origin":"ast","body_hash":"427ef733f8e6c30189048a4ea8d84b298ea1328bdfc6d2c22c48d3f86b99b70b","body_text":"def _fix_sqrt(string): \"\"\"Fix square root formatting in LaTeX strings.\"\"\" return re.sub(r\"\\\\sqrt(\\w+)\", r\"\\\\sqrt{\\1}\", string)","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#parse_digits","name":"parse_digits","kind":"function","path":"src/checker.py","span":"L245-L252","signature":"def parse_digits(num)","exported":true,"origin":"ast","body_hash":"a97053a8170fcf65eb2c0e55ed319d0ba7821310d39964008743d548d2e5011c","body_text":"def parse_digits(num): \"\"\"Parse numeric values from strings.\"\"\" num_str = strip_string(str(num)) num_str = regex.sub(\",\", \"\", num_str) try: return float(num_str) except: return None","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#is_digit","name":"is_digit","kind":"function","path":"src/checker.py","span":"L255-L257","signature":"def is_digit(num)","exported":true,"origin":"ast","body_hash":"69c42ca92a00d1eec04d0d45751c2eab2035ae23f32e5708388bc29fc5656284","body_text":"def is_digit(num): \"\"\"Check if a value can be parsed as a number.\"\"\" return parse_digits(num) is not None","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#symbolic_equal","name":"symbolic_equal","kind":"function","path":"src/checker.py","span":"L260-L282","signature":"def symbolic_equal(a, b)","exported":true,"origin":"ast","body_hash":"6c7a4e2bd2df84ced5fc9610c3aa27f7c9017969f6b4eb5602a2a8232a51340e","body_text":"def symbolic_equal(a, b): \"\"\"Check if two expressions are symbolically equal.\"\"\" if not LATEX2SYMPY_AVAILABLE: try: return simplify(a) == simplify(b) except: return False def _parse(s): try: return latex2sympy(s) except: try: return parse_expr(s) except: return s try: if simplify(_parse(a) - _parse(b)) == 0: return True except: pass return False","summary_state":"pending","summary":null,"crux":null},{"id":"src/checker.py#symbolic_equal._parse","name":"_parse","kind":"function","path":"src/checker.py","span":"L268-L275","signature":"def _parse(s)","exported":false,"origin":"ast","body_hash":"921308806091925f747fba0f40753f880af2a9d7edf5705bc2879f0a110b64a7","body_text":"def _parse(s): try: return latex2sympy(s) except: try: return parse_expr(s) except: return s","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"src/checker.py","relation":"imports","specifier":"re","file":"src/checker.py"},{"source":"src/checker.py","relation":"imports","specifier":"regex","file":"src/checker.py"},{"source":"src/checker.py","relation":"imports","specifier":"math","file":"src/checker.py"},{"source":"src/checker.py","relation":"imports","specifier":"abc","file":"src/checker.py"},{"source":"src/checker.py","relation":"imports","specifier":"sympy","file":"src/checker.py"},{"source":"src/checker.py","relation":"imports","specifier":"latex2sympy2","file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#BaseChecker","file":"src/checker.py"},{"source":"src/checker.py#BaseChecker","relation":"extends","name":"ABC","file":"src/checker.py"},{"source":"src/checker.py#BaseChecker","relation":"contains","targetId":"src/checker.py#BaseChecker.check_answer","file":"src/checker.py"},{"source":"src/checker.py#BaseChecker","relation":"contains","targetId":"src/checker.py#BaseChecker.parse_answer","file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#MATHChecker","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker","relation":"extends","name":"BaseChecker","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker","relation":"contains","targetId":"src/checker.py#MATHChecker.check_answer","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.check_answer","relation":"calls","name":"_math_equal","viaMember":true,"file":"src/checker.py","recvType":"MATHChecker"},{"source":"src/checker.py#MATHChecker.check_answer","relation":"calls","name":"print","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.check_answer","relation":"calls","name":"print","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.check_answer","relation":"calls","name":"print","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker","relation":"contains","targetId":"src/checker.py#MATHChecker.parse_answer","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"contains","targetId":"src/checker.py#MATHChecker.parse_answer._find_box","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer._find_box","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer._find_box","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"_find_box","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker.parse_answer","relation":"calls","name":"split","viaMember":true,"file":"src/checker.py","recvType":"str"},{"source":"src/checker.py#MATHChecker","relation":"contains","targetId":"src/checker.py#MATHChecker._math_equal","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"str","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"str","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"contains","targetId":"src/checker.py#MATHChecker._math_equal._check_pm_match","file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"match","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"split","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"group","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"group","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal._check_pm_match","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"_check_pm_match","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"_check_pm_match","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"is_digit","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"is_digit","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"parse_digits","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"parse_digits","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"isclose","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"split","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py","recvType":"strip_string"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"split","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py","recvType":"strip_string"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"sorted","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"sorted","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MATHChecker._math_equal","relation":"calls","name":"symbolic_equal","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#MMLUProChecker","file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker","relation":"extends","name":"BaseChecker","file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker","relation":"contains","targetId":"src/checker.py#MMLUProChecker.check_answer","file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.check_answer","relation":"calls","name":"_normalize_char","viaMember":true,"file":"src/checker.py","recvType":"MMLUProChecker"},{"source":"src/checker.py#MMLUProChecker.check_answer","relation":"calls","name":"_normalize_char","viaMember":true,"file":"src/checker.py","recvType":"MMLUProChecker"},{"source":"src/checker.py#MMLUProChecker","relation":"contains","targetId":"src/checker.py#MMLUProChecker.parse_answer","file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.parse_answer","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.parse_answer","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.parse_answer","relation":"calls","name":"findall","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.parse_answer","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker.parse_answer","relation":"calls","name":"lower","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker","relation":"contains","targetId":"src/checker.py#MMLUProChecker._normalize_char","file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"ord","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"chr","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"ord","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"chr"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"lower","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#MMLUProChecker._normalize_char","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py","recvType":"chr"},{"source":"src/checker.py","relation":"calls","name":"MATHChecker","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"calls","name":"MMLUProChecker","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"calls","name":"MMLUProChecker","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"calls","name":"MMLUProChecker","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#strip_string","file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"str","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"endswith","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"_fix_fracs","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"_fix_sqrt","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#strip_string","relation":"calls","name":"strip","viaMember":true,"file":"src/checker.py","recvType":"_fix_sqrt"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#_fix_fracs","file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"split","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#_fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#_fix_sqrt","file":"src/checker.py"},{"source":"src/checker.py#_fix_sqrt","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#parse_digits","file":"src/checker.py"},{"source":"src/checker.py#parse_digits","relation":"calls","name":"strip_string","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#parse_digits","relation":"calls","name":"str","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#parse_digits","relation":"calls","name":"sub","viaMember":true,"file":"src/checker.py"},{"source":"src/checker.py#parse_digits","relation":"calls","name":"float","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#is_digit","file":"src/checker.py"},{"source":"src/checker.py#is_digit","relation":"calls","name":"parse_digits","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py","relation":"contains","targetId":"src/checker.py#symbolic_equal","file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"calls","name":"simplify","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"calls","name":"simplify","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"contains","targetId":"src/checker.py#symbolic_equal._parse","file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal._parse","relation":"calls","name":"latex2sympy","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal._parse","relation":"calls","name":"parse_expr","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"calls","name":"simplify","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"calls","name":"_parse","viaMember":false,"file":"src/checker.py"},{"source":"src/checker.py#symbolic_equal","relation":"calls","name":"_parse","viaMember":false,"file":"src/checker.py"}]},"src/memory.py":{"size":15084,"mtimeMs":1789362123400.8032,"hash":"7d4950ed2c04a0e63bc17b12edd5bcaa2f20f41b26e36055f9258ef3f8227b9d","nodes":[{"id":"src/memory.py","name":"memory.py","kind":"file","path":"src/memory.py","span":"L1-L364","signature":null,"exported":true,"origin":"ast","body_hash":"7d4950ed2c04a0e63bc17b12edd5bcaa2f20f41b26e36055f9258ef3f8227b9d","chars":15084,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import chromadb import numpy as np from typing import List, Tuple, Dict, Callable, Optional from nltk.tokenize import word_tokenize from nltk.stem import SnowballStemmer from src.models import EmbeddingModel"},{"id":"src/memory.py#MemoryVectorDB","name":"MemoryVectorDB","kind":"class","path":"src/memory.py","span":"L11-L363","signature":"class MemoryVectorDB","exported":true,"origin":"ast","body_hash":"2944b94af69c968ddbdfa8c4f3777360a1df19187d5d3b552c0fd159c25f8419","body_text":"class MemoryVectorDB: def __init__(self, persistent_dir: str, collection_name: str, embedding_model: EmbeddingModel, tokenizer: Optional[Callable[[str], List[str]]] = word_tokenize, stemmer: Optional[Callable[[str], List[str]]] = SnowballStemmer(\"english\"), verbose: bool = True): self.client = chromadb.PersistentClient(path=persistent_dir) self.collection = self.client.get_or_create_collection( name=collection_name, embedding_function=None, metadata={\"hnsw:space\": \"cosine\"} ) self.resp_collection = self.client.get_or_create_collection( name=f\"{collection_name}_responses\", embedding_function=None, metadata={\"hnsw:space\": \"cosine\"}, ) self.embedding_model = embedding_model self.tokenizer = tokenizer self.stemmer = stemmer self.verbose = verbose # Pre-computed document embeddings: db_id -> float32 unit vector (1-D np.ndarray) # Loaded lazily via load_doc_embeddings(); None means not yet loaded. self._doc_embeddings: Optional[Dict[str, np.ndarray]] = None # ------------------------------------------------------------------ # Agent memory collection # ------------------------------------------------------------------ def load_doc_embeddings(self, npz_path: str) -> None: \"\"\"Load pre-computed document embeddings from a .npz file into memory.\"\"\" data = np.load(npz_path, allow_pickle=False) db_ids = data[\"db_ids\"].astype(str) embeddings = data[\"embeddings\"].astype(np.float32) # (N, dim) # Ensure unit norm (float16 storage may introduce tiny drift) norms = np.linalg.norm(embeddings, axis=1, keepdims=True) norms = np.where(norms == 0, 1.0, norms) embeddings = embeddings / norms self._doc_embeddings = {db_id: embeddings[i] for i, db_id in enumerate(db_ids)} if self.verbose: print(f\"Loaded {len(self._doc_embeddings)} document embeddings from {npz_path}\") def add_memory(self, db_id: str, embedding_text: str, document_text: str, add_meta_datas: Dict = None) -> None: try: assert isinstance(embedding_text, str), \"embedding_text must be str\" embedding = self.embedding_model.encode(embedding_text) meta_datas = dict() if add_meta_datas is not None: assert isinstance(add_meta_datas, dict), \"meta_datas must be Dict\" for key, val in add_meta_datas.items(): meta_datas[key] = val tokens = \" \".join([self.stemmer.stem(token) for token in self.tokenizer(embedding_text.lower())]) meta_datas[\"key_tokens\"] = tokens self.collection.upsert( embeddings=[embedding], documents=[document_text], metadatas=[meta_datas], ids=[db_id] ) if self.verbose: print(f\"Successfully added memory ID.: {db_id}\") except Exception as e: raise Exception(f\"Failed to add memory.: {str(e)}\") def add_batch_memories(self, db_ids: List[str], embedding_texts: List[str], document_texts: List[str], add_meta_datas: List[Dict] = None) -> None: try: assert len(db_ids) == len(embedding_texts) == len(document_texts), \\ \"db_ids, embedding_texts, and document_texts must have the same length\" embeddings = self.embedding_model.encode(embedding_texts, show_progress_bar=self.verbose) meta_datas_list = [] for i, emb_text in enumerate(embedding_texts): meta_datas = dict() if add_meta_datas is not None: assert isinstance(add_meta_datas[i], dict), \"each element of add_meta_datas must be Dict\" meta_datas.update(add_meta_datas[i]) tokens = \" \".join([self.stemmer.stem(token) for token in self.tokenizer(emb_text.lower())]) meta_datas[\"key_tokens\"] = tokens meta_datas_list.append(meta_datas) self.collection.upsert( embeddings=embeddings, documents=document_texts, metadatas=meta_datas_list, ids=db_ids, ) if self.verbose: print(f\"Successfully added {len(db_ids)} memories in batch.\") except Exception as e: raise Exception(f\"Failed to add memories in batch: {str(e)}\") def query_similar(self, query_text: str, n_results: int = 3, filter_metadata: Dict = None) -> List[Tuple[Dict, str, float, Optional[np.ndarray]]]: \"\"\" Returns a list of (metadata, document, distance, doc_embedding) tuples. doc_embedding is a float32 unit-norm np.ndarray when pre-computed embeddings are loaded via load_doc_embeddings(); otherwise None. \"\"\" try: query_embedding = self.embedding_model.encode(query_text) if filter_metadata is not None: assert isinstance(filter_metadata, dict), \"filter_metadata must be Dict\" results = self.collection.query( query_embeddings=[query_embedding], n_results=n_results, include=[\"metadatas\", \"documents\", \"distances\"], where=filter_metadata ) similar_pairs = [] for metadata, document, distance, db_id in zip( results['metadatas'][0], results['documents'][0], results['distances'][0], results['ids'][0], ): doc_emb = ( self._doc_embeddings.get(db_id) if self._doc_embeddings is not None else None ) similar_pairs.append((metadata, document, distance, doc_emb)) return similar_pairs except Exception as e: raise Exception(f\"Query failed.: {str(e)}\") def batch_query_similar( self, query_texts: List[str], n_results: int = 3, filter_metadata: Dict = None, ) -> List[List[Tuple[Dict, str, float, Optional[np.ndarray]]]]: \"\"\" Batch version of query_similar: encode all query_texts in one call, then issue a sin","summary_state":"pending","summary":null,"crux":null},{"id":"src/memory.py#MemoryVectorDB.__init__","name":"__init__","kind":"method","path":"src/memory.py","span":"L12-L35","signature":"def __init__(self, persistent_dir: str, collection_name: str, embedding_model: EmbeddingModel, tokenizer: Optional[Callable[[str], List[str]]] = word_tokenize, stemmer: Optional[Callable[[str], List[str]]] = SnowballStemmer(\"english\"), verbose: bool = True)","exported":false,"origin":"ast","body_hash":"acb50ec491d9b7c5ba9215f040630444b2668533e6a035b4050160706a9b1bc8","body_text":"def __init__(self, persistent_dir: str, collection_name: str, embedding_model: EmbeddingModel, tokenizer: Optional[Callable[[str], List[str]]] = word_tokenize, stemmer: Optional[Callable[[str], List[str]]] = SnowballStemmer(\"english\"), verbose: bool = True): self.client = chromadb.PersistentClient(path=persistent_dir) self.collection = self.client.get_or_create_collection( name=collection_name, embedding_function=None, metadata={\"hnsw:space\": \"cosine\"} ) self.resp_collection = self.client.get_or_create_collection( name=f\"{collection_name}_responses\", embedding_function=None, metadata={\"hnsw:space\": \"cosine\"}, ) self.embedding_model = embedding_model self.tokenizer = tokenizer self.stemmer = stemmer self.verbose = verbose # Pre-computed document embeddings: db_id -> float32 unit vector (1-D np.ndarray) # Loaded lazily via load_doc_embeddings(); None means not yet loaded. self._doc_embeddings: Optional[Dict[str, np.ndarray]] = None","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.load_doc_embeddings","name":"load_doc_embeddings","kind":"method","path":"src/memory.py","span":"L41-L54","signature":"def load_doc_embeddings(self, npz_path: str) -> None","exported":true,"origin":"ast","body_hash":"51445f953293705064187ec30c4fa82104ca306c31b7e919b23f6646dfab243b","body_text":"def load_doc_embeddings(self, npz_path: str) -> None: \"\"\"Load pre-computed document embeddings from a .npz file into memory.\"\"\" data = np.load(npz_path, allow_pickle=False) db_ids = data[\"db_ids\"].astype(str) embeddings = data[\"embeddings\"].astype(np.float32) # (N, dim) # Ensure unit norm (float16 storage may introduce tiny drift) norms = np.linalg.norm(embeddings, axis=1, keepdims=True) norms = np.where(norms == 0, 1.0, norms) embeddings = embeddings / norms self._doc_embeddings = {db_id: embeddings[i] for i, db_id in enumerate(db_ids)} if self.verbose: print(f\"Loaded {len(self._doc_embeddings)} document embeddings from {npz_path}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.add_memory","name":"add_memory","kind":"method","path":"src/memory.py","span":"L56-L80","signature":"def add_memory(self, db_id: str, embedding_text: str, document_text: str, add_meta_datas: Dict = None) -> None","exported":true,"origin":"ast","body_hash":"5937f923edcd23ed3f9b2b60dd1d59d778f7355d7a64a2d1b989b2d362b91fff","body_text":"def add_memory(self, db_id: str, embedding_text: str, document_text: str, add_meta_datas: Dict = None) -> None: try: assert isinstance(embedding_text, str), \"embedding_text must be str\" embedding = self.embedding_model.encode(embedding_text) meta_datas = dict() if add_meta_datas is not None: assert isinstance(add_meta_datas, dict), \"meta_datas must be Dict\" for key, val in add_meta_datas.items(): meta_datas[key] = val tokens = \" \".join([self.stemmer.stem(token) for token in self.tokenizer(embedding_text.lower())]) meta_datas[\"key_tokens\"] = tokens self.collection.upsert( embeddings=[embedding], documents=[document_text], metadatas=[meta_datas], ids=[db_id] ) if self.verbose: print(f\"Successfully added memory ID.: {db_id}\") except Exception as e: raise Exception(f\"Failed to add memory.: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.add_batch_memories","name":"add_batch_memories","kind":"method","path":"src/memory.py","span":"L82-L109","signature":"def add_batch_memories(self, db_ids: List[str], embedding_texts: List[str], document_texts: List[str], add_meta_datas: List[Dict] = None) -> None","exported":true,"origin":"ast","body_hash":"487d0945731336e9afac72c0e67f91d5881bc6d41340afd031b9e24a7aecc438","body_text":"def add_batch_memories(self, db_ids: List[str], embedding_texts: List[str], document_texts: List[str], add_meta_datas: List[Dict] = None) -> None: try: assert len(db_ids) == len(embedding_texts) == len(document_texts), \\ \"db_ids, embedding_texts, and document_texts must have the same length\" embeddings = self.embedding_model.encode(embedding_texts, show_progress_bar=self.verbose) meta_datas_list = [] for i, emb_text in enumerate(embedding_texts): meta_datas = dict() if add_meta_datas is not None: assert isinstance(add_meta_datas[i], dict), \"each element of add_meta_datas must be Dict\" meta_datas.update(add_meta_datas[i]) tokens = \" \".join([self.stemmer.stem(token) for token in self.tokenizer(emb_text.lower())]) meta_datas[\"key_tokens\"] = tokens meta_datas_list.append(meta_datas) self.collection.upsert( embeddings=embeddings, documents=document_texts, metadatas=meta_datas_list, ids=db_ids, ) if self.verbose: print(f\"Successfully added {len(db_ids)} memories in batch.\") except Exception as e: raise Exception(f\"Failed to add memories in batch: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.query_similar","name":"query_similar","kind":"method","path":"src/memory.py","span":"L111-L144","signature":"def query_similar(self, query_text: str, n_results: int = 3, filter_metadata: Dict = None) -> List[Tuple[Dict, str, float, Optional[np.ndarray]]]","exported":true,"origin":"ast","body_hash":"dfe13b4695039b7f9a90b9a1b05680ec8c926713adecd69428404c6aaec0906f","body_text":"def query_similar(self, query_text: str, n_results: int = 3, filter_metadata: Dict = None) -> List[Tuple[Dict, str, float, Optional[np.ndarray]]]: \"\"\" Returns a list of (metadata, document, distance, doc_embedding) tuples. doc_embedding is a float32 unit-norm np.ndarray when pre-computed embeddings are loaded via load_doc_embeddings(); otherwise None. \"\"\" try: query_embedding = self.embedding_model.encode(query_text) if filter_metadata is not None: assert isinstance(filter_metadata, dict), \"filter_metadata must be Dict\" results = self.collection.query( query_embeddings=[query_embedding], n_results=n_results, include=[\"metadatas\", \"documents\", \"distances\"], where=filter_metadata ) similar_pairs = [] for metadata, document, distance, db_id in zip( results['metadatas'][0], results['documents'][0], results['distances'][0], results['ids'][0], ): doc_emb = ( self._doc_embeddings.get(db_id) if self._doc_embeddings is not None else None ) similar_pairs.append((metadata, document, distance, doc_emb)) return similar_pairs except Exception as e: raise Exception(f\"Query failed.: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.batch_query_similar","name":"batch_query_similar","kind":"method","path":"src/memory.py","span":"L146-L189","signature":"def batch_query_similar( self, query_texts: List[str], n_results: int = 3, filter_metadata: Dict = None, ) -> List[List[Tuple[Dict, str, float, Optional[np.ndarray]]]]","exported":true,"origin":"ast","body_hash":"a4c925a97c4da30b077d94e5b66467e8cafed4fcb3ad7e4654a296be1a11adff","body_text":"def batch_query_similar( self, query_texts: List[str], n_results: int = 3, filter_metadata: Dict = None, ) -> List[List[Tuple[Dict, str, float, Optional[np.ndarray]]]]: \"\"\" Batch version of query_similar: encode all query_texts in one call, then issue a single ChromaDB query with multiple query embeddings. Returns one result list per query_text (same order). \"\"\" try: query_embeddings = self.embedding_model.encode(query_texts) if filter_metadata is not None: assert isinstance(filter_metadata, dict), \"filter_metadata must be Dict\" results = self.collection.query( query_embeddings=query_embeddings, n_results=n_results, include=[\"metadatas\", \"documents\", \"distances\"], where=filter_metadata, ) all_results = [] for i in range(len(query_texts)): similar_pairs = [] for metadata, document, distance, db_id in zip( results[\"metadatas\"][i], results[\"documents\"][i], results[\"distances\"][i], results[\"ids\"][i], ): doc_emb = ( self._doc_embeddings.get(db_id) if self._doc_embeddings is not None else None ) similar_pairs.append((metadata, document, distance, doc_emb)) all_results.append(similar_pairs) return all_results except Exception as e: raise Exception(f\"Batch query failed: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.batch_get_random","name":"batch_get_random","kind":"method","path":"src/memory.py","span":"L191-L241","signature":"def batch_get_random( self, n_queries: int, n_results: int = 3, filter_metadata: Dict = None, distance: float = 1.0, seed: Optional[int] = None, ) -> List[List[Tuple[Dict, str, float, Optional[np.ndarray]]]]","exported":true,"origin":"ast","body_hash":"7d3a7a493e060d21fde9de24e0595a53a9216802d2a73e3cbd8edaf87f9bbcfa","body_text":"def batch_get_random( self, n_queries: int, n_results: int = 3, filter_metadata: Dict = None, distance: float = 1.0, seed: Optional[int] = None, ) -> List[List[Tuple[Dict, str, float, Optional[np.ndarray]]]]: \"\"\" Draw n_queries independent random samples (without replacement) of n_results cases each from the collection, optionally restricted by filter_metadata. Fetches the matching pool once, then samples in Python. Returns one candidate list per query in the same 4-tuple shape as batch_query_similar: (metadata, document, distance, doc_embedding). `distance` is a constant placeholder (random cases have no similarity to a query). \"\"\" try: # Independent per-query RNG streams: sequential draws from one # Generator can collide for some seeds, so spawn a child stream per # query (reproducible given `seed`, distinct across queries). child_seeds = np.random.SeedSequence(seed).spawn(n_queries) pool = self.collection.get( where=filter_metadata, include=[\"metadatas\", \"documents\"], ) ids = pool[\"ids\"] metas = pool[\"metadatas\"] docs = pool[\"documents\"] all_results = [] for q in range(n_queries): candidates = [] if ids: rng = np.random.default_rng(child_seeds[q]) k = min(n_results, len(ids)) idxs = rng.choice(len(ids), size=k, replace=False) for j in idxs: doc_emb = ( self._doc_embeddings.get(ids[j]) if self._doc_embeddings is not None else None ) candidates.append((metas[j], docs[j], distance, doc_emb)) all_results.append(candidates) return all_results except Exception as e: raise Exception(f\"Random fetch failed: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.get_by_ids","name":"get_by_ids","kind":"method","path":"src/memory.py","span":"L243-L253","signature":"def get_by_ids(self, select_ids: List[str]) -> List[Tuple[Dict, str]]","exported":true,"origin":"ast","body_hash":"a8e8c23d08b6c65dcff405fdec755d165ec2d588edc0a93f8f0eb3f85e1d8228","body_text":"def get_by_ids(self, select_ids: List[str]) -> List[Tuple[Dict, str]]: try: results = self.collection.get(ids=select_ids, include=['metadatas', 'documents']) similar_pairs = [] for metadata, document in zip(results['metadatas'], results['documents']): similar_pairs.append((metadata, document)) return similar_pairs except Exception as e: raise Exception(f\"Query by IDs failed: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.update_memory","name":"update_memory","kind":"method","path":"src/memory.py","span":"L255-L281","signature":"def update_memory(self, db_id: str, embedding_text: str, document_text: str, add_meta_datas: Dict = None) -> None","exported":true,"origin":"ast","body_hash":"dd9f179d9957ff8374616416ace805c3f445e96c38f9a3f88e1c10035e260fc0","body_text":"def update_memory(self, db_id: str, embedding_text: str, document_text: str, add_meta_datas: Dict = None) -> None: try: assert isinstance(embedding_text, str), \"embedding_text must be str\" embedding = self.embedding_model.encode(embedding_text) meta_datas = dict() if add_meta_datas is not None: assert isinstance(add_meta_datas, dict), \"meta_datas must be Dict\" for key, val in add_meta_datas.items(): meta_datas[key] = val tokens = \" \".join([self.stemmer.stem(token) for token in self.tokenizer(embedding_text.lower())]) meta_datas[\"key_tokens\"] = tokens self.collection.update( embeddings=[embedding], documents=[document_text], metadatas=[meta_datas], ids=[db_id] ) if self.verbose: print(f\"Successfully updated memory ID: {db_id}\") except Exception as e: raise Exception(f\"Failed to update memory: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.delete_memory","name":"delete_memory","kind":"method","path":"src/memory.py","span":"L283-L289","signature":"def delete_memory(self, db_id: str) -> None","exported":true,"origin":"ast","body_hash":"8e7f1c5b95dc8718549aed91f28de1525cfd61c86941128d1a04eadda705d186","body_text":"def delete_memory(self, db_id: str) -> None: try: self.collection.delete(ids=[db_id]) if self.verbose: print(f\"Successfully deleted memory ID: {db_id}\") except Exception as e: raise Exception(f\"Failed to delete memory: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.get_memories_count","name":"get_memories_count","kind":"method","path":"src/memory.py","span":"L291-L292","signature":"def get_memories_count(self) -> int","exported":true,"origin":"ast","body_hash":"99868904ae61a73a17e82f1b60892b693030f85e4069d32f3c2a5e29bbd815cf","body_text":"def get_memories_count(self) -> int: return self.collection.count()","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.get_all_ids","name":"get_all_ids","kind":"method","path":"src/memory.py","span":"L294-L297","signature":"def get_all_ids(self, filter_metadata=None)","exported":true,"origin":"ast","body_hash":"6edc83c5017f9733e022a436680fdd9a9fc4aaa990a9985fbd199bfc9efd9ad3","body_text":"def get_all_ids(self, filter_metadata=None): data = self.collection.get(include=[\"metadatas\"], where=filter_metadata) ids = data['ids'] return ids","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB._make_resp_id","name":"_make_resp_id","kind":"method","path":"src/memory.py","span":"L304-L305","signature":"def _make_resp_id(query_id, round_idx: int, agent_id: str) -> str","exported":false,"origin":"ast","body_hash":"1ab18a4a85b0c40294caf467541129418f3436202d2da62f29ccfd597a4930e8","body_text":"def _make_resp_id(query_id, round_idx: int, agent_id: str) -> str: return f\"Q{query_id}_Round{round_idx}_{agent_id}\"","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.add_batch_responses","name":"add_batch_responses","kind":"method","path":"src/memory.py","span":"L307-L317","signature":"def add_batch_responses(self, db_ids: List[str], texts: List[str], meta_datas: List[Dict]) -> None","exported":true,"origin":"ast","body_hash":"bb07805fde2960a4e88921b05e1b0d48604da93a716096e640108743f1d8fe86","body_text":"def add_batch_responses(self, db_ids: List[str], texts: List[str], meta_datas: List[Dict]) -> None: embeddings = self.embedding_model.encode(texts, show_progress_bar=self.verbose) self.resp_collection.upsert( embeddings=embeddings, documents=texts, metadatas=meta_datas, ids=db_ids, ) if self.verbose: print(f\"Successfully added {len(db_ids)} response embeddings in batch.\")","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.get_response_embeddings","name":"get_response_embeddings","kind":"method","path":"src/memory.py","span":"L319-L339","signature":"def get_response_embeddings(self, query_id, round_idx: int, agent_ids: List[str]) -> Dict[str, Optional[np.ndarray]]","exported":true,"origin":"ast","body_hash":"9014f189d2fb3cff02580425fc13154eb3b5d39f2525c93b63ad278cbd6e349e","body_text":"def get_response_embeddings(self, query_id, round_idx: int, agent_ids: List[str]) -> Dict[str, Optional[np.ndarray]]: \"\"\" Return {agent_id: unit-norm embedding} for the given agents at a specific (query_id, round_idx). Missing entries map to None. \"\"\" ids = [self._make_resp_id(query_id, round_idx, aid) for aid in agent_ids] try: results = self.resp_collection.get(ids=ids, include=[\"embeddings\"]) except Exception: return {aid: None for aid in agent_ids} id_to_emb: Dict[str, np.ndarray] = { db_id: np.array(emb, dtype=np.float32) for db_id, emb in zip(results[\"ids\"], results[\"embeddings\"]) } out = {} for aid, db_id in zip(agent_ids, ids): out[aid] = id_to_emb.get(db_id) return out","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"},{"id":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","name":"batch_get_response_embeddings","kind":"method","path":"src/memory.py","span":"L341-L363","signature":"def batch_get_response_embeddings( self, keys: List[Tuple], # list of (query_id, round_idx, agent_id_str) ) -> Dict[Tuple, Optional[np.ndarray]]","exported":true,"origin":"ast","body_hash":"b5a43ee21967e296eb3fee2ad68f096fa05c79c1a99083b063f208e7bf6d1017","body_text":"def batch_get_response_embeddings( self, keys: List[Tuple], # list of (query_id, round_idx, agent_id_str) ) -> Dict[Tuple, Optional[np.ndarray]]: \"\"\" Fetch response embeddings for many (query_id, round_idx, agent_id) triples in a single ChromaDB get call. Returns {(query_id, round_idx, agent_id): vec}. \"\"\" ids = [self._make_resp_id(qid, r, aid) for qid, r, aid in keys] try: results = self.resp_collection.get(ids=ids, include=[\"embeddings\"]) except Exception: return {k: None for k in keys} id_to_emb: Dict[str, np.ndarray] = { db_id: np.array(emb, dtype=np.float32) for db_id, emb in zip(results[\"ids\"], results[\"embeddings\"]) } out = {} for key, db_id in zip(keys, ids): out[key] = id_to_emb.get(db_id) return out","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryVectorDB"}],"rawEdges":[{"source":"src/memory.py","relation":"imports","specifier":"os","file":"src/memory.py"},{"source":"src/memory.py","relation":"imports","specifier":"chromadb","file":"src/memory.py"},{"source":"src/memory.py","relation":"imports","specifier":"typing","file":"src/memory.py"},{"source":"src/memory.py","relation":"imports","specifier":"nltk.tokenize","file":"src/memory.py"},{"source":"src/memory.py","relation":"imports","specifier":"nltk.stem","file":"src/memory.py"},{"source":"src/memory.py","relation":"imports","specifier":"src.models","file":"src/memory.py"},{"source":"src/memory.py","relation":"contains","targetId":"src/memory.py#MemoryVectorDB","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.__init__","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.__init__","relation":"calls","name":"SnowballStemmer","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.__init__","relation":"calls","name":"PersistentClient","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.__init__","relation":"calls","name":"get_or_create_collection","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.__init__","relation":"calls","name":"get_or_create_collection","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.load_doc_embeddings","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"load","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"astype","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"astype","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"norm","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"where","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"enumerate","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.load_doc_embeddings","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.add_memory","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"dict","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"items","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"join","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"stem","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"tokenizer","viaMember":true,"file":"src/memory.py","recvType":"MemoryVectorDB"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"lower","viaMember":true,"file":"src/memory.py","recvType":"str"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"upsert","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_memory","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.add_batch_memories","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"enumerate","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"dict","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"update","viaMember":true,"file":"src/memory.py","recvType":"dict"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"join","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"stem","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"tokenizer","viaMember":true,"file":"src/memory.py","recvType":"MemoryVectorDB"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"lower","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"upsert","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_memories","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.query_similar","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"query","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.query_similar","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.batch_query_similar","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"query","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"range","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_query_similar","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.batch_get_random","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"spawn","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"SeedSequence","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"range","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"default_rng","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"min","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"choice","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_random","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.get_by_ids","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_by_ids","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_by_ids","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_by_ids","relation":"calls","name":"append","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_by_ids","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_by_ids","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.update_memory","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"dict","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"isinstance","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"items","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"join","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"stem","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"tokenizer","viaMember":true,"file":"src/memory.py","recvType":"MemoryVectorDB"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"lower","viaMember":true,"file":"src/memory.py","recvType":"str"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"update","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.update_memory","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.delete_memory","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.delete_memory","relation":"calls","name":"delete","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.delete_memory","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.delete_memory","relation":"calls","name":"Exception","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.delete_memory","relation":"calls","name":"str","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.get_memories_count","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_memories_count","relation":"calls","name":"count","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.get_all_ids","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_all_ids","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB._make_resp_id","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.add_batch_responses","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_responses","relation":"calls","name":"encode","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_responses","relation":"calls","name":"upsert","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_responses","relation":"calls","name":"print","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.add_batch_responses","relation":"calls","name":"len","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.get_response_embeddings","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"_make_resp_id","viaMember":true,"file":"src/memory.py","recvType":"MemoryVectorDB"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"array","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.get_response_embeddings","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB","relation":"contains","targetId":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"_make_resp_id","viaMember":true,"file":"src/memory.py","recvType":"MemoryVectorDB"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"array","viaMember":true,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"zip","viaMember":false,"file":"src/memory.py"},{"source":"src/memory.py#MemoryVectorDB.batch_get_response_embeddings","relation":"calls","name":"get","viaMember":true,"file":"src/memory.py"}]},"src/models.py":{"size":13629,"mtimeMs":1789362123400.8032,"hash":"713ccfac57a8375b54ac9b4b9cecad19d7288c01785a23505c54d1870a6d4c35","nodes":[{"id":"src/models.py","name":"models.py","kind":"file","path":"src/models.py","span":"L1-L339","signature":null,"exported":true,"origin":"ast","body_hash":"713ccfac57a8375b54ac9b4b9cecad19d7288c01785a23505c54d1870a6d4c35","chars":13629,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import time import random import json from concurrent.futures import ThreadPoolExecutor from typing import List, Dict, Any, Union, Optional, Tuple from src.utils import LLMConfig, EmbeddingConfig, TokenUsageTracker"},{"id":"src/models.py#build_language_model","name":"build_language_model","kind":"function","path":"src/models.py","span":"L11-L15","signature":"def build_language_model(llm_config: LLMConfig)","exported":true,"origin":"ast","body_hash":"dc33313a41b73c63ebeb9cf2c9384a9a525441eb2183f13c1f97e7591ba6397a","body_text":"def build_language_model(llm_config: LLMConfig): \"\"\"Return the language model implementation for the configured backend.\"\"\" if llm_config.backend == \"openrouter\": return OpenRouterLanguageModel(llm_config) return LanguageModel(llm_config)","summary_state":"pending","summary":null,"crux":null},{"id":"src/models.py#LanguageModel","name":"LanguageModel","kind":"class","path":"src/models.py","span":"L18-L120","signature":"class LanguageModel","exported":true,"origin":"ast","body_hash":"1f771dfd3dbdba9eabb25623541b41e84d342b0a2388daab17f65ba14504be3b","body_text":"class LanguageModel: \"\"\"Wraps a vLLM model; accepts message lists and returns structured results.\"\"\" def __init__(self, llm_config: LLMConfig): import torch from vllm import LLM, SamplingParams self.llm_config = llm_config self.model_name = llm_config.name or llm_config.model.split(\"/\")[-1] llm_name_path = ( self.llm_config.model_path if self.llm_config.model_path not in (None, \"None\") else self.llm_config.model ) print(f\">>>>>> LLM name path: {llm_name_path}\") self.llm = LLM( model=llm_name_path, tensor_parallel_size=getattr(self.llm_config, \"tensor_parallel_size\", 1), dtype=torch.bfloat16, max_model_len=getattr(self.llm_config, \"max_token_length\", 24576), gpu_memory_utilization=0.85, ) self.sampling_params = SamplingParams( temperature=getattr(self.llm_config, \"temperature\", 1), top_p=getattr(self.llm_config, \"top_p\", 1), max_tokens=getattr(self.llm_config, \"max_tokens\", 6144), stop_token_ids=[self.llm.get_tokenizer().eos_token_id], ) self.token_usage_tracker = TokenUsageTracker() def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]: \"\"\" Args: batch_messages: list of message lists, each message list is [{\"role\": ..., \"content\": ...}, ...] task_name: label used to group token usage statistics Returns dict with keys: results, cur_batch_input_tokens, cur_batch_output_tokens \"\"\" tokenizer = self.llm.get_tokenizer() tokenized_texts = tokenizer.apply_chat_template( batch_messages, tokenize=False, add_generation_prompt=True ) # Handle over-long queries max_len = self.llm.llm_engine.model_config.max_model_len valid_indices, valid_texts = [], [] for i, text in enumerate(tokenized_texts): token_len = len(tokenizer.encode(text)) if token_len > max_len: print( f\"[WARNING] [{task_name}] Skipping request {i}: \" f\"prompt length {token_len} > max_model_len {max_len}\" ) else: valid_indices.append(i) valid_texts.append(text) if not valid_texts: return {\"responses\": [\"\"] * len(batch_messages), \"cur_batch_input_tokens\": 0, \"cur_batch_output_tokens\": 0} outputs = self.llm.generate( valid_texts, sampling_params=self.sampling_params ) valid_responses, total_in, total_out = self._process_outputs(outputs) responses_list = [\"\"] * len(batch_messages) for idx, resp in zip(valid_indices, valid_responses): responses_list[idx] = resp self.token_usage_tracker.update_usage( self.llm_config.model, total_in, total_out, task_name=task_name ) print( f\">>>>>> Token usage [{task_name}]: input={total_in}, output={total_out}, \" f\"total={total_in + total_out}\" ) return { \"responses\": responses_list, \"cur_batch_input_tokens\": total_in, \"cur_batch_output_tokens\": total_out, } def _process_outputs(self, outputs): responses_list = [] total_in = total_out = 0 for output in outputs: response = output.outputs[0].text responses_list.append(response) total_in += len(output.prompt_token_ids) total_out += len(output.outputs[0].token_ids) return responses_list, total_in, total_out def get_token_usage_summary(self): return self.token_usage_tracker.get_summary()","summary_state":"pending","summary":null,"crux":null},{"id":"src/models.py#LanguageModel.__init__","name":"__init__","kind":"method","path":"src/models.py","span":"L20-L48","signature":"def __init__(self, llm_config: LLMConfig)","exported":false,"origin":"ast","body_hash":"dcafac8ea0b0c1f80ca606c41fa198ae2cc977f932f9352e80a390f26e816339","body_text":"def __init__(self, llm_config: LLMConfig): import torch from vllm import LLM, SamplingParams self.llm_config = llm_config self.model_name = llm_config.name or llm_config.model.split(\"/\")[-1] llm_name_path = ( self.llm_config.model_path if self.llm_config.model_path not in (None, \"None\") else self.llm_config.model ) print(f\">>>>>> LLM name path: {llm_name_path}\") self.llm = LLM( model=llm_name_path, tensor_parallel_size=getattr(self.llm_config, \"tensor_parallel_size\", 1), dtype=torch.bfloat16, max_model_len=getattr(self.llm_config, \"max_token_length\", 24576), gpu_memory_utilization=0.85, ) self.sampling_params = SamplingParams( temperature=getattr(self.llm_config, \"temperature\", 1), top_p=getattr(self.llm_config, \"top_p\", 1), max_tokens=getattr(self.llm_config, \"max_tokens\", 6144), stop_token_ids=[self.llm.get_tokenizer().eos_token_id], ) self.token_usage_tracker = TokenUsageTracker()","summary_state":"pending","summary":null,"crux":null,"owner":"LanguageModel"},{"id":"src/models.py#LanguageModel.__call__","name":"__call__","kind":"method","path":"src/models.py","span":"L50-L104","signature":"def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]","exported":false,"origin":"ast","body_hash":"0daac3545eabb3d3056df818135154f7e24bd0faeafce4e1a7cac8719c3cd43a","body_text":"def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]: \"\"\" Args: batch_messages: list of message lists, each message list is [{\"role\": ..., \"content\": ...}, ...] task_name: label used to group token usage statistics Returns dict with keys: results, cur_batch_input_tokens, cur_batch_output_tokens \"\"\" tokenizer = self.llm.get_tokenizer() tokenized_texts = tokenizer.apply_chat_template( batch_messages, tokenize=False, add_generation_prompt=True ) # Handle over-long queries max_len = self.llm.llm_engine.model_config.max_model_len valid_indices, valid_texts = [], [] for i, text in enumerate(tokenized_texts): token_len = len(tokenizer.encode(text)) if token_len > max_len: print( f\"[WARNING] [{task_name}] Skipping request {i}: \" f\"prompt length {token_len} > max_model_len {max_len}\" ) else: valid_indices.append(i) valid_texts.append(text) if not valid_texts: return {\"responses\": [\"\"] * len(batch_messages), \"cur_batch_input_tokens\": 0, \"cur_batch_output_tokens\": 0} outputs = self.llm.generate( valid_texts, sampling_params=self.sampling_params ) valid_responses, total_in, total_out = self._process_outputs(outputs) responses_list = [\"\"] * len(batch_messages) for idx, resp in zip(valid_indices, valid_responses): responses_list[idx] = resp self.token_usage_tracker.update_usage( self.llm_config.model, total_in, total_out, task_name=task_name ) print( f\">>>>>> Token usage [{task_name}]: input={total_in}, output={total_out}, \" f\"total={total_in + total_out}\" ) return { \"responses\": responses_list, \"cur_batch_input_tokens\": total_in, \"cur_batch_output_tokens\": total_out, }","summary_state":"pending","summary":null,"crux":null,"owner":"LanguageModel"},{"id":"src/models.py#LanguageModel._process_outputs","name":"_process_outputs","kind":"method","path":"src/models.py","span":"L106-L117","signature":"def _process_outputs(self, outputs)","exported":false,"origin":"ast","body_hash":"df99c1544e817afecd91a63a37a1bbce53284fde8deab8ba51d5a652a0b18da9","body_text":"def _process_outputs(self, outputs): responses_list = [] total_in = total_out = 0 for output in outputs: response = output.outputs[0].text responses_list.append(response) total_in += len(output.prompt_token_ids) total_out += len(output.outputs[0].token_ids) return responses_list, total_in, total_out","summary_state":"pending","summary":null,"crux":null,"owner":"LanguageModel"},{"id":"src/models.py#LanguageModel.get_token_usage_summary","name":"get_token_usage_summary","kind":"method","path":"src/models.py","span":"L119-L120","signature":"def get_token_usage_summary(self)","exported":true,"origin":"ast","body_hash":"444b6293462006eaf247aadfde10851cf8370ee3ff93feac0944afd7fd8c7018","body_text":"def get_token_usage_summary(self): return self.token_usage_tracker.get_summary()","summary_state":"pending","summary":null,"crux":null,"owner":"LanguageModel"},{"id":"src/models.py#EmbeddingModel","name":"EmbeddingModel","kind":"class","path":"src/models.py","span":"L123-L162","signature":"class EmbeddingModel","exported":true,"origin":"ast","body_hash":"8b522fd76470bb294d0ae82506dae7a303eceef196a41440cad3ac09d50cacbe","body_text":"class EmbeddingModel: \"\"\"Wraps a SentenceTransformer model for text embedding.\"\"\" def __init__(self, embedding_config: EmbeddingConfig, device: str = None): import torch from sentence_transformers import SentenceTransformer self.embedding_config = embedding_config model_name_path = ( self.embedding_config.model_path if self.embedding_config.model_path not in (None, \"None\") else self.embedding_config.model ) # An explicit device (e.g. \"cuda:2\" from --embed_gpu_id) wins if device is None: device = \"cuda\" if torch.cuda.is_available() else \"cpu\" print(f\">>>>>> Loading embedding model '{model_name_path}' on {device}\") self.model = SentenceTransformer( model_name_path, device=device, trust_remote_code=True, ) def encode(self, texts: Union[str, List[str]], normalize_embeddings: bool = True, show_progress_bar: bool = False) -> Union[List[float], List[List[float]]]: try: is_single = isinstance(texts, str) inputs = [texts] if is_single else texts embeddings = self.model.encode( inputs, normalize_embeddings=normalize_embeddings, show_progress_bar=show_progress_bar ) return embeddings[0].tolist() if is_single else [e.tolist() for e in embeddings] except Exception as e: raise Exception(f\"Failed to encode texts: {str(e)}\")","summary_state":"pending","summary":null,"crux":null},{"id":"src/models.py#EmbeddingModel.__init__","name":"__init__","kind":"method","path":"src/models.py","span":"L125-L145","signature":"def __init__(self, embedding_config: EmbeddingConfig, device: str = None)","exported":false,"origin":"ast","body_hash":"a186eceffd7ae516233773014838d65a835908234ef63d77d118f6bba8046c7c","body_text":"def __init__(self, embedding_config: EmbeddingConfig, device: str = None): import torch from sentence_transformers import SentenceTransformer self.embedding_config = embedding_config model_name_path = ( self.embedding_config.model_path if self.embedding_config.model_path not in (None, \"None\") else self.embedding_config.model ) # An explicit device (e.g. \"cuda:2\" from --embed_gpu_id) wins if device is None: device = \"cuda\" if torch.cuda.is_available() else \"cpu\" print(f\">>>>>> Loading embedding model '{model_name_path}' on {device}\") self.model = SentenceTransformer( model_name_path, device=device, trust_remote_code=True, )","summary_state":"pending","summary":null,"crux":null,"owner":"EmbeddingModel"},{"id":"src/models.py#EmbeddingModel.encode","name":"encode","kind":"method","path":"src/models.py","span":"L147-L162","signature":"def encode(self, texts: Union[str, List[str]], normalize_embeddings: bool = True, show_progress_bar: bool = False) -> Union[List[float], List[List[float]]]","exported":true,"origin":"ast","body_hash":"f5f36ddbcdf93f01b1a961bd369cdc2bcb37aff06eac6441deba5f7395c41f8b","body_text":"def encode(self, texts: Union[str, List[str]], normalize_embeddings: bool = True, show_progress_bar: bool = False) -> Union[List[float], List[List[float]]]: try: is_single = isinstance(texts, str) inputs = [texts] if is_single else texts embeddings = self.model.encode( inputs, normalize_embeddings=normalize_embeddings, show_progress_bar=show_progress_bar ) return embeddings[0].tolist() if is_single else [e.tolist() for e in embeddings] except Exception as e: raise Exception(f\"Failed to encode texts: {str(e)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"EmbeddingModel"},{"id":"src/models.py#OpenRouterLanguageModel","name":"OpenRouterLanguageModel","kind":"class","path":"src/models.py","span":"L165-L338","signature":"class OpenRouterLanguageModel","exported":true,"origin":"ast","body_hash":"ec16cae94add177e883ec5e07cdc34ed9600ff1a8158145154208318c441ca06","body_text":"class OpenRouterLanguageModel: \"\"\" OpenRouter (OpenAI-compatible) backend that mirrors the interface of ``LanguageModel``: model(batch_messages, task_name) -> {\"responses\", \"cur_batch_input_tokens\", \"cur_batch_output_tokens\"} model.model_name model.get_token_usage_summary() A batch is fanned out over a thread pool (each request is a blocking, IO-bound HTTP call). Concurrency is capped by ``max_workers`` and every request retries on rate-limit / transient errors with exponential backoff + jitter, honoring the server's ``Retry-After`` when present. A request that still fails yields \"\" so a single bad call never aborts the whole debate round (same fallback as the vLLM backend for over-long / empty prompts). \"\"\" def __init__(self, llm_config: LLMConfig): from openai import OpenAI, APIError, APITimeoutError, RateLimitError, APIConnectionError self.llm_config = llm_config self.model_name = llm_config.name or llm_config.model.split(\"/\")[-1] self.model_id = llm_config.model # Exception types stored on the instance so the lazy import stays local. self._api_error = APIError self._retryable = (RateLimitError, APITimeoutError, APIConnectionError) api_key = os.environ.get(llm_config.api_key_env) if not api_key: raise RuntimeError( f\"OpenRouter backend requires the '{llm_config.api_key_env}' environment \" f\"variable to be set for model '{self.model_name}'.\" ) # max_retries=0: we own retry/backoff so the SDK's built-in retries don't # multiply the request count on top of ours. self.client = OpenAI( base_url=llm_config.api_base, api_key=api_key, timeout=llm_config.request_timeout, max_retries=0, ) self.max_workers = max(1, int(llm_config.max_workers)) self.max_retries = max(0, int(llm_config.max_retries)) self.temperature = llm_config.temperature self.top_p = llm_config.top_p self.max_tokens = llm_config.max_tokens self.token_usage_tracker = TokenUsageTracker() print( f\">>>>>> OpenRouter model: {self.model_id} \" f\"(max_workers={self.max_workers}, max_retries={self.max_retries})\" ) def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]: if not batch_messages: return {\"responses\": [], \"cur_batch_input_tokens\": 0, \"cur_batch_output_tokens\": 0} from concurrent.futures import as_completed from tqdm import tqdm # Submit all requests, then consume them as they finish so tqdm reflects real progress. # Results are placed back by original index to keep output order. results: List[Optional[Tuple[str, int, int]]] = [None] * len(batch_messages) with ThreadPoolExecutor(max_workers=self.max_workers) as executor: future_to_idx = { executor.submit(self._complete_one_query, messages, task_name): i for i, messages in enumerate(batch_messages) } for future in tqdm( as_completed(future_to_idx), total=len(batch_messages), desc=f\"[{task_name}] {self.model_name}\", leave=False, ): idx = future_to_idx[future] try: results[idx] = future.result() except Exception as e: print(f\"[WARNING] [{task_name}] request worker failed at index {idx}: {e}\") results[idx] = (\"\", 0, 0) responses_list: List[str] = [] total_in = total_out = 0 for text, in_tok, out_tok in results: responses_list.append(text) total_in += in_tok total_out += out_tok self.token_usage_tracker.update_usage( self.model_id, total_in, total_out, task_name=task_name ) print( f\">>>>>> Token usage [{task_name}]: input={total_in}, output={total_out}, \" f\"total={total_in + total_out}\" ) return { \"responses\": responses_list, \"cur_batch_input_tokens\": total_in, \"cur_batch_output_tokens\": total_out, } def _complete_one_query(self, messages: List[Dict[str, str]], task_name: str) -> Tuple[str, int, int]: \"\"\"Single chat completion with retry/backoff. Returns (text, in_tokens, out_tokens).\"\"\" for attempt in range(self.max_retries + 1): try: resp = self.client.chat.completions.create( model=self.model_id, messages=messages, temperature=self.temperature, top_p=self.top_p, max_tokens=self.max_tokens, ) text = resp.choices[0].message.content or \"\" usage = getattr(resp, \"usage\", None) in_tok = getattr(usage, \"prompt_tokens\", 0) or 0 out_tok = getattr(usage, \"completion_tokens\", 0) or 0 return text, in_tok, out_tok except self._retryable as e: if attempt >= self.max_retries: print(f\"[WARNING] [{task_name}] request failed after {attempt + 1} attempts: {e}\") return \"\", 0, 0 time.sleep(self._backoff_seconds(attempt, e)) except json.JSONDecodeError as e: # Gateway occasionally returns HTML/empty body under heavy load. if attempt >= self.max_retries: print( f\"[WARNING] [{task_name}] malformed JSON response after \" f\"{attempt + 1} attempts: {e}\" ) return \"\", 0, 0 time.sleep(self._backoff_seconds(attempt, e)) except self._api_error as e: # Non-retryable API error (bad request, auth, etc.): give up on this one. print(f\"[WARNING] [{task_name}] non-retryable API error: {e}\") return \"\", 0, 0 except Exception as e: if attempt >= self.max_retries: print( f\"[WARNING] [{task_name}] unexpected error after \" f\"{attempt + 1} attempts: {type(e).__name__}: {e}\" ) return \"\", 0, 0","summary_state":"pending","summary":null,"crux":null},{"id":"src/models.py#OpenRouterLanguageModel.__init__","name":"__init__","kind":"method","path":"src/models.py","span":"L182-L220","signature":"def __init__(self, llm_config: LLMConfig)","exported":false,"origin":"ast","body_hash":"479d514551f36fc59c7336df936fc08e3619cd0991763e6aebcd31ae1309a9ce","body_text":"def __init__(self, llm_config: LLMConfig): from openai import OpenAI, APIError, APITimeoutError, RateLimitError, APIConnectionError self.llm_config = llm_config self.model_name = llm_config.name or llm_config.model.split(\"/\")[-1] self.model_id = llm_config.model # Exception types stored on the instance so the lazy import stays local. self._api_error = APIError self._retryable = (RateLimitError, APITimeoutError, APIConnectionError) api_key = os.environ.get(llm_config.api_key_env) if not api_key: raise RuntimeError( f\"OpenRouter backend requires the '{llm_config.api_key_env}' environment \" f\"variable to be set for model '{self.model_name}'.\" ) # max_retries=0: we own retry/backoff so the SDK's built-in retries don't # multiply the request count on top of ours. self.client = OpenAI( base_url=llm_config.api_base, api_key=api_key, timeout=llm_config.request_timeout, max_retries=0, ) self.max_workers = max(1, int(llm_config.max_workers)) self.max_retries = max(0, int(llm_config.max_retries)) self.temperature = llm_config.temperature self.top_p = llm_config.top_p self.max_tokens = llm_config.max_tokens self.token_usage_tracker = TokenUsageTracker() print( f\">>>>>> OpenRouter model: {self.model_id} \" f\"(max_workers={self.max_workers}, max_retries={self.max_retries})\" )","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"},{"id":"src/models.py#OpenRouterLanguageModel.__call__","name":"__call__","kind":"method","path":"src/models.py","span":"L222-L269","signature":"def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]","exported":false,"origin":"ast","body_hash":"e14c4f0346e22ebf9af00cf83461f3782b0533e5d0be027d1075b3531e20813d","body_text":"def __call__(self, batch_messages: List[List[Dict[str, str]]], task_name: str = \"default\") -> Dict[str, Any]: if not batch_messages: return {\"responses\": [], \"cur_batch_input_tokens\": 0, \"cur_batch_output_tokens\": 0} from concurrent.futures import as_completed from tqdm import tqdm # Submit all requests, then consume them as they finish so tqdm reflects real progress. # Results are placed back by original index to keep output order. results: List[Optional[Tuple[str, int, int]]] = [None] * len(batch_messages) with ThreadPoolExecutor(max_workers=self.max_workers) as executor: future_to_idx = { executor.submit(self._complete_one_query, messages, task_name): i for i, messages in enumerate(batch_messages) } for future in tqdm( as_completed(future_to_idx), total=len(batch_messages), desc=f\"[{task_name}] {self.model_name}\", leave=False, ): idx = future_to_idx[future] try: results[idx] = future.result() except Exception as e: print(f\"[WARNING] [{task_name}] request worker failed at index {idx}: {e}\") results[idx] = (\"\", 0, 0) responses_list: List[str] = [] total_in = total_out = 0 for text, in_tok, out_tok in results: responses_list.append(text) total_in += in_tok total_out += out_tok self.token_usage_tracker.update_usage( self.model_id, total_in, total_out, task_name=task_name ) print( f\">>>>>> Token usage [{task_name}]: input={total_in}, output={total_out}, \" f\"total={total_in + total_out}\" ) return { \"responses\": responses_list, \"cur_batch_input_tokens\": total_in, \"cur_batch_output_tokens\": total_out, }","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"},{"id":"src/models.py#OpenRouterLanguageModel._complete_one_query","name":"_complete_one_query","kind":"method","path":"src/models.py","span":"L271-L314","signature":"def _complete_one_query(self, messages: List[Dict[str, str]], task_name: str) -> Tuple[str, int, int]","exported":false,"origin":"ast","body_hash":"b23c05e4e9c31ef52a6b215944998d3d3b9b3a8edc6ebbd851f88adadb29335a","body_text":"def _complete_one_query(self, messages: List[Dict[str, str]], task_name: str) -> Tuple[str, int, int]: \"\"\"Single chat completion with retry/backoff. Returns (text, in_tokens, out_tokens).\"\"\" for attempt in range(self.max_retries + 1): try: resp = self.client.chat.completions.create( model=self.model_id, messages=messages, temperature=self.temperature, top_p=self.top_p, max_tokens=self.max_tokens, ) text = resp.choices[0].message.content or \"\" usage = getattr(resp, \"usage\", None) in_tok = getattr(usage, \"prompt_tokens\", 0) or 0 out_tok = getattr(usage, \"completion_tokens\", 0) or 0 return text, in_tok, out_tok except self._retryable as e: if attempt >= self.max_retries: print(f\"[WARNING] [{task_name}] request failed after {attempt + 1} attempts: {e}\") return \"\", 0, 0 time.sleep(self._backoff_seconds(attempt, e)) except json.JSONDecodeError as e: # Gateway occasionally returns HTML/empty body under heavy load. if attempt >= self.max_retries: print( f\"[WARNING] [{task_name}] malformed JSON response after \" f\"{attempt + 1} attempts: {e}\" ) return \"\", 0, 0 time.sleep(self._backoff_seconds(attempt, e)) except self._api_error as e: # Non-retryable API error (bad request, auth, etc.): give up on this one. print(f\"[WARNING] [{task_name}] non-retryable API error: {e}\") return \"\", 0, 0 except Exception as e: if attempt >= self.max_retries: print( f\"[WARNING] [{task_name}] unexpected error after \" f\"{attempt + 1} attempts: {type(e).__name__}: {e}\" ) return \"\", 0, 0 time.sleep(self._backoff_seconds(attempt, e)) return \"\", 0, 0","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"},{"id":"src/models.py#OpenRouterLanguageModel._backoff_seconds","name":"_backoff_seconds","kind":"method","path":"src/models.py","span":"L316-L321","signature":"def _backoff_seconds(self, attempt: int, error: Exception) -> float","exported":false,"origin":"ast","body_hash":"8505645ed01ef0adc8191a001b113326e667e59123db5260ee70933cee0775d2","body_text":"def _backoff_seconds(self, attempt: int, error: Exception) -> float: \"\"\"Honor Retry-After when the server sends it; otherwise exponential backoff + jitter.\"\"\" retry_after = self._parse_retry_after(error) if retry_after is not None: return retry_after + random.random() return min(2.0 ** attempt, 30.0) + random.random()","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"},{"id":"src/models.py#OpenRouterLanguageModel._parse_retry_after","name":"_parse_retry_after","kind":"method","path":"src/models.py","span":"L324-L335","signature":"def _parse_retry_after(error: Exception) -> Optional[float]","exported":false,"origin":"ast","body_hash":"ab20914a2001ed4f3e27db2a2c5c72b9404361407b409ea5b2d049c932de68f8","body_text":"def _parse_retry_after(error: Exception) -> Optional[float]: response = getattr(error, \"response\", None) headers = getattr(response, \"headers\", None) if not headers: return None value = headers.get(\"Retry-After\") or headers.get(\"retry-after\") if value is None: return None try: return float(value) except (TypeError, ValueError): return None","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"},{"id":"src/models.py#OpenRouterLanguageModel.get_token_usage_summary","name":"get_token_usage_summary","kind":"method","path":"src/models.py","span":"L337-L338","signature":"def get_token_usage_summary(self)","exported":true,"origin":"ast","body_hash":"444b6293462006eaf247aadfde10851cf8370ee3ff93feac0944afd7fd8c7018","body_text":"def get_token_usage_summary(self): return self.token_usage_tracker.get_summary()","summary_state":"pending","summary":null,"crux":null,"owner":"OpenRouterLanguageModel"}],"rawEdges":[{"source":"src/models.py","relation":"imports","specifier":"os","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"time","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"random","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"json","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"concurrent.futures","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"typing","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"src.utils","file":"src/models.py"},{"source":"src/models.py","relation":"contains","targetId":"src/models.py#build_language_model","file":"src/models.py"},{"source":"src/models.py#build_language_model","relation":"calls","name":"OpenRouterLanguageModel","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#build_language_model","relation":"calls","name":"LanguageModel","viaMember":false,"file":"src/models.py"},{"source":"src/models.py","relation":"contains","targetId":"src/models.py#LanguageModel","file":"src/models.py"},{"source":"src/models.py#LanguageModel","relation":"contains","targetId":"src/models.py#LanguageModel.__init__","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"torch","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"vllm","file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"split","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"LLM","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"SamplingParams","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"get_tokenizer","viaMember":true,"file":"src/models.py","recvType":"LLM"},{"source":"src/models.py#LanguageModel.__init__","relation":"calls","name":"TokenUsageTracker","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel","relation":"contains","targetId":"src/models.py#LanguageModel.__call__","file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"get_tokenizer","viaMember":true,"file":"src/models.py","recvType":"LLM"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"apply_chat_template","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"enumerate","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"encode","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"append","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"append","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"generate","viaMember":true,"file":"src/models.py","recvType":"LLM"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"_process_outputs","viaMember":true,"file":"src/models.py","recvType":"LanguageModel"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"zip","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"update_usage","viaMember":true,"file":"src/models.py","recvType":"TokenUsageTracker"},{"source":"src/models.py#LanguageModel.__call__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel","relation":"contains","targetId":"src/models.py#LanguageModel._process_outputs","file":"src/models.py"},{"source":"src/models.py#LanguageModel._process_outputs","relation":"calls","name":"append","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#LanguageModel._process_outputs","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel._process_outputs","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#LanguageModel","relation":"contains","targetId":"src/models.py#LanguageModel.get_token_usage_summary","file":"src/models.py"},{"source":"src/models.py#LanguageModel.get_token_usage_summary","relation":"calls","name":"get_summary","viaMember":true,"file":"src/models.py","recvType":"TokenUsageTracker"},{"source":"src/models.py","relation":"contains","targetId":"src/models.py#EmbeddingModel","file":"src/models.py"},{"source":"src/models.py#EmbeddingModel","relation":"contains","targetId":"src/models.py#EmbeddingModel.__init__","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"torch","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"sentence_transformers","file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.__init__","relation":"calls","name":"is_available","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.__init__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.__init__","relation":"calls","name":"SentenceTransformer","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel","relation":"contains","targetId":"src/models.py#EmbeddingModel.encode","file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"isinstance","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"encode","viaMember":true,"file":"src/models.py","recvType":"SentenceTransformer"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"tolist","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"tolist","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"Exception","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#EmbeddingModel.encode","relation":"calls","name":"str","viaMember":false,"file":"src/models.py"},{"source":"src/models.py","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel.__init__","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"openai","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"split","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"RuntimeError","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"OpenAI","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"max","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"int","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"max","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"int","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"TokenUsageTracker","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__init__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel.__call__","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"concurrent.futures","file":"src/models.py"},{"source":"src/models.py","relation":"imports","specifier":"tqdm","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"ThreadPoolExecutor","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"submit","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"enumerate","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"tqdm","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"as_completed","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"len","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"result","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"append","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"update_usage","viaMember":true,"file":"src/models.py","recvType":"TokenUsageTracker"},{"source":"src/models.py#OpenRouterLanguageModel.__call__","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel._complete_one_query","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"range","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"create","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"sleep","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"_backoff_seconds","viaMember":true,"file":"src/models.py","recvType":"OpenRouterLanguageModel"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"sleep","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"_backoff_seconds","viaMember":true,"file":"src/models.py","recvType":"OpenRouterLanguageModel"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"print","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"type","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"sleep","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._complete_one_query","relation":"calls","name":"_backoff_seconds","viaMember":true,"file":"src/models.py","recvType":"OpenRouterLanguageModel"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel._backoff_seconds","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._backoff_seconds","relation":"calls","name":"_parse_retry_after","viaMember":true,"file":"src/models.py","recvType":"OpenRouterLanguageModel"},{"source":"src/models.py#OpenRouterLanguageModel._backoff_seconds","relation":"calls","name":"random","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._backoff_seconds","relation":"calls","name":"min","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._backoff_seconds","relation":"calls","name":"random","viaMember":true,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel._parse_retry_after","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._parse_retry_after","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._parse_retry_after","relation":"calls","name":"getattr","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel._parse_retry_after","relation":"calls","name":"get","viaMember":true,"file":"src/models.py","recvType":"getattr"},{"source":"src/models.py#OpenRouterLanguageModel._parse_retry_after","relation":"calls","name":"get","viaMember":true,"file":"src/models.py","recvType":"getattr"},{"source":"src/models.py#OpenRouterLanguageModel._parse_retry_after","relation":"calls","name":"float","viaMember":false,"file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel","relation":"contains","targetId":"src/models.py#OpenRouterLanguageModel.get_token_usage_summary","file":"src/models.py"},{"source":"src/models.py#OpenRouterLanguageModel.get_token_usage_summary","relation":"calls","name":"get_summary","viaMember":true,"file":"src/models.py","recvType":"TokenUsageTracker"}]},"src/multi_agent_debate.py":{"size":19919,"mtimeMs":1789362123401.8032,"hash":"75cf85f0c06c03eda250c47d3195fecaf563f72825d29ff911c5a002c135ebca","nodes":[{"id":"src/multi_agent_debate.py","name":"multi_agent_debate.py","kind":"file","path":"src/multi_agent_debate.py","span":"L1-L461","signature":null,"exported":true,"origin":"ast","body_hash":"75cf85f0c06c03eda250c47d3195fecaf563f72825d29ff911c5a002c135ebca","chars":19919,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import re import json from datetime import datetime from typing import List, Dict, Any, Optional from collections import defaultdict, Counter import numpy as np from src.agent import Query, Agent from src.models import LanguageModel from src.checker import ANSWER_CHECKER from src.paths import MEMORY_DATA_DIR from src.prompts import DEBATE_SUMMARY_PROMPT"},{"id":"src/multi_agent_debate.py#remove_think_tags","name":"remove_think_tags","kind":"function","path":"src/multi_agent_debate.py","span":"L17-L18","signature":"def remove_think_tags(text: str) -> str","exported":true,"origin":"ast","body_hash":"100e88f7a94003ca9e6d0a377e4c5dccaa6c51a3b54fe0ab0682b7a93960f1ca","body_text":"def remove_think_tags(text: str) -> str: return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)","summary_state":"pending","summary":null,"crux":null},{"id":"src/multi_agent_debate.py#MultiAgentDebate","name":"MultiAgentDebate","kind":"class","path":"src/multi_agent_debate.py","span":"L21-L460","signature":"class MultiAgentDebate","exported":true,"origin":"ast","body_hash":"acb8c5309cb801dbfa0eb000fd078814b8d03139376c392661f0c13ab70ee02c","body_text":"class MultiAgentDebate: \"\"\" Orchestrates a multi-round debate for a batch of queries. A fixed list of Agent instances is shared across all queries; each agent stores conversation history per query_id internally, so no duplication of agent objects is needed. Round 0 : every agent answers independently. Round 1+: every agent sees the other agents' previous responses and revises. All queries are processed in a single batched model call per round; the batch sent to LanguageModel has shape (num_active_queries * num_agents). \"\"\" def __init__( self, agents: List[Agent], model: LanguageModel, task_type: str, max_round: int = 3, consensus_threshold: float = 0.8, if_summarize_state: bool = False, if_confidence_weight: bool = False, ) -> None: self.agents = agents self.model = model self.max_round = max_round self.consensus_threshold = consensus_threshold self.if_summarize_state = if_summarize_state self.if_confidence_weight = if_confidence_weight self.task_type = task_type self._checker = ANSWER_CHECKER[task_type] # Persistent debate state (reset each `run`) self.current_rounds = {} self.agent_responses = defaultdict(lambda: defaultdict(dict)) self.consensus_flags = {} self.debate_log = {} # prev_resp_embeddings[qid][agent_id_str] = unit-norm float32 ndarray self._prev_resp_embeddings = defaultdict(lambda: dict()) # ------------------------------------------------------------------ # Public entry point # ------------------------------------------------------------------ def run(self, queries: List[Query], is_training: bool = True, batch_idx: int = 0) -> List[Dict[str, Any]]: \"\"\" Run a full debate for a batch of queries. Args: queries: list of Query objects Returns: list of debate_log entries (one per query), each with keys: query, rounds, final_answer, consensus \"\"\" self._clear_debate() self._initialize_debate(queries) for round_idx in range(self.max_round): if is_training: # force all queries to run all rounds during training active_queries = queries else: active_queries = [q for q in queries if not self.consensus_flags[q.query_id]] if not active_queries: break print(f\"\\n>>>>>> Round {round_idx}\\n\") self._run_debate_round(active_queries) for query in queries: qid = query.query_id last_round = self.current_rounds[qid] - 1 correct_count, error_count = 0, 0 final_true_answer, final_false_answer = None, None for agent in self.agents: if self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"is_correct\"]: correct_count += 1 final_true_answer = self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"pred_answer\"] else: error_count += 1 final_false_answer = self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"pred_answer\"] is_final_correct = (correct_count > error_count) final_answer = final_true_answer if is_final_correct else final_false_answer self.debate_log[qid][\"final_answer\"] = final_answer self.debate_log[qid][\"is_final_correct\"] = is_final_correct self.debate_log[qid][\"consensus\"] = self.consensus_flags[qid] if is_training: self._save_memory_data(queries, batch_idx=batch_idx) return [self.debate_log[q.query_id] for q in queries] # ------------------------------------------------------------------ # Round execution # ------------------------------------------------------------------ def _run_debate_round(self, queries: List[Query]) -> None: \"\"\" Run one debate round for the given active queries. For round 0, each agent receives no prior responses (initial answer). For round > 0, each agent receives the other agents' responses from the previous round, read from self.agent_responses. \"\"\" # Build batched messages: [q0_a0, q0_a1, ..., q1_a0, q1_a1, ...] batch_messages = [] flat_meta = [] # (query, agent_idx) # Precompute per-query context maps for batch memory retrieval prev_responses_map: Dict[int, Any] = {} prev_round_summaries: Dict[int, Any] = {} prev_consensus_ratios: Dict[int, Any] = {} for query in queries: qid = query.query_id round_idx = self.current_rounds[qid] rounds_log = self.debate_log[qid][\"rounds\"] if round_idx == 0: prev_responses_map[qid] = None prev_round_summaries[qid] = None prev_consensus_ratios[qid] = None else: emb_cache = self._prev_resp_embeddings.get(qid, {}) if self.if_confidence_weight else {} prev_responses_map[qid] = { agent.agent_id: ( self.agent_responses[qid][agent.agent_id][f\"round_{round_idx - 1}\"].get(\"response\", \"\"), emb_cache.get(agent.agent_id, None), ) for agent in self.agents } prev_round_summaries[qid] = rounds_log[-1].get(\"summary\", None) prev_consensus_ratios[qid] = rounds_log[-1].get(\"consensus_ratio\", None) # Batch memory retrieval: one DB round-trip per agent instead of per (agent, query) agent_prefetched = {} for agent in self.agents: agent_prefetched[agent.agent_id] = agent.batch_retrieve_all( queries=queries, prev_responses_map={ qid: resp for qid, resp in prev_responses_map.items() if resp is not None }, prev_round_summaries=prev_round_summaries, prev_consensus_ratios=prev_consensus_ratios, ) for query in queries: qid = query.query_id ","summary_state":"pending","summary":null,"crux":null},{"id":"src/multi_agent_debate.py#MultiAgentDebate.__init__","name":"__init__","kind":"method","path":"src/multi_agent_debate.py","span":"L36-L62","signature":"def __init__( self, agents: List[Agent], model: LanguageModel, task_type: str, max_round: int = 3, consensus_threshold: float = 0.8, if_summarize_state: bool = False, if_confidence_weight: bool = False, ) -> None","exported":false,"origin":"ast","body_hash":"170a598d9d8154d623eae2efef3465d6e19e02ffe09a60cbf7e0cade26176e82","body_text":"def __init__( self, agents: List[Agent], model: LanguageModel, task_type: str, max_round: int = 3, consensus_threshold: float = 0.8, if_summarize_state: bool = False, if_confidence_weight: bool = False, ) -> None: self.agents = agents self.model = model self.max_round = max_round self.consensus_threshold = consensus_threshold self.if_summarize_state = if_summarize_state self.if_confidence_weight = if_confidence_weight self.task_type = task_type self._checker = ANSWER_CHECKER[task_type] # Persistent debate state (reset each `run`) self.current_rounds = {} self.agent_responses = defaultdict(lambda: defaultdict(dict)) self.consensus_flags = {} self.debate_log = {} # prev_resp_embeddings[qid][agent_id_str] = unit-norm float32 ndarray self._prev_resp_embeddings = defaultdict(lambda: dict())","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate.run","name":"run","kind":"method","path":"src/multi_agent_debate.py","span":"L68-L117","signature":"def run(self, queries: List[Query], is_training: bool = True, batch_idx: int = 0) -> List[Dict[str, Any]]","exported":true,"origin":"ast","body_hash":"ed77736ce74891bbdba922ab8726fbdf39855da68029487865eb2e17c1d93db6","body_text":"def run(self, queries: List[Query], is_training: bool = True, batch_idx: int = 0) -> List[Dict[str, Any]]: \"\"\" Run a full debate for a batch of queries. Args: queries: list of Query objects Returns: list of debate_log entries (one per query), each with keys: query, rounds, final_answer, consensus \"\"\" self._clear_debate() self._initialize_debate(queries) for round_idx in range(self.max_round): if is_training: # force all queries to run all rounds during training active_queries = queries else: active_queries = [q for q in queries if not self.consensus_flags[q.query_id]] if not active_queries: break print(f\"\\n>>>>>> Round {round_idx}\\n\") self._run_debate_round(active_queries) for query in queries: qid = query.query_id last_round = self.current_rounds[qid] - 1 correct_count, error_count = 0, 0 final_true_answer, final_false_answer = None, None for agent in self.agents: if self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"is_correct\"]: correct_count += 1 final_true_answer = self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"pred_answer\"] else: error_count += 1 final_false_answer = self.agent_responses[qid][agent.agent_id][f\"round_{last_round}\"][\"pred_answer\"] is_final_correct = (correct_count > error_count) final_answer = final_true_answer if is_final_correct else final_false_answer self.debate_log[qid][\"final_answer\"] = final_answer self.debate_log[qid][\"is_final_correct\"] = is_final_correct self.debate_log[qid][\"consensus\"] = self.consensus_flags[qid] if is_training: self._save_memory_data(queries, batch_idx=batch_idx) return [self.debate_log[q.query_id] for q in queries]","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","name":"_run_debate_round","kind":"method","path":"src/multi_agent_debate.py","span":"L123-L273","signature":"def _run_debate_round(self, queries: List[Query]) -> None","exported":false,"origin":"ast","body_hash":"2ff2af579b650f185ef35ca9003ddcb08ca9414103e9a77419bf8641106243d9","body_text":"def _run_debate_round(self, queries: List[Query]) -> None: \"\"\" Run one debate round for the given active queries. For round 0, each agent receives no prior responses (initial answer). For round > 0, each agent receives the other agents' responses from the previous round, read from self.agent_responses. \"\"\" # Build batched messages: [q0_a0, q0_a1, ..., q1_a0, q1_a1, ...] batch_messages = [] flat_meta = [] # (query, agent_idx) # Precompute per-query context maps for batch memory retrieval prev_responses_map: Dict[int, Any] = {} prev_round_summaries: Dict[int, Any] = {} prev_consensus_ratios: Dict[int, Any] = {} for query in queries: qid = query.query_id round_idx = self.current_rounds[qid] rounds_log = self.debate_log[qid][\"rounds\"] if round_idx == 0: prev_responses_map[qid] = None prev_round_summaries[qid] = None prev_consensus_ratios[qid] = None else: emb_cache = self._prev_resp_embeddings.get(qid, {}) if self.if_confidence_weight else {} prev_responses_map[qid] = { agent.agent_id: ( self.agent_responses[qid][agent.agent_id][f\"round_{round_idx - 1}\"].get(\"response\", \"\"), emb_cache.get(agent.agent_id, None), ) for agent in self.agents } prev_round_summaries[qid] = rounds_log[-1].get(\"summary\", None) prev_consensus_ratios[qid] = rounds_log[-1].get(\"consensus_ratio\", None) # Batch memory retrieval: one DB round-trip per agent instead of per (agent, query) agent_prefetched = {} for agent in self.agents: agent_prefetched[agent.agent_id] = agent.batch_retrieve_all( queries=queries, prev_responses_map={ qid: resp for qid, resp in prev_responses_map.items() if resp is not None }, prev_round_summaries=prev_round_summaries, prev_consensus_ratios=prev_consensus_ratios, ) for query in queries: qid = query.query_id prev_responses = prev_responses_map[qid] prev_round_summary = prev_round_summaries[qid] prev_consensus_ratio = prev_consensus_ratios[qid] for a_idx, agent in enumerate(self.agents): prefetched = agent_prefetched[agent.agent_id].get(qid) messages = agent.build_messages( query=query, prev_responses=prev_responses, prev_round_summary=prev_round_summary, prev_consensus_ratio=prev_consensus_ratio, prefetched=prefetched, ) batch_messages.append(messages) flat_meta.append((query, a_idx)) model_output = self.model(batch_messages, task_name=\"debate\") responses = [remove_think_tags(r).strip() for r in model_output[\"responses\"]] # Batch-encode all responses once for confidence weighting in the next round resp_embs = None if self.if_confidence_weight: embed_fn = next( (a.memory_db.embedding_model.encode for a in self.agents if a.memory_db is not None), None, ) if embed_fn is not None: resp_embs = np.array(embed_fn(responses), dtype=np.float32) # Process and store per-agent results into agent_responses for i, (query, a_idx) in enumerate(flat_meta): qid = query.query_id round_idx = self.current_rounds[qid] pred_answer = self._checker.parse_answer(responses[i]) is_correct = self._checker.check_answer(pred_answer, query.answer) agent = self.agents[a_idx] agent.record_response(query, responses[i]) self.agent_responses[qid][agent.agent_id][f\"round_{round_idx}\"] = { \"response\": responses[i], \"pred_answer\": pred_answer, \"is_correct\": is_correct, } if resp_embs is not None: self._prev_resp_embeddings[qid][agent.agent_id] = resp_embs[i] # Check consensus and log each active query for query in queries: qid = query.query_id round_idx = self.current_rounds[qid] pred_answers = [ self.agent_responses[qid][agent.agent_id][f\"round_{round_idx}\"][\"pred_answer\"] for agent in self.agents ] is_consensus, consensus_ratio, common_answer = self._check_consensus(query, pred_answers) agents_log = [] for agent in self.agents: round_data = self.agent_responses[qid][agent.agent_id][f\"round_{round_idx}\"] entry = { \"agent_id\": str(agent.agent_id), \"response\": round_data[\"response\"], \"pred_answer\": round_data[\"pred_answer\"], \"is_correct\": round_data[\"is_correct\"], } # prefetched is (retrieved_memories, other_confidences), or None at 0 prefetched = agent_prefetched.get(agent.agent_id, {}).get(qid) retrieved, confidences = prefetched if prefetched else (None, None) # Memories this agent retrieved to produce this round's response. if agent.if_use_memory: entry[\"retrieved_memories\"] = [{ \"content\": rm.content, \"is_correct\": rm.is_correct, \"similarity\": rm.similarity, } for rm in (retrieved or [])] # Confidence this agent assigned to each peer's previous response. if self.if_confidence_weight: entry[\"peer_confidences\"] = { str(peer_id): float(conf) for peer_id, conf in (confidences or {}).items() } agents_log.append(entry) self.debate_log[qid][\"rounds\"].append({ \"round\": round_idx, \"agents\": agents_log, \"common_answer\": common_answer, \"consensus\": is_consensus, \"consensus_ratio\": consensus_ratio, }) if is_consensus: self.consensus_flags[qid] = True self.current_rounds[qid] += 1 if self.if_summarize_state: print(f\"\\n>>>>>> Summarize Debate State\\n\") self._summarize_debate_state(queries)","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","name":"_summarize_debate_state","kind":"method","path":"src/multi_agent_debate.py","span":"L275-L326","signature":"def _summarize_debate_state(self, queries: List[Query]) -> None","exported":false,"origin":"ast","body_hash":"8057bd6b6106b81c97e4bda1cab5be442db18881b4fac4f52624b74ceeb3c380","body_text":"def _summarize_debate_state(self, queries: List[Query]) -> None: \"\"\" Generate debate summaries for the just-completed round across all queries in a single batched model call, then write results back to debate_log. \"\"\" # Build one summary prompt per query, tracking which round each refers to batch_messages = [] meta = [] # (qid, round_log_index) for query in queries: rounds = self.debate_log[query.query_id][\"rounds\"] round_log = rounds[-1] round_idx = len(rounds) - 1 prev_summary = rounds[-2][\"summary\"] if round_idx > 0 else \"None\" agent_responses = [] for i, entry in enumerate(round_log[\"agents\"]): if round_idx > 0: prev_entry = rounds[-2][\"agents\"][i] stance_changed = str(entry[\"pred_answer\"] != prev_entry[\"pred_answer\"]).lower() agent_responses.append( f'<agent index=\"{i}\">\\n' f'<response>\\n{entry[\"response\"]}\\n</response>\\n' f'<answer>\\n{entry[\"pred_answer\"]}\\n</answer>\\n' f'<previous_answer>\\n{prev_entry[\"pred_answer\"]}\\n</previous_answer>\\n' f'<stance_changed>\\n{stance_changed}\\n</stance_changed>\\n' f'</agent>' ) else: agent_responses.append( f'<agent index=\"{i}\">\\n' f'<response>\\n{entry[\"response\"]}\\n</response>\\n' f'<answer>\\n{entry[\"pred_answer\"]}\\n</answer>\\n' f'</agent>' ) agent_responses_text = \"\\n\".join(agent_responses) prompt = DEBATE_SUMMARY_PROMPT.format( QUESTION=query.content, PREV_SUMMARY=prev_summary, AGENT_RESPONSES=agent_responses_text, CONSENSUS_RATIO=f\"{round_log['consensus_ratio']:.2f}\", ) batch_messages.append([{\"role\": \"user\", \"content\": prompt}]) meta.append((query.query_id, round_idx)) if not batch_messages: return model_output = self.model(batch_messages, task_name=\"summary\") for (qid, round_idx), summary in zip(meta, model_output[\"responses\"]): self.debate_log[qid][\"rounds\"][round_idx][\"summary\"] = remove_think_tags(summary).strip()","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","name":"_check_consensus","kind":"method","path":"src/multi_agent_debate.py","span":"L328-L357","signature":"def _check_consensus(self, query: Query, pred_answers: List[str]) -> tuple[bool, float, str]","exported":false,"origin":"ast","body_hash":"6dd427e7e8d015560f5e251b342fc624c7d2080c76a13d13050e4aae53877670","body_text":"def _check_consensus(self, query: Query, pred_answers: List[str]) -> tuple[bool, float, str]: \"\"\" Check if all agents have reached consensus on the same answer. Returns: is_consensus: bool consensus_ratio: float (fraction of agents that agreed on the most common answer) common_answer: str (the answer that the majority agreed on, or None if no consensus) \"\"\" if not pred_answers: return False, 0.0, \"\" qid = query.query_id round_idx = self.current_rounds[qid] normalized_answers = [answer.lower().strip() for answer in pred_answers] correct_count = 0 for agent in self.agents: if self.agent_responses[qid][agent.agent_id][f\"round_{round_idx}\"][\"is_correct\"]: correct_count += 1 if correct_count == len(pred_answers): return True, 1.0, normalized_answers[0] answer_counts = Counter(normalized_answers) common_answer, max_count = answer_counts.most_common(1)[0] consensus_ratio = max_count / len(pred_answers) is_consensus = consensus_ratio >= self.consensus_threshold return is_consensus, consensus_ratio, common_answer","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._initialize_debate","name":"_initialize_debate","kind":"method","path":"src/multi_agent_debate.py","span":"L363-L373","signature":"def _initialize_debate(self, queries: List[Query]) -> None","exported":false,"origin":"ast","body_hash":"de216185871f8011dbca1cb1679ae11bd6745556b9f45f627aed5174977679e3","body_text":"def _initialize_debate(self, queries: List[Query]) -> None: for query in queries: self.current_rounds[query.query_id] = 0 self.debate_log[query.query_id] = { \"query\": query.content, \"answer\": query.answer, \"rounds\": [] } self.consensus_flags[query.query_id] = False for agent in self.agents: agent.initialize_conversation(query)","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","name":"_clear_debate","kind":"method","path":"src/multi_agent_debate.py","span":"L375-L382","signature":"def _clear_debate(self) -> None","exported":false,"origin":"ast","body_hash":"270ce9f7a246709d312fb2c2b2d5fcf64f844823d74fc946994ad0b17bc5df85","body_text":"def _clear_debate(self) -> None: self.current_rounds.clear() self.agent_responses.clear() self.debate_log.clear() self.consensus_flags.clear() self._prev_resp_embeddings.clear() for agent in self.agents: agent.clear_all_conversations()","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","name":"_save_memory_data","kind":"method","path":"src/multi_agent_debate.py","span":"L388-L452","signature":"def _save_memory_data(self, queries: List[Query], batch_idx: int = 0) -> None","exported":false,"origin":"ast","body_hash":"736be302119f63037a9f7e9222eb407cccc17381866d59079b35f7871493f56a","body_text":"def _save_memory_data(self, queries: List[Query], batch_idx: int = 0) -> None: \"\"\" After a completed debate, write one JSON line per (query, agent, round) to MEMORY_DATA_DIR/<task_type>/agent{n}_batch{idx:03d}_{timestamp}.jsonl. \"\"\" task_dir = os.path.join(MEMORY_DATA_DIR, self.task_type, self.model.model_name) os.makedirs(task_dir, exist_ok=True) timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\") filename = f\"agent{len(self.agents)}_batch{batch_idx:03d}_{timestamp}.jsonl\" filepath = os.path.join(task_dir, filename) with open(filepath, \"w\", encoding=\"utf-8\") as f: for query in queries: qid = query.query_id log = self.debate_log[qid] is_final_correct = log[\"is_final_correct\"] rounds_log = log[\"rounds\"] num_rounds = len(rounds_log) for agent in self.agents: aid = agent.agent_id for round_idx in range(num_rounds): round_data = self.agent_responses[qid][aid][f\"round_{round_idx}\"] current_response = round_data[\"response\"] if not current_response: continue is_current_correct = round_data[\"is_correct\"] if round_idx > 0: prev_response = self.agent_responses[qid][aid][f\"round_{round_idx - 1}\"][\"response\"] prev_pred = self.agent_responses[qid][aid][f\"round_{round_idx - 1}\"][\"pred_answer\"] stance_changed = (prev_pred != round_data[\"pred_answer\"]) prev_consensus_ratio = rounds_log[round_idx - 1][\"consensus_ratio\"] prev_round_summary = rounds_log[round_idx - 1].get(\"summary\", None) else: prev_response = None stance_changed = None prev_consensus_ratio = None prev_round_summary = None memory_case = { \"db_id\": f\"Q{qid}_Round{round_idx}_{aid}\", \"query_id\": qid, \"agent_id\": str(aid), \"round\": round_idx, \"query\": query.content, \"answer\": query.answer, \"solution\": query.solution, \"category\": query.category, \"prev_response\": prev_response, \"prev_round_summary\": prev_round_summary, \"prev_consensus_ratio\": prev_consensus_ratio, \"current_response\": current_response, \"is_current_correct\": is_current_correct, \"stance_changed\": stance_changed, \"is_final_correct\": is_final_correct, } f.write(json.dumps(memory_case, ensure_ascii=False) + \"\\n\")","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","name":"save_debate_log","kind":"method","path":"src/multi_agent_debate.py","span":"L454-L457","signature":"def save_debate_log(self, filename: str) -> None","exported":true,"origin":"ast","body_hash":"55df815a09208eb287f1673f79ec48ace2ac212630c32d66c90229a1a27a1cd2","body_text":"def save_debate_log(self, filename: str) -> None: os.makedirs(os.path.dirname(filename), exist_ok=True) with open(filename, \"w\", encoding=\"utf-8\") as f: json.dump(self.debate_log, f, ensure_ascii=False, indent=2)","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"},{"id":"src/multi_agent_debate.py#MultiAgentDebate.get_token_usage_summary","name":"get_token_usage_summary","kind":"method","path":"src/multi_agent_debate.py","span":"L459-L460","signature":"def get_token_usage_summary(self)","exported":true,"origin":"ast","body_hash":"61cb8dd073485b509287cb1bf05553edd4cd4fde28c40e096d0616cd81dee054","body_text":"def get_token_usage_summary(self): return self.model.get_token_usage_summary()","summary_state":"pending","summary":null,"crux":null,"owner":"MultiAgentDebate"}],"rawEdges":[{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"os","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"re","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"json","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"datetime","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"typing","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"collections","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"src.agent","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"src.models","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"src.checker","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"src.paths","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"imports","specifier":"src.prompts","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"contains","targetId":"src/multi_agent_debate.py#remove_think_tags","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#remove_think_tags","relation":"calls","name":"sub","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate.__init__","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.__init__","relation":"calls","name":"defaultdict","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.__init__","relation":"calls","name":"defaultdict","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.__init__","relation":"calls","name":"defaultdict","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.__init__","relation":"calls","name":"dict","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate.run","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"_clear_debate","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"_initialize_debate","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"range","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"print","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"_run_debate_round","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.run","relation":"calls","name":"_save_memory_data","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"defaultdict"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"batch_retrieve_all","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"items","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"enumerate","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"build_messages","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"model","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"strip","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"remove_think_tags","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"next","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"array","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"embed_fn","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"enumerate","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"parse_answer","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"check_answer","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"record_response","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"_check_consensus","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"str","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"str","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"float","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"items","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"print","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._run_debate_round","relation":"calls","name":"_summarize_debate_state","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"len","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"enumerate","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"lower","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"str","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"join","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"format","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"append","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"model","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"MultiAgentDebate"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"zip","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"strip","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._summarize_debate_state","relation":"calls","name":"remove_think_tags","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"strip","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"lower","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"len","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"Counter","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"most_common","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"Counter"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._check_consensus","relation":"calls","name":"len","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._initialize_debate","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._initialize_debate","relation":"calls","name":"initialize_conversation","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"defaultdict"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear","viaMember":true,"file":"src/multi_agent_debate.py","recvType":"defaultdict"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._clear_debate","relation":"calls","name":"clear_all_conversations","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"join","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"makedirs","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"strftime","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"now","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"len","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"join","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"open","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"len","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"range","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"get","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"str","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"write","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate._save_memory_data","relation":"calls","name":"dumps","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","relation":"calls","name":"makedirs","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","relation":"calls","name":"dirname","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","relation":"calls","name":"open","viaMember":false,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.save_debate_log","relation":"calls","name":"dump","viaMember":true,"file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate","relation":"contains","targetId":"src/multi_agent_debate.py#MultiAgentDebate.get_token_usage_summary","file":"src/multi_agent_debate.py"},{"source":"src/multi_agent_debate.py#MultiAgentDebate.get_token_usage_summary","relation":"calls","name":"get_token_usage_summary","viaMember":true,"file":"src/multi_agent_debate.py"}]},"src/paths.py":{"size":786,"mtimeMs":1789362123401.8032,"hash":"f9a1a7bb9a0c875f7137dfd63158dd641436250c5e1ba47f136af7d8481a1f99","nodes":[{"id":"src/paths.py","name":"paths.py","kind":"file","path":"src/paths.py","span":"L1-L21","signature":null,"exported":true,"origin":"ast","body_hash":"f9a1a7bb9a0c875f7137dfd63158dd641436250c5e1ba47f136af7d8481a1f99","chars":786,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os # Project root (this file lives at src/paths.py, so root is one level up) ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # Raw datasets downloaded from HuggingFace RAW_DATA_DIR = os.path.join(ROOT_DIR, \"data\") # Processed / formatted datasets ready for DataLoader PROCESSED_DATA_DIR = os.path.join(ROOT_DIR, \"processed_data\") # Debate result JSON logs LOG_DIR = os.path.join(ROOT_DIR, \"logs\") # JSONL files written during training for memory construction MEMORY_DATA_DIR = os.path.join(ROOT_DIR, \"memory_data\") # ChromaDB persistent vector store for the memory module; also holds pre-computed # document embedding .npz files (one per collection, named {collection}_doc_embeddings.npz) PERSISTENT_DIR = os.path.join(ROOT_DIR, \"persistent_memory\")"}],"rawEdges":[{"source":"src/paths.py","relation":"imports","specifier":"os","file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"dirname","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"dirname","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"abspath","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"join","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"join","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"join","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"join","viaMember":true,"file":"src/paths.py"},{"source":"src/paths.py","relation":"calls","name":"join","viaMember":true,"file":"src/paths.py"}]},"src/prompts.py":{"size":19564,"mtimeMs":1789362123401.8032,"hash":"3fb6dcf3d0f043a590af9800ed0d6ce28bd8cd6d18d58138de77e0fb50057a87","nodes":[{"id":"src/prompts.py","name":"prompts.py","kind":"file","path":"src/prompts.py","span":"L1-L244","signature":null,"exported":true,"origin":"ast","body_hash":"3fb6dcf3d0f043a590af9800ed0d6ce28bd8cd6d18d58138de77e0fb50057a87","chars":19560,"summary_state":"pending","summary":null,"crux":null,"body_text":"MATH500_SYS_PROMPTS = [ ( \"You are a theoretical mathematics professor with a rigorous approach to problem-solving. You excel in formal proofs and mathematical reasoning. \" \"Always verify assumptions, consider edge cases, and provide step-by-step logical arguments. Focus on theoretical foundations and mathematical principles. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), ( \"You are a practical mathematics problem-solving expert with extensive experience in competitive mathematics. You excel at finding efficient solutions and spotting patterns quickly. \" \"Focus on problem-solving strategies, shortcuts, and alternative approaches. Challenge assumptions when necessary. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), ( \"You are an experienced mathematics educator who excels at breaking down complex problems. You focus on clear explanations, visual representations, and multiple solution methods. \" \"Always connect concepts to fundamental principles and similar problems. Validate solutions through different approaches. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), ( \"You are an intuitive mathematics expert who excels at rapid problem-solving and pattern recognition. You have exceptional ability to see the core of problems and find elegant solutions. \" \"While maintaining mathematical rigor, you prefer concise and creative approaches over lengthy formal proofs. Think fast, be decisive, and trust your mathematical intuition. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), ( \"You are an analytical mathematics expert who specializes in systematic problem decomposition and rigorous logical reasoning. You excel at breaking complex problems into manageable steps and identifying key mathematical relationships. \" \"Always validate your reasoning through careful analysis and consider potential edge cases. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), ( \"You are a critical thinking mathematics expert who excels at questioning assumptions and validating solutions. You carefully examine problems from multiple angles, challenge conventional approaches, and verify conclusions. \" \"Focus on finding potential flaws in reasoning and exploring alternative solutions. Always consider whether the answer makes sense in context. Be concise and focus only on essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ) ] ENGINEERING_SYS_PROMPTS = [ ( \"You're a theoretical engineering expert with deep knowledge in engineering principles, physics, and mathematical modeling. Focus on fundamental laws, governing equations, and conceptual frameworks when analyzing problems. \" \"Always ground your answers in established engineering theories and first principles. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're a hands-on engineering practitioner with extensive experience in design, troubleshooting, and real-world systems. Focus on practical constraints, material properties, manufacturing considerations, and industry standards when analyzing problems. \" \"Draw on engineering experience and domain-specific heuristics to identify the most feasible solution. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're a multidisciplinary engineering consultant with expertise spanning mechanical, electrical, civil, and systems engineering. Approach problems by integrating cross-domain knowledge and considering interactions between subsystems. \" \"Balance theoretical rigor with practical engineering judgment in your analysis. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an engineering expert with strong critical thinking and safety-oriented mindset. Always evaluate failure modes, identify hidden assumptions, and rigorously check unit consistency and boundary conditions. \" \"Challenge seemingly obvious answers by considering edge cases and worst-case scenarios. When uncertain, explicitly state your confidence level. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an engineering expert who excels at structured problem-solving. Approach each question by: 1) Identifying known quantities and governing principles, 2) Setting up the appropriate equations or models, 3) Solving step-by-step with dimensional analysis, 4) Validating the result against physical intuition and constraints. \" \"Always organize your reasoning methodically. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an engineering expert who specializes in estimation and order-of-magnitude reasoning. Evaluate problems by identifying dominant effects, simplifying secondary factors, and bounding the solution space before committing to a precise answer. \" \"Use dimensional analysis and scaling arguments to quickly eliminate implausible options. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ) ] ECONOMICS_SYS_PROMPTS = [ ( \"You're a theoretical economics expert with deep knowledge in economic principles and models. Focus on fundamental theories, mathematical relationships, and conceptual frameworks when analyzing problems. \" \"Always support your answers with established economic theories. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an empirical economics researcher specializing in data analysis and real-world economic phenomena. Focus on historical examples, empirical evidence, and practical applications when analyzing problems. \" \"Consider real market behaviors and outcomes in your reasoning. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're a comprehensive economics consultant with expertise in both theoretical and applied economics. Approach problems by considering multiple perspectives, including behavioral economics insights and institutional factors. \" \"Balance theoretical principles with practical implications in your analysis. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an economics expert with strong critical thinking skills. Always evaluate multiple possibilities before making decisions, identify potential logical flaws, and challenge common assumptions. \" \"Consider edge cases and counterarguments in your analysis. When uncertain, explicitly state your confidence level and reasoning. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an economics expert who excels at structured problem-solving. Approach each question by: 1) Breaking down complex concepts into basic components, 2) Analyzing each component systematically, 3) Identifying key relationships and dependencies, 4) Drawing logical conclusions based on the analysis. \" \"Always organize your thoughts step-by-step. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ), ( \"You're an economics expert who specializes in probabilistic reasoning. Evaluate problems by considering multiple scenarios and their likelihood. Think in terms of probability distributions rather than absolutes. \" \"Consider both the most likely outcome and potential alternative scenarios. Use Bayesian-style updating when processing information. Keep your reasoning concise and focus only on the essential steps necessary to reach the conclusion. \" \"Provide your final answer in double parentheses: ((answer)), where answer can be A, B, C, D, E, F, G, H, I, or J.\" ) ] TRUTHFULQA_SYS_PROMPTS = [ ( \"You're a fact-checking expert trained to distinguish truth from popular misconceptions. Many widely believed statements are false — your job is to identify what is actually true, not what sounds plausible or is commonly assumed. \" \"Be especially skeptical of answers that align with urban legends, folk wisdom, or intuitive-sounding claims that lack factual basis. Keep your reasoning concise and focus only on the essential steps. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ), ( \"You're a critical epistemologist who specializes in identifying false beliefs that are widely held. Your primary instinct is to question what 'everyone knows' and verify claims against evidence. \" \"When a question seems to have an obvious answer, treat that as a warning sign — TruthfulQA questions are specifically designed to test whether you will echo misinformation. Reason carefully before committing. Keep your reasoning concise. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ), ( \"You're a domain-spanning knowledge expert with deep familiarity across science, history, law, medicine, and culture. Your strength is knowing the actual consensus or established facts in each domain, even when they contradict popular belief. \" \"Approach each question by recalling the authoritative understanding of the topic, then select the option that aligns with verified truth rather than common assumption. Keep your reasoning concise. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ), ( \"You're a logical analyst who evaluates claims with rigorous skepticism. You never accept a statement just because it 'sounds right' or is frequently repeated. Instead, you trace each claim back to its logical foundations and identify which option is actually defensible. \" \"Pay particular attention to options that are partially true but misleading in the given context. Keep your reasoning concise and focus only on the essential steps. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ), ( \"You're a calibrated reasoner who is acutely aware of your own potential to reproduce misinformation. Before answering, you actively consider: 'Is this a common misconception I might be inclined to repeat?' \" \"You deliberately slow down on questions that feel easy or obvious, and verify your reasoning against what is actually known to be true rather than what is culturally prevalent. Keep your reasoning concise. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ), ( \"You're a truthfulness-focused expert who prioritizes factual accuracy above all else, even when the truth is counterintuitive or goes against majority belief. You are trained to resist sycophancy and never select an answer simply because it is popular or expected. \" \"Evaluate each option on its factual merit alone, and explicitly reject options that are plausible-sounding but empirically unsupported. Keep your reasoning concise and focus only on the essential steps. \" \"Provide your final answer in double parentheses: ((answer)), where answer is the letter of the correct option from the given choices.\" ) ] MAD_SYS_PROMPTS = { \"MATH500\": MATH500_SYS_PROMPTS, \"MMLUPro_Engineering\": ENGINEERING_SYS_PROMPTS, \"MMLUPro_Economics\": ECONOMICS_SYS_PROMPTS, \"TruthfulQA\": TRUTHFULQA_SYS_PROMPTS, } TASK_PROMPT = { \"MATH500\": ( \"Can you solve the following math question as accurately as possible?\\n\\n\" \"<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Present your analysis concisely using only essential reasoning steps. \" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), \"MMLUPro_Engineering\": ( \"Please analyze the following engineering question:\\n\\n\" \"<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Present your analysis concisely using only essential reasoning steps and select the correct answer. \" \"Provide the final answer in double parentheses at the end of your response: The answer is ((answer)).\" ), \"MMLUPro_Economics\": ( \"Please analyze the following economics question:\\n\\n\" \"<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Present your analysis concisely using only essential reasoning steps and select the correct answer. \" \"Provide the final answer in double parentheses at the end of your response: The answer is ((answer)).\" ), \"TruthfulQA\": ( \"Please carefully evaluate the following question and select the most truthful answer:\\n\\n\" \"<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Be wary of options that sound plausible but reflect common misconceptions. Present your analysis concisely and focus only on what is factually accurate. \" \"Provide the final answer in double parentheses at the end of your response: The answer is ((answer)).\" ), } DEBATE_PROMPT = { \"MATH500\": ( \"Use the solutions from other agents as additional information, can you give an updated answer? \" \"The original question is: \\n\\n<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Provide the final answer in the following format at the end of your response: The answer is \\\\boxed{{answer}}.\" ), \"MMLUPro_Engineering\": ( \"Use the solutions from other agents as additional information, can you give an updated answer? \" \"The original question is: \\n\\n<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Provide the final answer in the following format at the end of your response: The answer is ((answer)).\" ), \"MMLUPro_Economics\": ( \"Use the solutions from other agents as additional information, can you give an updated answer? \" \"The original question is: \\n\\n<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Provide the final answer in the following format at the end of your response: The answer is ((answer)).\" ), \"TruthfulQA\": ( \"Use the solutions from other agents as additional information, can you give an updated answer? \" \"The original question is: \\n\\n<question>\\n{QUESTION}\\n</question>\\n\\n\" \"Provide the final answer in double"}],"rawEdges":[]},"src/retrieval_policy.py":{"size":12593,"mtimeMs":1789362123401.8032,"hash":"1558b7bfef2267803bdc1a09c1c8cdfabc88e04a7d0d1aeb04c59785d02c8492","nodes":[{"id":"src/retrieval_policy.py","name":"retrieval_policy.py","kind":"file","path":"src/retrieval_policy.py","span":"L1-L329","signature":null,"exported":true,"origin":"ast","body_hash":"1558b7bfef2267803bdc1a09c1c8cdfabc88e04a7d0d1aeb04c59785d02c8492","chars":12585,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"\"\" Pluggable memory-retrieval policies. A retrieval policy owns the *whole* two-stage retrieval — candidate generation (the ChromaDB query) and selection — and returns, per query, both the retrieved cases (which get formatted into the prompt) and the full candidate pool (which the Agent reuses for memory-estimated confidence). Keeping the candidates in the return value is what keeps confidence coupled to whatever cases the policy actually retrieved (see Agent.batch_retrieve_all). Policies operate purely on the raw ChromaDB candidate tuples Candidate = (metadata: Dict, document: str, distance: float, doc_embedding: Optional[np.ndarray]) and primitives, so this module does not import from `agent.py` (avoids a circular import). The Agent converts selected candidates into RetrievedMemory objects. The debate-state-aware policy is the method described in the paper; the other three (similarity / random / diversity) exist as retrieval-policy ablations. \"\"\" import inspect import numpy as np from typing import Dict, List, Optional, Tuple Candidate = Tuple[Dict, str, float, Optional[np.ndarray]] # Sentinel distance used for cases retrieved without a similarity notion (random # policy). 1.0 keeps the logged similarity (1 - distance) at 0.0. _NO_SIMILARITY_DISTANCE = 1.0 # ---------------------------------------------------------------------- # Shared selection helpers # ---------------------------------------------------------------------- # ---------------------------------------------------------------------- # Policies # ---------------------------------------------------------------------- POLICY_REGISTRY = { DebateStateAwarePolicy.name: DebateStateAwarePolicy, SimilarityPolicy.name: SimilarityPolicy, DiversityPolicy.name: DiversityPolicy, RandomPolicy.name: RandomPolicy, FixedLambdaPolicy.name: FixedLambdaPolicy, PositiveOnlyPolicy.name: PositiveOnlyPolicy, }"},{"id":"src/retrieval_policy.py#_doc_embedding_matrix","name":"_doc_embedding_matrix","kind":"function","path":"src/retrieval_policy.py","span":"L37-L49","signature":"def _doc_embedding_matrix(candidates: List[Candidate], embedding_model=None) -> np.ndarray","exported":false,"origin":"ast","body_hash":"051374ee79d1a7859c5d713d368c8f74faf93231ad32acf3554ab0f5aeb2a6f9","body_text":"def _doc_embedding_matrix(candidates: List[Candidate], embedding_model=None) -> np.ndarray: \"\"\" Stack candidate document embeddings into an (n, dim) matrix. Uses the pre-computed embeddings loaded from the .npz sidecar when all are present; otherwise falls back to encoding the documents on the fly. \"\"\" precomputed = [c[3] for c in candidates] if all(e is not None for e in precomputed): return np.stack(precomputed) if embedding_model is None: raise ValueError(\"Document embeddings missing and no embedding_model provided for fallback\") doc_texts = [c[1] for c in candidates] return np.array(embedding_model.encode(doc_texts))","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#mmr_select","name":"mmr_select","kind":"function","path":"src/retrieval_policy.py","span":"L52-L100","signature":"def mmr_select( candidates: List[Candidate], n: int, lam: float, embedding_model=None, ) -> List[Candidate]","exported":true,"origin":"ast","body_hash":"f68d23967ae804b0e85af3f0c5a73e978e99a5a1b86d3a14052524220943de47","body_text":"def mmr_select( candidates: List[Candidate], n: int, lam: float, embedding_model=None, ) -> List[Candidate]: \"\"\" Maximal Marginal Relevance greedy selection (debate-state-aware policy). Greedily pick n cases maximizing lambda * sim(e, q) * r(e) - (1 - lambda) * max_{e' in selected} sim(e', e) where sim(e, q) = 1 - distance, r(e) is case correctness (0/1), and the inter-candidate similarity is cosine of the document embeddings. Returns the selected candidate tuples in selection order. \"\"\" if not candidates: return [] sim_to_query = np.array([1.0 - c[2] for c in candidates]) case_correct = np.array([ 1.0 if c[0].get(\"response_correct\", False) else 0.0 for c in candidates ]) doc_embeddings = _doc_embedding_matrix(candidates, embedding_model) sim_matrix = doc_embeddings @ doc_embeddings.T # (n_candidates, n_candidates) n = min(n, len(candidates)) selected_indices = [] remaining = list(range(len(candidates))) for _ in range(n): if not selected_indices: max_sim_to_selected = np.zeros(len(candidates)) else: max_sim_to_selected = sim_matrix[:, selected_indices].max(axis=1) scores = np.full(len(candidates), -np.inf) for i in remaining: scores[i] = ( lam * sim_to_query[i] * case_correct[i] - (1.0 - lam) * max_sim_to_selected[i] ) best = int(np.argmax(scores)) selected_indices.append(best) remaining.remove(best) return [candidates[i] for i in selected_indices]","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#diversity_select","name":"diversity_select","kind":"function","path":"src/retrieval_policy.py","span":"L103-L131","signature":"def diversity_select( candidates: List[Candidate], n: int, embedding_model=None, seed: Optional[int] = None, ) -> List[Candidate]","exported":true,"origin":"ast","body_hash":"379a737d109924a63050fcbc8dced49799485b7675f06864b942d2be582685e6","body_text":"def diversity_select( candidates: List[Candidate], n: int, embedding_model=None, seed: Optional[int] = None, ) -> List[Candidate]: \"\"\" Pure diversity selection: anchor on a random candidate, then greedily add the candidate that minimizes its maximum cosine similarity to those already selected. Ignores query relevance and case correctness. \"\"\" if not candidates: return [] doc_embeddings = _doc_embedding_matrix(candidates, embedding_model) sim_matrix = doc_embeddings @ doc_embeddings.T # Anchor the diverse set on a randomly chosen candidate, so the selection carries no query-similarity bias. rng = np.random.default_rng(seed) anchor = int(rng.integers(len(candidates))) n = min(n, len(candidates)) selected_indices = [anchor] remaining = [i for i in range(len(candidates)) if i != anchor] while len(selected_indices) < n and remaining: best = min(remaining, key=lambda i: float(sim_matrix[i, selected_indices].max())) selected_indices.append(best) remaining.remove(best) return [candidates[i] for i in selected_indices]","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#RetrievalPolicy","name":"RetrievalPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L138-L179","signature":"class RetrievalPolicy","exported":true,"origin":"ast","body_hash":"16fcc257ea18e5dbd2211127974448a306ff153ee59797d4d3951695cc3705e0","body_text":"class RetrievalPolicy: \"\"\"Base class. Subclasses implement the full two-stage retrieval.\"\"\" name: str = \"base\" def __init__(self, seed: Optional[int] = None): # Stateful RNG seeded once from the base seed. self._rng = np.random.default_rng(seed) def batch_retrieve( self, memory_db, agent_id: str, queries: List[Tuple[int, str]], n_retrieval: int, prev_consensus_ratios: Dict[int, Optional[float]], ) -> Dict[int, Tuple[List[Candidate], List[Candidate]]]: \"\"\" Args: memory_db: MemoryVectorDB instance. agent_id: str(agent_id) used to filter the collection. queries: [(query_id, embedding_text)] for queries needing retrieval. n_retrieval: number of cases to select per query. prev_consensus_ratios: {query_id: prev-round consensus ratio or None}. Returns: {query_id: (retrieved, candidates)} where `retrieved` are the cases selected for the prompt and `candidates` is the full candidate pool (retrieved ⊆ candidates). \"\"\" raise NotImplementedError def _next_seed(self) -> int: \"\"\"Draw a fresh child seed from the policy's RNG.\"\"\" return int(self._rng.integers(np.iinfo(np.int64).max)) @staticmethod def _filter(agent_id: str, correct_only: bool = False) -> Dict: base = {\"agent_id\": {\"$eq\": agent_id}} if correct_only: return {\"$and\": [base, {\"response_correct\": {\"$eq\": True}}]} return base","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#RetrievalPolicy.__init__","name":"__init__","kind":"method","path":"src/retrieval_policy.py","span":"L143-L145","signature":"def __init__(self, seed: Optional[int] = None): # Stateful RNG seeded once from the base seed.","exported":false,"origin":"ast","body_hash":"8518c3ab8ab0762ffb18608ced1e1e5886789736dad048ef7867b161b8f369e6","body_text":"def __init__(self, seed: Optional[int] = None): # Stateful RNG seeded once from the base seed. self._rng = np.random.default_rng(seed)","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievalPolicy"},{"id":"src/retrieval_policy.py#RetrievalPolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L147-L168","signature":"def batch_retrieve( self, memory_db, agent_id: str, queries: List[Tuple[int, str]], n_retrieval: int, prev_consensus_ratios: Dict[int, Optional[float]], ) -> Dict[int, Tuple[List[Candidate], List[Candidate]]]","exported":true,"origin":"ast","body_hash":"0677ad6bda9d9a071447fcf48c07860e0f7ca2434b605234d20e1f26256131f5","body_text":"def batch_retrieve( self, memory_db, agent_id: str, queries: List[Tuple[int, str]], n_retrieval: int, prev_consensus_ratios: Dict[int, Optional[float]], ) -> Dict[int, Tuple[List[Candidate], List[Candidate]]]: \"\"\" Args: memory_db: MemoryVectorDB instance. agent_id: str(agent_id) used to filter the collection. queries: [(query_id, embedding_text)] for queries needing retrieval. n_retrieval: number of cases to select per query. prev_consensus_ratios: {query_id: prev-round consensus ratio or None}. Returns: {query_id: (retrieved, candidates)} where `retrieved` are the cases selected for the prompt and `candidates` is the full candidate pool (retrieved ⊆ candidates). \"\"\" raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievalPolicy"},{"id":"src/retrieval_policy.py#RetrievalPolicy._next_seed","name":"_next_seed","kind":"method","path":"src/retrieval_policy.py","span":"L170-L172","signature":"def _next_seed(self) -> int","exported":false,"origin":"ast","body_hash":"54c86ed39419c65b9544b4222809d695e63bed421f7c57252a7347a3d0af6953","body_text":"def _next_seed(self) -> int: \"\"\"Draw a fresh child seed from the policy's RNG.\"\"\" return int(self._rng.integers(np.iinfo(np.int64).max))","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievalPolicy"},{"id":"src/retrieval_policy.py#RetrievalPolicy._filter","name":"_filter","kind":"method","path":"src/retrieval_policy.py","span":"L175-L179","signature":"def _filter(agent_id: str, correct_only: bool = False) -> Dict","exported":false,"origin":"ast","body_hash":"a45f18bcf835361b4e32d72e8eef0e9d0b7cd903204afbffe80c157ed6efa1ca","body_text":"def _filter(agent_id: str, correct_only: bool = False) -> Dict: base = {\"agent_id\": {\"$eq\": agent_id}} if correct_only: return {\"$and\": [base, {\"response_correct\": {\"$eq\": True}}]} return base","summary_state":"pending","summary":null,"crux":null,"owner":"RetrievalPolicy"},{"id":"src/retrieval_policy.py#DebateStateAwarePolicy","name":"DebateStateAwarePolicy","kind":"class","path":"src/retrieval_policy.py","span":"L182-L218","signature":"class DebateStateAwarePolicy(RetrievalPolicy)","exported":true,"origin":"ast","body_hash":"a0db61d106b4a0dc6c7dbb2adcfbfa9a81f08d773ec0028c9ed8470eef2b6542","body_text":"class DebateStateAwarePolicy(RetrievalPolicy): \"\"\"Semantic top-3n candidates + adaptive MMR re-ranking. With adaptive_lambda=True (paper method), the MMR relevance/diversity tradeoff lambda = 1 - gamma * consensus_ratio adapts to the previous round's consensus. With adaptive_lambda=False (fixed_lambda ablation), lambda is held constant at ref_lambda regardless of consensus. \"\"\" name = \"debate_state_aware\" def __init__(self, gamma: float = 0.9, ref_lambda: float = 0.7, adaptive_lambda: bool = True, seed: Optional[int] = None): super().__init__(seed) self.gamma = gamma self.ref_lambda = ref_lambda self.adaptive_lambda = adaptive_lambda def _lambda_for(self, consensus_ratio: Optional[float]) -> float: if self.adaptive_lambda and consensus_ratio is not None: return 1.0 - self.gamma * consensus_ratio return self.ref_lambda def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) out = {} for (qid, _), candidates in zip(queries, batch_candidates): lam = self._lambda_for(prev_consensus_ratios.get(qid)) retrieved = mmr_select(candidates, n_retrieval, lam, memory_db.embedding_model) out[qid] = (retrieved, candidates) return out","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#DebateStateAwarePolicy.__init__","name":"__init__","kind":"method","path":"src/retrieval_policy.py","span":"L192-L197","signature":"def __init__(self, gamma: float = 0.9, ref_lambda: float = 0.7, adaptive_lambda: bool = True, seed: Optional[int] = None)","exported":false,"origin":"ast","body_hash":"b347c44f2abcf2a33768adaf3b85018a2c1802f0ff4cda54fe5884a322e661e5","body_text":"def __init__(self, gamma: float = 0.9, ref_lambda: float = 0.7, adaptive_lambda: bool = True, seed: Optional[int] = None): super().__init__(seed) self.gamma = gamma self.ref_lambda = ref_lambda self.adaptive_lambda = adaptive_lambda","summary_state":"pending","summary":null,"crux":null,"owner":"DebateStateAwarePolicy"},{"id":"src/retrieval_policy.py#DebateStateAwarePolicy._lambda_for","name":"_lambda_for","kind":"method","path":"src/retrieval_policy.py","span":"L199-L202","signature":"def _lambda_for(self, consensus_ratio: Optional[float]) -> float","exported":false,"origin":"ast","body_hash":"585b408209eaff5d78db715e4f76200b8388d3c512fbce526a2903db954611eb","body_text":"def _lambda_for(self, consensus_ratio: Optional[float]) -> float: if self.adaptive_lambda and consensus_ratio is not None: return 1.0 - self.gamma * consensus_ratio return self.ref_lambda","summary_state":"pending","summary":null,"crux":null,"owner":"DebateStateAwarePolicy"},{"id":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L204-L218","signature":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios)","exported":true,"origin":"ast","body_hash":"09ddf4c24d3782e75422c9b340d21feb67c5a7b41db8cb0f8d56781a53679c4e","body_text":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) out = {} for (qid, _), candidates in zip(queries, batch_candidates): lam = self._lambda_for(prev_consensus_ratios.get(qid)) retrieved = mmr_select(candidates, n_retrieval, lam, memory_db.embedding_model) out[qid] = (retrieved, candidates) return out","summary_state":"pending","summary":null,"crux":null,"owner":"DebateStateAwarePolicy"},{"id":"src/retrieval_policy.py#FixedLambdaPolicy","name":"FixedLambdaPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L221-L226","signature":"class FixedLambdaPolicy(DebateStateAwarePolicy)","exported":true,"origin":"ast","body_hash":"d3a6f5616cdd1cccfed4b2dff04d10cc1b4d26f88e4b6f9adead607232800ace","body_text":"class FixedLambdaPolicy(DebateStateAwarePolicy): \"\"\"Debate_state_aware MMR with a constant lambda (no consensus adaptation).\"\"\" name = \"fixed_lambda\" def __init__(self, ref_lambda: float = 0.7, seed: Optional[int] = None): super().__init__(ref_lambda=ref_lambda, adaptive_lambda=False, seed=seed)","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#FixedLambdaPolicy.__init__","name":"__init__","kind":"method","path":"src/retrieval_policy.py","span":"L225-L226","signature":"def __init__(self, ref_lambda: float = 0.7, seed: Optional[int] = None)","exported":false,"origin":"ast","body_hash":"79a16f95e31d0f8ee5dd04b6f9489ee7460358f6c79d77fb1c747995ced4f2e8","body_text":"def __init__(self, ref_lambda: float = 0.7, seed: Optional[int] = None): super().__init__(ref_lambda=ref_lambda, adaptive_lambda=False, seed=seed)","summary_state":"pending","summary":null,"crux":null,"owner":"FixedLambdaPolicy"},{"id":"src/retrieval_policy.py#SimilarityPolicy","name":"SimilarityPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L229-L243","signature":"class SimilarityPolicy(RetrievalPolicy)","exported":true,"origin":"ast","body_hash":"2977f8db6085550f085aa0295e18e8cced281f00108563d7cf36c678d257a634","body_text":"class SimilarityPolicy(RetrievalPolicy): \"\"\"Plain semantic top-n retrieval (no re-ranking, no correctness).\"\"\" name = \"similarity\" def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#SimilarityPolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L233-L243","signature":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios)","exported":true,"origin":"ast","body_hash":"8c6d45b89937d3a2fdb56db040630067383711f5aa00855a9e9180e4ff279bf1","body_text":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null,"owner":"SimilarityPolicy"},{"id":"src/retrieval_policy.py#DiversityPolicy","name":"DiversityPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L246-L263","signature":"class DiversityPolicy(RetrievalPolicy)","exported":true,"origin":"ast","body_hash":"edaba640087d895b438d9979f09511a5f24cc0b29c53e5f621450338d06a2c5b","body_text":"class DiversityPolicy(RetrievalPolicy): \"\"\"Semantic top-3n candidates + pure diversity selection.\"\"\" name = \"diversity\" def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) out = {} for (qid, _), candidates in zip(queries, batch_candidates): retrieved = diversity_select(candidates, n_retrieval, memory_db.embedding_model, seed=self._next_seed()) out[qid] = (retrieved, candidates) return out","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L250-L263","signature":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios)","exported":true,"origin":"ast","body_hash":"ebc706b88245f78ec82e5c798e7355dcc9a480d5ea410e5808f2f74ebe35ed70","body_text":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), ) out = {} for (qid, _), candidates in zip(queries, batch_candidates): retrieved = diversity_select(candidates, n_retrieval, memory_db.embedding_model, seed=self._next_seed()) out[qid] = (retrieved, candidates) return out","summary_state":"pending","summary":null,"crux":null,"owner":"DiversityPolicy"},{"id":"src/retrieval_policy.py#RandomPolicy","name":"RandomPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L266-L281","signature":"class RandomPolicy(RetrievalPolicy)","exported":true,"origin":"ast","body_hash":"d8aa53ffaa49aed53dcf4cafc1639123639b8c73df812d15d7389bee89a9c721","body_text":"class RandomPolicy(RetrievalPolicy): \"\"\"Uniformly random n cases from the whole collection (ignores query).\"\"\" name = \"random\" def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): batch_candidates = memory_db.batch_get_random( n_queries=len(queries), n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), distance=_NO_SIMILARITY_DISTANCE, seed=self._next_seed(), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L270-L281","signature":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios)","exported":true,"origin":"ast","body_hash":"0477f7d254cbe15737ecd0d1e605fab7d398837b10eb68bfdfb598700300ac77","body_text":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): batch_candidates = memory_db.batch_get_random( n_queries=len(queries), n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id), distance=_NO_SIMILARITY_DISTANCE, seed=self._next_seed(), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null,"owner":"RandomPolicy"},{"id":"src/retrieval_policy.py#PositiveOnlyPolicy","name":"PositiveOnlyPolicy","kind":"class","path":"src/retrieval_policy.py","span":"L284-L298","signature":"class PositiveOnlyPolicy(RetrievalPolicy)","exported":true,"origin":"ast","body_hash":"cbc3bb5f6ac89b5d0260c152fd70156c2242205e9fad35c24a5e1d14e8166b3d","body_text":"class PositiveOnlyPolicy(RetrievalPolicy): \"\"\"Semantic retrieval restricted to correct (positive) cases only.\"\"\" name = \"positive_only\" def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id, correct_only=True), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null},{"id":"src/retrieval_policy.py#PositiveOnlyPolicy.batch_retrieve","name":"batch_retrieve","kind":"method","path":"src/retrieval_policy.py","span":"L288-L298","signature":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios)","exported":true,"origin":"ast","body_hash":"649a04a62f607f281687059fb7859f5f3922642829c1593e9f9301e5b0a4c0f0","body_text":"def batch_retrieve(self, memory_db, agent_id, queries, n_retrieval, prev_consensus_ratios): embedding_texts = [emb_text for _, emb_text in queries] batch_candidates = memory_db.batch_query_similar( query_texts=embedding_texts, n_results=n_retrieval * 3, filter_metadata=self._filter(agent_id, correct_only=True), ) return { qid: (candidates[:n_retrieval], candidates) for (qid, _), candidates in zip(queries, batch_candidates) }","summary_state":"pending","summary":null,"crux":null,"owner":"PositiveOnlyPolicy"},{"id":"src/retrieval_policy.py#build_retrieval_policy","name":"build_retrieval_policy","kind":"function","path":"src/retrieval_policy.py","span":"L311-L328","signature":"def build_retrieval_policy(name: str, seed: Optional[int] = None, **policy_kwargs) -> RetrievalPolicy","exported":true,"origin":"ast","body_hash":"cd614b4aca064a0e86a9307efc0e109d2f094bb0773358fec10f8e17d01d6612","body_text":"def build_retrieval_policy(name: str, seed: Optional[int] = None, **policy_kwargs) -> RetrievalPolicy: \"\"\" Instantiate a policy by name. `seed` is passed to every policy; any extra policy_kwargs (e.g. ref_lambda, gamma) are forwarded only to policies whose constructor accepts them, and only when not None — so callers can pass a superset of knobs regardless of which policy is selected. \"\"\" if name not in POLICY_REGISTRY: raise ValueError( f\"Unknown retrieval_policy '{name}'. Available: {sorted(POLICY_REGISTRY)}\" ) cls = POLICY_REGISTRY[name] accepted = inspect.signature(cls.__init__).parameters kwargs = {\"seed\": seed} for key, value in policy_kwargs.items(): if value is not None and key in accepted: kwargs[key] = value return cls(**kwargs)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"src/retrieval_policy.py","relation":"imports","specifier":"inspect","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"imports","specifier":"typing","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#_doc_embedding_matrix","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#_doc_embedding_matrix","relation":"calls","name":"all","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#_doc_embedding_matrix","relation":"calls","name":"stack","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#_doc_embedding_matrix","relation":"calls","name":"ValueError","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#_doc_embedding_matrix","relation":"calls","name":"array","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#_doc_embedding_matrix","relation":"calls","name":"encode","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#mmr_select","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"array","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"array","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"get","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"_doc_embedding_matrix","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"min","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"list","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"range","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"range","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"zeros","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"max","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"full","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"int","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"argmax","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"append","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#mmr_select","relation":"calls","name":"remove","viaMember":true,"file":"src/retrieval_policy.py","recvType":"list"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#diversity_select","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"_doc_embedding_matrix","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"default_rng","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"int","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"integers","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"min","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"range","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"min","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"float","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"max","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"append","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#diversity_select","relation":"calls","name":"remove","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy","relation":"contains","targetId":"src/retrieval_policy.py#RetrievalPolicy.__init__","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy.__init__","relation":"calls","name":"default_rng","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy","relation":"contains","targetId":"src/retrieval_policy.py#RetrievalPolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy","relation":"contains","targetId":"src/retrieval_policy.py#RetrievalPolicy._next_seed","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy._next_seed","relation":"calls","name":"int","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy._next_seed","relation":"calls","name":"integers","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy._next_seed","relation":"calls","name":"iinfo","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RetrievalPolicy","relation":"contains","targetId":"src/retrieval_policy.py#RetrievalPolicy._filter","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#DebateStateAwarePolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy","relation":"extends","name":"RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy","relation":"contains","targetId":"src/retrieval_policy.py#DebateStateAwarePolicy.__init__","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.__init__","relation":"calls","name":"super","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy","relation":"contains","targetId":"src/retrieval_policy.py#DebateStateAwarePolicy._lambda_for","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy","relation":"contains","targetId":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"batch_query_similar","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"_filter","viaMember":true,"file":"src/retrieval_policy.py","recvType":"DebateStateAwarePolicy"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"zip","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"_lambda_for","viaMember":true,"file":"src/retrieval_policy.py","recvType":"DebateStateAwarePolicy"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"get","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DebateStateAwarePolicy.batch_retrieve","relation":"calls","name":"mmr_select","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#FixedLambdaPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#FixedLambdaPolicy","relation":"extends","name":"DebateStateAwarePolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#FixedLambdaPolicy","relation":"contains","targetId":"src/retrieval_policy.py#FixedLambdaPolicy.__init__","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#FixedLambdaPolicy.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#FixedLambdaPolicy.__init__","relation":"calls","name":"super","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#SimilarityPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#SimilarityPolicy","relation":"extends","name":"RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#SimilarityPolicy","relation":"contains","targetId":"src/retrieval_policy.py#SimilarityPolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#SimilarityPolicy.batch_retrieve","relation":"calls","name":"batch_query_similar","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#SimilarityPolicy.batch_retrieve","relation":"calls","name":"_filter","viaMember":true,"file":"src/retrieval_policy.py","recvType":"SimilarityPolicy"},{"source":"src/retrieval_policy.py#SimilarityPolicy.batch_retrieve","relation":"calls","name":"zip","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#DiversityPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy","relation":"extends","name":"RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy","relation":"contains","targetId":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","relation":"calls","name":"batch_query_similar","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","relation":"calls","name":"_filter","viaMember":true,"file":"src/retrieval_policy.py","recvType":"DiversityPolicy"},{"source":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","relation":"calls","name":"zip","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","relation":"calls","name":"diversity_select","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#DiversityPolicy.batch_retrieve","relation":"calls","name":"_next_seed","viaMember":true,"file":"src/retrieval_policy.py","recvType":"DiversityPolicy"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#RandomPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RandomPolicy","relation":"extends","name":"RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RandomPolicy","relation":"contains","targetId":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","relation":"calls","name":"batch_get_random","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","relation":"calls","name":"len","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","relation":"calls","name":"_filter","viaMember":true,"file":"src/retrieval_policy.py","recvType":"RandomPolicy"},{"source":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","relation":"calls","name":"_next_seed","viaMember":true,"file":"src/retrieval_policy.py","recvType":"RandomPolicy"},{"source":"src/retrieval_policy.py#RandomPolicy.batch_retrieve","relation":"calls","name":"zip","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#PositiveOnlyPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#PositiveOnlyPolicy","relation":"extends","name":"RetrievalPolicy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#PositiveOnlyPolicy","relation":"contains","targetId":"src/retrieval_policy.py#PositiveOnlyPolicy.batch_retrieve","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#PositiveOnlyPolicy.batch_retrieve","relation":"calls","name":"batch_query_similar","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#PositiveOnlyPolicy.batch_retrieve","relation":"calls","name":"_filter","viaMember":true,"file":"src/retrieval_policy.py","recvType":"PositiveOnlyPolicy"},{"source":"src/retrieval_policy.py#PositiveOnlyPolicy.batch_retrieve","relation":"calls","name":"zip","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py","relation":"contains","targetId":"src/retrieval_policy.py#build_retrieval_policy","file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#build_retrieval_policy","relation":"calls","name":"ValueError","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#build_retrieval_policy","relation":"calls","name":"sorted","viaMember":false,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#build_retrieval_policy","relation":"calls","name":"signature","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#build_retrieval_policy","relation":"calls","name":"items","viaMember":true,"file":"src/retrieval_policy.py"},{"source":"src/retrieval_policy.py#build_retrieval_policy","relation":"calls","name":"cls","viaMember":false,"file":"src/retrieval_policy.py"}]},"src/utils/__init__.py":{"size":353,"mtimeMs":1789362123401.8032,"hash":"05e808dbbc28bfe1949c29d8285a3757f4d8f151d3181655b95e54cc3ebe81a7","nodes":[{"id":"src/utils/__init__.py","name":"__init__.py","kind":"file","path":"src/utils/__init__.py","span":"L1-L15","signature":null,"exported":true,"origin":"ast","body_hash":"05e808dbbc28bfe1949c29d8285a3757f4d8f151d3181655b95e54cc3ebe81a7","chars":353,"summary_state":"pending","summary":null,"crux":null,"body_text":"from src.utils.config_utils import ( load_configs_from_yaml, load_env_file, LLMConfig, EmbeddingConfig, ) from src.utils.data_utils import DataLoader from src.utils.model_utils import TokenUsageTracker from src.utils.memory_utils import ( read_memory_data, sample_memories, construct_memory, construct_response_memory, )"}],"rawEdges":[{"source":"src/utils/__init__.py","relation":"imports","specifier":"src.utils.config_utils","file":"src/utils/__init__.py"},{"source":"src/utils/__init__.py","relation":"imports","specifier":"src.utils.data_utils","file":"src/utils/__init__.py"},{"source":"src/utils/__init__.py","relation":"imports","specifier":"src.utils.model_utils","file":"src/utils/__init__.py"},{"source":"src/utils/__init__.py","relation":"imports","specifier":"src.utils.memory_utils","file":"src/utils/__init__.py"}]},"src/utils/config_utils.py":{"size":2779,"mtimeMs":1789362123402.8035,"hash":"603a2525252af49e23e7dfef349e2a65fb06ccbce77d9a5255ee1cb05f11b5c7","nodes":[{"id":"src/utils/config_utils.py","name":"config_utils.py","kind":"file","path":"src/utils/config_utils.py","span":"L1-L63","signature":null,"exported":true,"origin":"ast","body_hash":"603a2525252af49e23e7dfef349e2a65fb06ccbce77d9a5255ee1cb05f11b5c7","chars":2779,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import yaml"},{"id":"src/utils/config_utils.py#load_configs_from_yaml","name":"load_configs_from_yaml","kind":"function","path":"src/utils/config_utils.py","span":"L5-L7","signature":"def load_configs_from_yaml(file_path)","exported":true,"origin":"ast","body_hash":"7a6021cadc21664d1a7319acf8182aed72e4f8d99ee684333d52777a1554744d","body_text":"def load_configs_from_yaml(file_path): with open(file_path, \"r\") as f: return yaml.safe_load(f)","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/config_utils.py#load_env_file","name":"load_env_file","kind":"function","path":"src/utils/config_utils.py","span":"L10-L35","signature":"def load_env_file(path: str = \".env\") -> None","exported":true,"origin":"ast","body_hash":"befb1758d4bcca7a7505b3b5ee548cc79118f048f373a08cf96374e3086fbbb4","body_text":"def load_env_file(path: str = \".env\") -> None: \"\"\" Lightweight, zero-dependency loader for a `.env` file (KEY=VALUE lines). Reads simple `KEY=value` / `export KEY=value` lines, ignoring blanks and `#` comments and stripping surrounding quotes. Variables already present in the environment are NOT overwritten, so an explicit shell `export` still wins. Missing file is a no-op. \"\"\" if not os.path.exists(path): return with open(path, \"r\", encoding=\"utf-8\") as f: for raw_line in f: line = raw_line.strip() if not line or line.startswith(\"#\"): continue if line.startswith(\"export \"): line = line[len(\"export \"):].strip() if \"=\" not in line: continue key, value = line.split(\"=\", 1) key = key.strip() value = value.strip().strip('\"').strip(\"'\") if key and key not in os.environ: os.environ[key] = value","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/config_utils.py#LLMConfig","name":"LLMConfig","kind":"class","path":"src/utils/config_utils.py","span":"L38-L56","signature":"class LLMConfig","exported":true,"origin":"ast","body_hash":"6685f1133d2d51274cb5c725bec3354ffc77b2aa630d36505d42d7e6560c296d","body_text":"class LLMConfig: def __init__(self, general_config: dict, model_config: dict, name: str = \"\"): self.name = name # Serving backend: \"vllm\" (local GPU) or \"openrouter\" (remote API). self.backend = model_config.get(\"backend\", general_config.get(\"backend\", \"vllm\")) self.model = model_config.get(\"model\", \"Qwen/Qwen2.5-7B-Instruct\") self.model_path = model_config.get(\"model_path\", None) self.max_token_length = model_config.get(\"max_token_length\", general_config.get(\"max_token_length\", 24576)) self.temperature = model_config.get(\"temperature\", general_config.get(\"temperature\", 1)) self.top_p = model_config.get(\"top_p\", general_config.get(\"top_p\", 1)) self.max_tokens = model_config.get(\"max_tokens\", general_config.get(\"max_tokens\", 6144)) # API backend (openrouter) settings; ignored by the vLLM backend. self.api_base = model_config.get(\"api_base\", general_config.get(\"api_base\", \"https://openrouter.ai/api/v1\")) self.api_key_env = model_config.get(\"api_key_env\", general_config.get(\"api_key_env\", \"OPENROUTER_API_KEY\")) self.max_workers = model_config.get(\"max_workers\", general_config.get(\"max_workers\", 12)) self.max_retries = model_config.get(\"max_retries\", general_config.get(\"max_retries\", 5)) self.request_timeout = model_config.get(\"request_timeout\", general_config.get(\"request_timeout\", 120))","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/config_utils.py#LLMConfig.__init__","name":"__init__","kind":"method","path":"src/utils/config_utils.py","span":"L39-L56","signature":"def __init__(self, general_config: dict, model_config: dict, name: str = \"\")","exported":false,"origin":"ast","body_hash":"e188f40705b8b6f65aff2cbf211f02ca2f41ffe95a09e823ae5d0e9aebb8557c","body_text":"def __init__(self, general_config: dict, model_config: dict, name: str = \"\"): self.name = name # Serving backend: \"vllm\" (local GPU) or \"openrouter\" (remote API). self.backend = model_config.get(\"backend\", general_config.get(\"backend\", \"vllm\")) self.model = model_config.get(\"model\", \"Qwen/Qwen2.5-7B-Instruct\") self.model_path = model_config.get(\"model_path\", None) self.max_token_length = model_config.get(\"max_token_length\", general_config.get(\"max_token_length\", 24576)) self.temperature = model_config.get(\"temperature\", general_config.get(\"temperature\", 1)) self.top_p = model_config.get(\"top_p\", general_config.get(\"top_p\", 1)) self.max_tokens = model_config.get(\"max_tokens\", general_config.get(\"max_tokens\", 6144)) # API backend (openrouter) settings; ignored by the vLLM backend. self.api_base = model_config.get(\"api_base\", general_config.get(\"api_base\", \"https://openrouter.ai/api/v1\")) self.api_key_env = model_config.get(\"api_key_env\", general_config.get(\"api_key_env\", \"OPENROUTER_API_KEY\")) self.max_workers = model_config.get(\"max_workers\", general_config.get(\"max_workers\", 12)) self.max_retries = model_config.get(\"max_retries\", general_config.get(\"max_retries\", 5)) self.request_timeout = model_config.get(\"request_timeout\", general_config.get(\"request_timeout\", 120))","summary_state":"pending","summary":null,"crux":null,"owner":"LLMConfig"},{"id":"src/utils/config_utils.py#EmbeddingConfig","name":"EmbeddingConfig","kind":"class","path":"src/utils/config_utils.py","span":"L59-L62","signature":"class EmbeddingConfig","exported":true,"origin":"ast","body_hash":"b5af39439d8aad4736f7632d232cf00b8d68e2c5127905ef7163d42b94dd99bf","body_text":"class EmbeddingConfig: def __init__(self, model_config: dict): self.model = model_config.get(\"model\", \"bge-m3\") self.model_path = model_config.get(\"model_path\", None)","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/config_utils.py#EmbeddingConfig.__init__","name":"__init__","kind":"method","path":"src/utils/config_utils.py","span":"L60-L62","signature":"def __init__(self, model_config: dict)","exported":false,"origin":"ast","body_hash":"0c7214c369751b1f1830b6d812e748a587e84efd3d437a1f8e50d5596e26c11e","body_text":"def __init__(self, model_config: dict): self.model = model_config.get(\"model\", \"bge-m3\") self.model_path = model_config.get(\"model_path\", None)","summary_state":"pending","summary":null,"crux":null,"owner":"EmbeddingConfig"}],"rawEdges":[{"source":"src/utils/config_utils.py","relation":"imports","specifier":"os","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py","relation":"imports","specifier":"yaml","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py","relation":"contains","targetId":"src/utils/config_utils.py#load_configs_from_yaml","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_configs_from_yaml","relation":"calls","name":"open","viaMember":false,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_configs_from_yaml","relation":"calls","name":"safe_load","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py","relation":"contains","targetId":"src/utils/config_utils.py#load_env_file","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"exists","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"open","viaMember":false,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"startswith","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"startswith","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"len","viaMember":false,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"split","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#load_env_file","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py","relation":"contains","targetId":"src/utils/config_utils.py#LLMConfig","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#LLMConfig","relation":"contains","targetId":"src/utils/config_utils.py#LLMConfig.__init__","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#LLMConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py","relation":"contains","targetId":"src/utils/config_utils.py#EmbeddingConfig","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#EmbeddingConfig","relation":"contains","targetId":"src/utils/config_utils.py#EmbeddingConfig.__init__","file":"src/utils/config_utils.py"},{"source":"src/utils/config_utils.py#EmbeddingConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"},{"source":"src/utils/config_utils.py#EmbeddingConfig.__init__","relation":"calls","name":"get","viaMember":true,"file":"src/utils/config_utils.py","recvType":"dict"}]},"src/utils/data_utils.py":{"size":3247,"mtimeMs":1789362123402.8035,"hash":"313e2b80449313b7b56e329e4318bd2a8ef2efa9c7201b6fae518c685acb21eb","nodes":[{"id":"src/utils/data_utils.py","name":"data_utils.py","kind":"file","path":"src/utils/data_utils.py","span":"L1-L87","signature":null,"exported":true,"origin":"ast","body_hash":"313e2b80449313b7b56e329e4318bd2a8ef2efa9c7201b6fae518c685acb21eb","chars":3247,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import json from typing import Dict from src.paths import PROCESSED_DATA_DIR"},{"id":"src/utils/data_utils.py#read_json","name":"read_json","kind":"function","path":"src/utils/data_utils.py","span":"L8-L13","signature":"def read_json(file_path)","exported":true,"origin":"ast","body_hash":"7d82552eadbaf954df59bfae303d3208942dd61e8310212f66b2d7d883fd4420","body_text":"def read_json(file_path): file_path = str(file_path) with open(file_path, \"r\", encoding=\"utf-8\") as f: if file_path.endswith(\".jsonl\"): return [json.loads(line) for line in f if line.strip()] return json.load(f)","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/data_utils.py#format_choices","name":"format_choices","kind":"function","path":"src/utils/data_utils.py","span":"L16-L17","signature":"def format_choices(options: Dict[str, str]) -> str","exported":true,"origin":"ast","body_hash":"c7f63b6c25d42134a1ff7d063d8419faca3e38e0bb81f153e0d6b7cea2c9e963","body_text":"def format_choices(options: Dict[str, str]) -> str: return '\\n'.join([idx + '. ' + item for idx, item in list(options.items())])","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/data_utils.py#DataLoader","name":"DataLoader","kind":"class","path":"src/utils/data_utils.py","span":"L20-L86","signature":"class DataLoader","exported":true,"origin":"ast","body_hash":"63d5a6df97381c28bdfaaf98ea25df0b966876a47bf01e220b28607f10525c40","body_text":"class DataLoader: @staticmethod def load_dataset(question_type: str) -> Dict: dataset_paths = { \"MATH500\": \"math500/test.jsonl\", \"MATH500_TRAIN\": \"math500/train.jsonl\", \"MMLUPro_Engineering\": \"mmlu_pro/engineering/test.jsonl\", \"MMLUPro_Engineering_TRAIN\": \"mmlu_pro/engineering/train.jsonl\", \"MMLUPro_Economics\": \"mmlu_pro/economics/test.jsonl\", \"MMLUPro_Economics_TRAIN\": \"mmlu_pro/economics/train.jsonl\", \"TruthfulQA\": \"truthfulqa/test.jsonl\", \"TruthfulQA_TRAIN\": \"truthfulqa/train.jsonl\", } if question_type not in dataset_paths: raise ValueError(f\"DataLoader: Unsupported question type in reading data: {question_type}\") dataset = read_json(os.path.join(PROCESSED_DATA_DIR, dataset_paths[question_type])) return dataset @staticmethod def format_question(item: Dict, question_type: str) -> tuple: formatters = { \"MATH500\": lambda x: (x['question'], x['answer'], x['solution'], x['category']), \"MATH500_TRAIN\": lambda x: (x['question'], x['answer'], x['solution'], x['category']), \"MMLUPro_Engineering\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"engineering\" ), \"MMLUPro_Engineering_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"engineering\" ), \"MMLUPro_Economics\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"economics\" ), \"MMLUPro_Economics_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"economics\" ), \"TruthfulQA\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"truthfulqa\" ), \"TruthfulQA_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"truthfulqa\" ), } if question_type not in formatters: raise ValueError(f\"DataLoader: Unsupported question type in format question: {question_type}\") return formatters[question_type](item)","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/data_utils.py#DataLoader.load_dataset","name":"load_dataset","kind":"method","path":"src/utils/data_utils.py","span":"L22-L38","signature":"def load_dataset(question_type: str) -> Dict","exported":true,"origin":"ast","body_hash":"5c4ad755f0ea1b82198f4ce497e76c05866baf69ddefe0598a909e40e2ee797a","body_text":"def load_dataset(question_type: str) -> Dict: dataset_paths = { \"MATH500\": \"math500/test.jsonl\", \"MATH500_TRAIN\": \"math500/train.jsonl\", \"MMLUPro_Engineering\": \"mmlu_pro/engineering/test.jsonl\", \"MMLUPro_Engineering_TRAIN\": \"mmlu_pro/engineering/train.jsonl\", \"MMLUPro_Economics\": \"mmlu_pro/economics/test.jsonl\", \"MMLUPro_Economics_TRAIN\": \"mmlu_pro/economics/train.jsonl\", \"TruthfulQA\": \"truthfulqa/test.jsonl\", \"TruthfulQA_TRAIN\": \"truthfulqa/train.jsonl\", } if question_type not in dataset_paths: raise ValueError(f\"DataLoader: Unsupported question type in reading data: {question_type}\") dataset = read_json(os.path.join(PROCESSED_DATA_DIR, dataset_paths[question_type])) return dataset","summary_state":"pending","summary":null,"crux":null,"owner":"DataLoader"},{"id":"src/utils/data_utils.py#DataLoader.format_question","name":"format_question","kind":"method","path":"src/utils/data_utils.py","span":"L41-L86","signature":"def format_question(item: Dict, question_type: str) -> tuple","exported":true,"origin":"ast","body_hash":"d3b4c3d278537982d4efa7851e71759f6806cfb686b62360263aa4794dc62932","body_text":"def format_question(item: Dict, question_type: str) -> tuple: formatters = { \"MATH500\": lambda x: (x['question'], x['answer'], x['solution'], x['category']), \"MATH500_TRAIN\": lambda x: (x['question'], x['answer'], x['solution'], x['category']), \"MMLUPro_Engineering\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"engineering\" ), \"MMLUPro_Engineering_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"engineering\" ), \"MMLUPro_Economics\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"economics\" ), \"MMLUPro_Economics_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"economics\" ), \"TruthfulQA\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"truthfulqa\" ), \"TruthfulQA_TRAIN\": lambda x: ( f\"Question: {x['question']}\\n\\n{format_choices(x['options'])}\", x['answer'], x['answer'], \"truthfulqa\" ), } if question_type not in formatters: raise ValueError(f\"DataLoader: Unsupported question type in format question: {question_type}\") return formatters[question_type](item)","summary_state":"pending","summary":null,"crux":null,"owner":"DataLoader"}],"rawEdges":[{"source":"src/utils/data_utils.py","relation":"imports","specifier":"os","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"imports","specifier":"json","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"imports","specifier":"typing","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"imports","specifier":"src.paths","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"contains","targetId":"src/utils/data_utils.py#read_json","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"str","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"open","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"endswith","viaMember":true,"file":"src/utils/data_utils.py","recvType":"str"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"loads","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#read_json","relation":"calls","name":"load","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"contains","targetId":"src/utils/data_utils.py#format_choices","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#format_choices","relation":"calls","name":"join","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#format_choices","relation":"calls","name":"list","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#format_choices","relation":"calls","name":"items","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py","relation":"contains","targetId":"src/utils/data_utils.py#DataLoader","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader","relation":"contains","targetId":"src/utils/data_utils.py#DataLoader.load_dataset","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.load_dataset","relation":"calls","name":"ValueError","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.load_dataset","relation":"calls","name":"read_json","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.load_dataset","relation":"calls","name":"join","viaMember":true,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader","relation":"contains","targetId":"src/utils/data_utils.py#DataLoader.format_question","file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"format_choices","viaMember":false,"file":"src/utils/data_utils.py"},{"source":"src/utils/data_utils.py#DataLoader.format_question","relation":"calls","name":"ValueError","viaMember":false,"file":"src/utils/data_utils.py"}]},"src/utils/memory_utils.py":{"size":5555,"mtimeMs":1789362123402.8035,"hash":"92f21d688833626c850a88ae741d197128c742f52a18d1a34cfa290135010a24","nodes":[{"id":"src/utils/memory_utils.py","name":"memory_utils.py","kind":"file","path":"src/utils/memory_utils.py","span":"L1-L163","signature":null,"exported":true,"origin":"ast","body_hash":"92f21d688833626c850a88ae741d197128c742f52a18d1a34cfa290135010a24","chars":5555,"summary_state":"pending","summary":null,"crux":null,"body_text":"import json import random from pathlib import Path from typing import Dict, List, Tuple # --------------------------------------------------------------------------- # Memory construction # ---------------------------------------------------------------------------"},{"id":"src/utils/memory_utils.py#read_memory_data","name":"read_memory_data","kind":"function","path":"src/utils/memory_utils.py","span":"L7-L22","signature":"def read_memory_data(data_dir: str) -> List[Dict]","exported":true,"origin":"ast","body_hash":"bf7c30d24bb48869b24979d61202805e2059cb8c07df1476c4dbf9a59626dded","body_text":"def read_memory_data(data_dir: str) -> List[Dict]: \"\"\" Read all JSONL files under data_dir into a list of dicts, sorted by db_id. \"\"\" rows = [] for jsonl_file in sorted(Path(data_dir).glob(\"*.jsonl\")): with open(jsonl_file, \"r\", encoding=\"utf-8\") as f: for line in f: line = line.strip() if not line: continue rows.append(json.loads(line)) rows.sort(key=lambda r: r[\"db_id\"]) return rows","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#sample_memories","name":"sample_memories","kind":"function","path":"src/utils/memory_utils.py","span":"L25-L45","signature":"def sample_memories(records: List[Dict], sample_size: int = -1, seed: int = 42) -> List[Dict]","exported":true,"origin":"ast","body_hash":"5e541819fd34a08ad8019cb2a3063cdc378968e5abe72280dcbf364696b86b6a","body_text":"def sample_memories(records: List[Dict], sample_size: int = -1, seed: int = 42) -> List[Dict]: \"\"\" Sample up to sample_size unique questions (all records for each sampled question are retained). sample_size == -1 means keep all. Round-0 cases are excluded before sampling since memory is only used from round 1. \"\"\" random.seed(seed) records = [r for r in records if int(r.get(\"round\", 0)) > 0] question_ids = sorted({r[\"query_id\"] for r in records}) random.shuffle(question_ids) if sample_size == -1 or sample_size >= len(question_ids): print(\"All memory questions sampled.\") sample_size = len(question_ids) selected_query_ids = set(question_ids[:sample_size]) records_sample = [r for r in records if r[\"query_id\"] in selected_query_ids] return records_sample","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#_build_embedding_text","name":"_build_embedding_text","kind":"function","path":"src/utils/memory_utils.py","span":"L52-L74","signature":"def _build_embedding_text(row: Dict) -> str","exported":false,"origin":"ast","body_hash":"deec4471a2b074f0fb01a9ccd45c24fa8ef8b4f2a67c45b161f39d16f5cff206","body_text":"def _build_embedding_text(row: Dict) -> str: \"\"\" Construct the retrieval key (debate state) from: 1. task x (query) 2. agent previous response z_i^{t-1} 3. previous debate summary \"\"\" embedding_text = \"\" # 1. task embedding_text += f\"{row['query']}\\n\" # 2. agent previous response (None for round 0) prev_response = row.get(\"prev_response\") if prev_response: embedding_text += f\"{prev_response}\\n\" # 3. previous debate summary (None for round 0) prev_summary = row.get(\"prev_round_summary\") if prev_summary: embedding_text += f\"{prev_summary}\\n\" return embedding_text.strip()","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#_build_value_text","name":"_build_value_text","kind":"function","path":"src/utils/memory_utils.py","span":"L77-L114","signature":"def _build_value_text(row: Dict) -> str","exported":false,"origin":"ast","body_hash":"737bb8b1fd3e97d826821fa7c00dfe0d09e8cdfa69719636fef2774b501ad608","body_text":"def _build_value_text(row: Dict) -> str: \"\"\" Construct the retrieval value (prompt context) from: 1. task x (query) 2. correct solution z^* 3. agent current response z_i^t 4. whether agent changed response 5. agent current correctness All parts are wrapped in XML tags. \"\"\" value_text = \"\" # 1. task value_text += f\"<example_question>\\n{row['query']}\\n</example_question>\\n\\n\" # 2. correct solution value_text += f\"<example_solution>\\n{row['solution']}\\n</example_solution>\\n\\n\" # 3. agent current response value_text += f\"<example_agent_response>\\n{row['current_response']}\\n</example_agent_response>\\n\\n\" # 4. whether agent changed response (N/A for round 0) stance_changed = row.get(\"stance_changed\") if stance_changed is None: changed_text = \"N/A (initial round)\" else: changed_text = \"The agent changed its stance.\" if stance_changed else \"The agent did not change its stance.\" value_text += f\"<example_agent_stance_changed>\\n{changed_text}\\n</example_agent_stance_changed>\\n\\n\" # 5. agent current correctness is_correct = row.get(\"is_current_correct\", False) if is_correct: correct_text = \"The response is correct.\" else: correct_text = \"The response is incorrect.\" value_text += f\"<example_agent_response_correctness>\\n{correct_text}\\n</example_agent_response_correctness>\" return value_text","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#_build_meta_data","name":"_build_meta_data","kind":"function","path":"src/utils/memory_utils.py","span":"L117-L132","signature":"def _build_meta_data(row: Dict) -> Dict","exported":false,"origin":"ast","body_hash":"e580541cd371d7d686abed035aebb241184606caa2e13c3333b8ee167986a1d2","body_text":"def _build_meta_data(row: Dict) -> Dict: \"\"\" Construct the ChromaDB metadata dict. All values must be ChromaDB-compatible scalars (str / int / float / bool). \"\"\" prev_consensus = row.get(\"prev_consensus_ratio\") return { \"category\": str(row.get(\"category\", \"\")), \"question_id\": str(row.get(\"query_id\", \"\")), \"agent_id\": str(row.get(\"agent_id\", \"\")), \"current_round\": int(row.get(\"round\", 0)), \"previous_consensus_ratio\": float(prev_consensus) if prev_consensus is not None else -1.0, \"response_correct\": bool(row.get(\"is_current_correct\", False)), \"final_correct\": bool(row.get(\"is_final_correct\", False)), }","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#construct_memory","name":"construct_memory","kind":"function","path":"src/utils/memory_utils.py","span":"L135-L145","signature":"def construct_memory(row: Dict) -> Tuple[str, str, str, Dict]","exported":true,"origin":"ast","body_hash":"3f0558ac870afba53aad5c98c63ceb960dee470a1c26bc010517a106e43c9955","body_text":"def construct_memory(row: Dict) -> Tuple[str, str, str, Dict]: \"\"\" Build all four components needed to store one memory entry: db_id, embedding_text, value_text, meta_data \"\"\" db_id = row[\"db_id\"] embedding_text = _build_embedding_text(row) value_text = _build_value_text(row) meta_data = _build_meta_data(row) return db_id, embedding_text, value_text, meta_data","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/memory_utils.py#construct_response_memory","name":"construct_response_memory","kind":"function","path":"src/utils/memory_utils.py","span":"L148-L162","signature":"def construct_response_memory(row: Dict) -> Tuple[str, str, Dict]","exported":true,"origin":"ast","body_hash":"ba65a0d9337bfb03c35525dc6474fd69b9f25d2908298bdc62caafec656554e6","body_text":"def construct_response_memory(row: Dict) -> Tuple[str, str, Dict]: \"\"\" Build the (db_id, response_text, metadata) triple for a single memory row's current_response, to be stored in the response-embedding collection. \"\"\" response_text = row[\"current_response\"] meta_data = { \"category\": str(row.get(\"category\", \"\")), \"question_id\": str(row.get(\"query_id\", \"\")), \"agent_id\": str(row.get(\"agent_id\", \"\")), \"round\": int(row.get(\"round\", 0)), \"response_correct\": bool(row.get(\"is_current_correct\", False)), \"final_correct\": bool(row.get(\"is_final_correct\", False)), } return row[\"db_id\"], response_text, meta_data","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"src/utils/memory_utils.py","relation":"imports","specifier":"json","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"imports","specifier":"random","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"imports","specifier":"pathlib","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"imports","specifier":"typing","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#read_memory_data","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"sorted","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"glob","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"Path","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"open","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"append","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"loads","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#read_memory_data","relation":"calls","name":"sort","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#sample_memories","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"seed","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"int","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"sorted","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"shuffle","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"len","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"print","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"len","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#sample_memories","relation":"calls","name":"set","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#_build_embedding_text","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_embedding_text","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_embedding_text","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_embedding_text","relation":"calls","name":"strip","viaMember":true,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#_build_value_text","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_value_text","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_value_text","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#_build_meta_data","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"int","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"float","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"bool","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"bool","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#_build_meta_data","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#construct_memory","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_memory","relation":"calls","name":"_build_embedding_text","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_memory","relation":"calls","name":"_build_value_text","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_memory","relation":"calls","name":"_build_meta_data","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py","relation":"contains","targetId":"src/utils/memory_utils.py#construct_response_memory","file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"str","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"int","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"bool","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"bool","viaMember":false,"file":"src/utils/memory_utils.py"},{"source":"src/utils/memory_utils.py#construct_response_memory","relation":"calls","name":"get","viaMember":true,"file":"src/utils/memory_utils.py","recvType":"Dict"}]},"src/utils/model_utils.py":{"size":1775,"mtimeMs":1789362123402.8035,"hash":"98726ac6541d4594806f30d8541e942556fbea39b5094fcb1b7f68083621f599","nodes":[{"id":"src/utils/model_utils.py","name":"model_utils.py","kind":"file","path":"src/utils/model_utils.py","span":"L1-L45","signature":null,"exported":true,"origin":"ast","body_hash":"98726ac6541d4594806f30d8541e942556fbea39b5094fcb1b7f68083621f599","chars":1775,"summary_state":"pending","summary":null,"crux":null,"body_text":""},{"id":"src/utils/model_utils.py#TokenUsageTracker","name":"TokenUsageTracker","kind":"class","path":"src/utils/model_utils.py","span":"L1-L44","signature":"class TokenUsageTracker","exported":true,"origin":"ast","body_hash":"3e4f9f8cbbb9b71c6a46811fe109b46f51248fe7fe6737e00e3427551a5f65da","body_text":"class TokenUsageTracker: \"\"\"Tracks token usage and calculates costs, optionally broken down by task.\"\"\" def __init__(self): self.total_input_tokens = 0 self.total_output_tokens = 0 self.usage_history = [] self.task_usage = {} # task_name -> {\"input_tokens\": int, \"output_tokens\": int} def update_usage(self, model: str, input_tokens: int, output_tokens: int, task_name: str = \"default\"): usage_record = { \"model\": model, \"task\": task_name, \"input_tokens\": input_tokens, \"output_tokens\": output_tokens, \"total_tokens\": input_tokens + output_tokens, } self.total_input_tokens += input_tokens self.total_output_tokens += output_tokens self.usage_history.append(usage_record) if task_name not in self.task_usage: self.task_usage[task_name] = {\"input_tokens\": 0, \"output_tokens\": 0} self.task_usage[task_name][\"input_tokens\"] += input_tokens self.task_usage[task_name][\"output_tokens\"] += output_tokens def get_summary(self): by_task = { task: { \"input_tokens\": stats[\"input_tokens\"], \"output_tokens\": stats[\"output_tokens\"], \"total_tokens\": stats[\"input_tokens\"] + stats[\"output_tokens\"], } for task, stats in self.task_usage.items() } return { \"total_input_tokens\": self.total_input_tokens, \"total_output_tokens\": self.total_output_tokens, \"total_tokens\": self.total_input_tokens + self.total_output_tokens, \"call_count\": len(self.usage_history), \"by_task\": by_task, \"history\": self.usage_history, }","summary_state":"pending","summary":null,"crux":null},{"id":"src/utils/model_utils.py#TokenUsageTracker.__init__","name":"__init__","kind":"method","path":"src/utils/model_utils.py","span":"L3-L8","signature":"def __init__(self)","exported":false,"origin":"ast","body_hash":"7c666694739bb2a5778458a0db131941a7e41ccbbb292d223c858394d1c2feea","body_text":"def __init__(self): self.total_input_tokens = 0 self.total_output_tokens = 0 self.usage_history = [] self.task_usage = {} # task_name -> {\"input_tokens\": int, \"output_tokens\": int}","summary_state":"pending","summary":null,"crux":null,"owner":"TokenUsageTracker"},{"id":"src/utils/model_utils.py#TokenUsageTracker.update_usage","name":"update_usage","kind":"method","path":"src/utils/model_utils.py","span":"L10-L26","signature":"def update_usage(self, model: str, input_tokens: int, output_tokens: int, task_name: str = \"default\")","exported":true,"origin":"ast","body_hash":"7eeaaeaeb4962cf53890b36e04bceaf6c9b60100c3d05f2f205f189eea72482d","body_text":"def update_usage(self, model: str, input_tokens: int, output_tokens: int, task_name: str = \"default\"): usage_record = { \"model\": model, \"task\": task_name, \"input_tokens\": input_tokens, \"output_tokens\": output_tokens, \"total_tokens\": input_tokens + output_tokens, } self.total_input_tokens += input_tokens self.total_output_tokens += output_tokens self.usage_history.append(usage_record) if task_name not in self.task_usage: self.task_usage[task_name] = {\"input_tokens\": 0, \"output_tokens\": 0} self.task_usage[task_name][\"input_tokens\"] += input_tokens self.task_usage[task_name][\"output_tokens\"] += output_tokens","summary_state":"pending","summary":null,"crux":null,"owner":"TokenUsageTracker"},{"id":"src/utils/model_utils.py#TokenUsageTracker.get_summary","name":"get_summary","kind":"method","path":"src/utils/model_utils.py","span":"L28-L44","signature":"def get_summary(self)","exported":true,"origin":"ast","body_hash":"cb5cc405cf0369e235122826445f509ba8a7cc5acd0d9fba75cf12abe9c51632","body_text":"def get_summary(self): by_task = { task: { \"input_tokens\": stats[\"input_tokens\"], \"output_tokens\": stats[\"output_tokens\"], \"total_tokens\": stats[\"input_tokens\"] + stats[\"output_tokens\"], } for task, stats in self.task_usage.items() } return { \"total_input_tokens\": self.total_input_tokens, \"total_output_tokens\": self.total_output_tokens, \"total_tokens\": self.total_input_tokens + self.total_output_tokens, \"call_count\": len(self.usage_history), \"by_task\": by_task, \"history\": self.usage_history, }","summary_state":"pending","summary":null,"crux":null,"owner":"TokenUsageTracker"}],"rawEdges":[{"source":"src/utils/model_utils.py","relation":"contains","targetId":"src/utils/model_utils.py#TokenUsageTracker","file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker","relation":"contains","targetId":"src/utils/model_utils.py#TokenUsageTracker.__init__","file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker","relation":"contains","targetId":"src/utils/model_utils.py#TokenUsageTracker.update_usage","file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker.update_usage","relation":"calls","name":"append","viaMember":true,"file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker","relation":"contains","targetId":"src/utils/model_utils.py#TokenUsageTracker.get_summary","file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker.get_summary","relation":"calls","name":"items","viaMember":true,"file":"src/utils/model_utils.py"},{"source":"src/utils/model_utils.py#TokenUsageTracker.get_summary","relation":"calls","name":"len","viaMember":false,"file":"src/utils/model_utils.py"}]}}}