script_query.py

import json
import argparse
import os
import asyncio
# from dataclasses import is_dataclass, asdict
from datetime import datetime
import gc
import pathlib


from PRoH import PRoH, QueryParam


os.environ["OPENAI_API_KEY"] = open("openai_api_key.txt").read().strip()


# CLI
parser = argparse.ArgumentParser()
parser.add_argument('--data_source', default='example')
parser.add_argument('--part', default='')
# parser.add_argument('--qp', default='', help='Path to JSON with QueryParam overrides')
parser.add_argument('--ts', default='', help='Run timestamp decided by input')
args = parser.parse_args()

data_source = args.data_source
part = args.part.strip()
part_tag = part if part else "orig"

MAX_CONCURRENCY = 32  
MAX_RETRYS = 3      


WORKING_DIR = os.path.abspath(f"./data/work/{data_source}")
if not os.path.exists(WORKING_DIR):
    exit(f"Error: working dir not exists: {WORKING_DIR}")

QUESTION_DIR = f"questions/{data_source}/questions.json"
if part:
    QUESTION_DIR = f"questions/{data_source}/questions_{part}.json"

RUN_ROOT = f"results/{data_source}" if not part else f"results/{data_source}_{part}"

# RUN_TS / RUN_DIR
RUN_TS = args.ts.strip() or datetime.now().strftime("%Y%m%d-%H%M%S")
RUN_DIR = os.path.join(RUN_ROOT, RUN_TS)
os.makedirs(RUN_DIR, exist_ok=True)


# File paths inside the run dir
LOG_FILE = os.path.join(RUN_DIR, "test.log")
ANSWER_SAVE_PATH = os.path.join(RUN_DIR, "answers.jsonl")
FINAL_JSON_PATH = os.path.join(RUN_DIR, "generated_answer.json")

# # Registry / params snapshot
# REGISTRY_PATH = os.path.join(RUN_ROOT, "run_registry.jsonl")
# PARAMS_SNAPSHOT_PATH = os.path.join(RUN_DIR, "qp_used.json")


# RAG instance (single-process async like original code 1) -
rag = PRoH(
    working_dir=WORKING_DIR,
    log_level="INFO",
    log_file=LOG_FILE,
    llm_model_max_async=MAX_CONCURRENCY,
    max_concurrency=MAX_CONCURRENCY 
)



async def query_with_semaphore(sem: asyncio.Semaphore, q_data: dict):
    async with sem:

        qd = dict(q_data) 
        # if "answer" in qd:
        #     qd["golden_answers"] = [qd.pop("answer")]
        q = qd["question"]

        result = {
            "gen_answer": "",
            "generation": "",
            "retrieved": [],
            "dag": "",
            "timer": None
        }

        last_err = None
        for _ in range(MAX_RETRYS):
            try:
                response = await rag.aquery_reasoning(q)
                if response:
                    result.update(response)
                    last_err = None
                    break
            except Exception as e:
                last_err = repr(e)

        if last_err is not None:
            result["error"] = last_err

        # prune heavy context in provided reasoning paths
        rp = qd.get("reasoning_paths")
        if isinstance(rp, list):
            for path in rp:
                if isinstance(path, dict):
                    path.pop("context", None)

        qd.update(result)

        # console echo
        print(f"Question: {q}")
        print(f"Answer: {qd.get('gen_answer', '')}")
        print(f"Timer: {qd.get('timer', None)}")

        return qd



async def amain():
    # Load questions
    with open(QUESTION_DIR, "r", encoding="utf-8") as f:
        data = json.load(f)

    # fresh stream file
    open(ANSWER_SAVE_PATH, "w", encoding="utf-8").close()

    sem = asyncio.Semaphore(MAX_CONCURRENCY)
    write_lock = asyncio.Lock()

    tasks = [asyncio.create_task(query_with_semaphore(sem, d)) for d in data]

    results_all = []
    for fut in asyncio.as_completed(tasks):
        updated = await fut
        results_all.append(updated)
        # incremental append
        async with write_lock:
            with open(ANSWER_SAVE_PATH, "a", encoding="utf-8") as f:
                f.write(json.dumps(updated, ensure_ascii=False) + "\n")

    # final consolidated snapshot
    with open(FINAL_JSON_PATH, "w", encoding="utf-8") as f:
        json.dump(results_all, f, indent=4, ensure_ascii=False)

    print(f"Run dir -> {RUN_DIR}")
    print(f"Log     -> {LOG_FILE}")
    print(f"Out     -> {ANSWER_SAVE_PATH}")
    print(f"Final   -> {FINAL_JSON_PATH}")
    # print(f"Registry-> {REGISTRY_PATH}")


def main():
    asyncio.run(amain())


if __name__ == "__main__":

    main()
    pathlib.Path(RUN_DIR, ".READY").touch()