from core.settings import get_settings
from core.openai_used_attrs import rounds
from openai import OpenAI
import base64, json
# import pubchempy as pcp
from loguru import logger
import os
import glob

settings = get_settings()

client = OpenAI(
    api_key=settings.api_key,
    base_url=settings.base_url,
)

# ============== Read File ==============

# new_datasets = ["crossed_barrel", "dye_lasers", "lnp3", "perovskites", "redoxmers"]

dataset = "crossed_barrel"
json_path = f"/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/train_regression/data4regression/{dataset}/options.json"

dataset_name = dataset
save_path = f"/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/Rag-Cluster/sci_files/{dataset}"

os.makedirs(save_path, exist_ok=True)


templates = {
    "role": "user",
    "content": []
}

# ============== Read Targets ==============

with open(f"sci_files/prompts/{dataset}.txt", "r") as f:
    queries = f.read()

templates["content"].append({
    "type": "input_text",
    "text": queries,
})

logger.info("Send Request To API")

response = client.responses.create(
    model=settings.model_name,
    tools=[{"type": "web_search_preview"}] if "gpt" in settings.model_name else None,
    input=[templates],
)
try:
    reply = response.output[1].content[0].text
except:
    import pdb;pdb.set_trace()

print(reply)

# 保存结果
with open(os.path.join(save_path, f"openai_{dataset}_rank.md"), "w") as f:
    f.write(reply)