from openai import OpenAI
import json
import time

# 初始化DeepSeek AI客户端
client = OpenAI(api_key="sk-ee3", base_url="https://api.deepseek.com")

def gen(prompt):
    """模拟APPL的gen函数，使用DeepSeek AI生成回复"""
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "You are a helpful assistant skilled in step-by-step reasoning."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=1024,
        temperature=0.7,
        stream=False
    )
    return response.choices[0].message.content

def cot_sc(question: str, num_samples: int = 5):
    # 设置系统提示和用户问题
    system_prompt = "You are a helpful assistant skilled in step-by-step reasoning."
    user_prompt = f"Question: {question}\nLet's approach this step-by-step:"
    
    # 生成多个CoT推理样本
    samples = [gen(user_prompt) for _ in range(num_samples)]
    
    # 从每个样本中提取最终答案
    final_answers = []
    for sample in samples:
        # 解析每个样本的最后一行作为最终答案
        final_answer = sample.split('\n')[-1].strip()
        if final_answer.startswith("Therefore, "):
            final_answer = final_answer[len("Therefore, "):]
        final_answers.append(final_answer)
    
    # 进行自洽性检查
    consistency_prompt = "Now, let's analyze the consistency of our reasoning:\n"
    for i, answer in enumerate(final_answers):
        consistency_prompt += f"Sample {i+1}: {answer}\n"
    consistency_prompt += "Based on the above samples, the most consistent answer is:"
    
    # 生成最终结论
    final_conclusion = gen(consistency_prompt)
    
    return final_conclusion

# 示例使用
question = "If a train travels 120 km in 2 hours, what is its average speed in km/h?"
start_time = time.time()
result = cot_sc(question)
end_time = time.time()

print(f"Final conclusion: {result}")
print(f"Total time taken: {end_time - start_time:.2f} seconds")