poojagoyal3072/DABstep

1import json2from tqdm import tqdm3import logging4import threading5from smolagents import CodeAgent6from custom_agent import CustomCodeAgent7from custom_litellm import LiteLLMModelWithBackOff8from huggingface_hub import hf_hub_download9from constants import REPO_ID, ADDITIONAL_AUTHORIZED_IMPORTS10from pathlib import Path11from prompts import reasoning_llm_system_prompt, chat_llm_system_prompt12 13append_answer_lock = threading.Lock()14append_console_output_lock = threading.Lock()15 16class TqdmLoggingHandler(logging.Handler):17    def emit(self, record):18        tqdm.write(self.format(record))19 20def read_only_open(*a, **kw):21    if (len(a) > 1 and isinstance(a[1], str) and a[1] != 'r') or kw.get('mode', 'r') != 'r':22        raise Exception("Only mode='r' allowed for the function open")23    return open(*a, **kw)24 25def download_context(base_dir: str) -> str:26    ctx_files = [27        "data/context/acquirer_countries.csv",28        "data/context/payments.csv",29        "data/context/merchant_category_codes.csv",30        "data/context/fees.json",31        "data/context/merchant_data.json",32        "data/context/manual.md",33        "data/context/payments-readme.md"34    ]35    for f in ctx_files:36        hf_hub_download(REPO_ID, repo_type="dataset", filename=f, local_dir=base_dir, force_download=True)37 38    root_dir = Path(__file__).resolve().parent.parent39    full_path = Path(base_dir) / Path(ctx_files[0]).parent40    relative_path = full_path.relative_to(root_dir)41    return str(relative_path)42 43def is_reasoning_llm(model_id: str) -> bool:44    reasoning_llm_list = [45        "openai/o1",46        "openai/o3",47        "openai/o3-mini",48        "deepseek/deepseek-reasoner"49    ]50    return model_id in reasoning_llm_list51 52def get_tasks_to_run(data, total: int, base_filename: Path, tasks_ids: list[int]):53    import json54    f = base_filename.parent / f"{base_filename.stem}_answers.jsonl"55    done = set()56    if f.exists():57        with open(f, encoding="utf-8") as fh:58            done = {json.loads(line)["task_id"] for line in fh if line.strip()}59 60    tasks = []61    for i in range(total):62        task_id = int(data[i]["task_id"])63        if task_id not in done:64            if tasks_ids is not None:65                if task_id in tasks_ids:66                    tasks.append(data[i])67            else:68                tasks.append(data[i])69    return tasks70 71 72def append_answer(entry: dict, jsonl_file: Path) -> None:73    jsonl_file.parent.mkdir(parents=True, exist_ok=True)74    with append_answer_lock, open(jsonl_file, "a", encoding="utf-8") as fp:75        fp.write(json.dumps(entry) + "\n")76 77 78def append_console_output(captured_text: str, txt_file: Path) -> None:79    txt_file.parent.mkdir(parents=True, exist_ok=True)80    with append_console_output_lock, open(txt_file, "a", encoding="utf-8") as fp:81        fp.write(captured_text + "\n")82 83def create_code_agent_with_reasoning_llm(model_id: str, api_base=None, api_key=None, max_steps=10, ctx_path=None):84    agent = CustomCodeAgent(85        system_prompt=reasoning_llm_system_prompt,86        tools=[],87        model=LiteLLMModelWithBackOff(88            model_id=model_id, api_base=api_base, api_key=api_key, max_tokens=None, max_completion_tokens=3000),89        additional_authorized_imports=ADDITIONAL_AUTHORIZED_IMPORTS,90        max_steps=max_steps,91        verbosity_level=3,92    )93    agent.python_executor.static_tools.update({"open": read_only_open})94 95    agent.system_prompt = agent.system_prompt.format(ctx_path=ctx_path)96    return agent97 98def create_code_agent_with_chat_llm(model_id: str, api_base=None, api_key=None, max_steps=10):99    agent = CodeAgent(100        system_prompt=chat_llm_system_prompt,101        tools=[],102        model=LiteLLMModelWithBackOff(model_id=model_id, api_base=api_base, api_key=api_key, max_tokens=3000),103        additional_authorized_imports=ADDITIONAL_AUTHORIZED_IMPORTS,104        max_steps=max_steps,105        verbosity_level=3,106    )107 108    agent.python_executor.static_tools.update({"open": read_only_open})109    return agent110