import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
import json


# 默认参数配置
default_params_config = {
    'temperature': 0.8,
    'max_tokens': 15000,
    'top_p': 0.9,
    'top_k': -1,
    'presence_penalty': 1.0,
    'repetition_penalty': 1.05,
    'n': 1,
    'stop': ['</python>', '</search>', '</answer>'],
    'include_stop_str_in_output': True,
}

def last_boxed_only_string(string):
    try:
        idx = string.rfind("\\boxed")
    except:
        import pdb; pdb.set_trace()
    if idx < 0:
        idx = string.rfind("\\fbox")
        if idx < 0:
            return None

    i = idx
    right_brace_idx = None
    num_left_braces_open = 0
    while i < len(string):
        if string[i] == "{":
            num_left_braces_open += 1
        if string[i] == "}":
            num_left_braces_open -= 1
            if num_left_braces_open == 0:
                right_brace_idx = i
                break
        i += 1

    if right_brace_idx is None:
        retval = string[idx:]
    else:
        retval = string[idx:right_brace_idx + 1]

    return retval

def seed_api(query):
    """调用API获取答案"""
    try:
        response = requests.post("http://0.0.0.0:8008/v1/chat/completions", 
                               json={"question": query, "default_params_config": default_params_config}, timeout=600)
        response.raise_for_status()
        data = response.json()
        # 完整输出
        ans = data["answer"]
        # 只输出<answer>中的\\boxed{}中的内容
        final_answer = data["final_answer"]

        answer_text = ans.rfind("</think>")
        if answer_text != -1:
            answer_text = ans[answer_text+len("</think>"):]
            answer_text = last_boxed_only_string(answer_text)
            if answer_text:
                final_answer = answer_text[7:-1]
            else:
                final_answer = ""

        return ans, final_answer
    except Exception as e:
        print(f"API调用失败: {e}")
        return f"Error: {str(e)}"

def make_concurrent_requests(queries, max_workers=5):
    """多线程并发请求"""
    results = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_query = {executor.submit(seed_api, query): query for query in queries}
        
        for future in as_completed(future_to_query):
            query = future_to_query[future]
            try:
                result, final_answer = future.result()
                results.append((query, result, final_answer))
            except Exception as e:
                print(f"请求'{query}'时发生异常: {e}")
                results.append((query, f"Error: {str(e)}"))
    return results

if __name__ == "__main__":
    question = "The molecule has 36 carbon atoms, 1 tellurium atom, and 2 boron atoms.\n"
    pre_instruction = """You are working as an assistant of a chemist user. Please follow the instruction of the chemist and generate a molecule that satisfies the requirements of the chemist user. You could think step by step, but your final response should be a SMILES string. 

Final Result Format:
- Place the final calculation or derived answer within the symbol \\boxed{ SMILES string }.

Questions:
{question}
"""
    ans, final_answer = seed_api(pre_instruction.replace("{question}", question))
    print(ans)
    print(f"*"*100)
    print(final_answer)