from dotenv import load_dotenv

import os

# 加载 .env 文件

load_dotenv()

# 设置代理（如果需要）

def set_proxy():

    http_proxy = os.getenv('HTTP_PROXY')

    https_proxy = os.getenv('HTTPS_PROXY')

    if http_proxy:

        os.environ['HTTP_PROXY'] = http_proxy

    if https_proxy:

        os.environ['HTTPS_PROXY'] = https_proxy

from pydantic import BaseModel

# 定义生物标志物注释的 Schema

class BiomarkerAnnotation(BaseModel):

    Biomarker_Std: str  # 标准化后的生物标志物名称

    GPTAnno: str  # 标准化后的名称及其生物功能注释

    Symbol: str  # 相关基因的 SYMBOL 或 "/"

from openai import OpenAI

# 初始化 OpenAI 客户端

def init_openai_client(api_key: str):

    return OpenAI(api_key=api_key)

# 注释生物标志物函数

def annotate_biomarker(biomarker_name: str, model: str, api_key: str) -> BiomarkerAnnotation:

    prompt = f"""

    I have a biomarker related to aging.

    1. Please standardize the name of this biomarker if it's unclear, and provide a concise annotation describing its function in biological processes, especially related to aging.

    2. If this biomarker represents a specific gene, provide the official SYMBOL for the gene. If it's a chemical substance or small molecule, provide the directly related gene SYMBOL(s). If no relevant gene is associated, return "/".

    3. If there are multiple SYMBOLs, separate them with a comma followed by a space (", ").

    The output should be in JSON format with the following fields:

    
    {{
      "Biomarker_Std": "the standardized name of this biomarker",
      "GPTAnno": "standardized biomarker name and its biological function",
      "Symbol": "associated gene SYMBOL(s) or '/' if none"
    }}
    

    Here is the biomarker: '{biomarker_name}'

    """

    # 使用 OpenAI 客户端请求

    client = init_openai_client(api_key=api_key)

    completion = client.beta.chat.completions.parse(

        model=model,

        messages=[

            {"role": "system", "content": "You are a knowledgeable assistant specialized in aging-related biomarkers."},

            {"role": "user", "content": prompt},

        ],

        response_format=BiomarkerAnnotation,  # 使用自定义 Schema

    )

    # 解析结果

    message = completion.choices[0].message

    if message.parsed:

        return message.parsed

    else:

        return message.refusal

import pandas as pd

import time

# 批量处理生物标志物并保存结果

def process_biomarkers(df: pd.DataFrame, model: str, api_key: str, output_file: str, batch_size: int = 10):

    df_ans = pd.DataFrame(columns=['PMID', 'Biomarker', 'Biomarker_Std', 'GPTAnno', 'Symbol'])

    for i in range(min(batch_size, len(df))):

        if i > 0 and i % 50 == 0:

            time.sleep(5)

        p = df.iloc[i]['PMID']

        biomarker = df.iloc[i]['Biomarker']

        print(f'正在处理第 {i + 1}/{len(df)} 个生物标志物，PMID 为：{p}')

        response = None

        for attempt in range(5):

            try:

                response = annotate_biomarker(biomarker, model=model, api_key=api_key)

                break

            except Exception as e:

                print(f"Error on attempt {attempt + 1} for PMID {p}: {e}")

                time.sleep(2)

                continue

        if response is None:

            print(f"Failed to process biomarker '{biomarker}' after 5 attempts.")

            continue

        # 保存结果

        dfrow = {

            'PMID': p,

            'Biomarker': biomarker,  

            'Biomarker_Std': response.Biomarker_Std, 

            'GPTAnno': response.GPTAnno, 

            'Symbol': response.Symbol

        }

        df_ans = pd.concat([df_ans, pd.DataFrame([dfrow])], ignore_index=True)

        print(f"Processed: {response}\n")

        if (i + 1) % 500 == 0:

            df_ans.to_csv(f'intermediate_result_{i+1}.csv', index=False, encoding='utf-8')

        time.sleep(1)

    df_ans.to_csv(output_file, index=False, encoding='utf-8')

    print(f"All biomarkers processed and saved to {output_file}")