import requests

# Ollama API 地址
OLLAMA_URL = "http://localhost:11434/api/generate"

def extract_triples(text):
    """
    调用 DeepSeek 模型提取三元组信息
    :param text: 输入的文本
    :return: 提取的三元组列表
    """
    # 构造提示词，明确要求提取三元组
    prompt = f"""
    请从以下文本中提取知识图谱的三元组信息，格式为 (实体1, 关系, 实体2)。
    文本：{text}
    请直接返回三元组，不要添加额外解释。
    """

    # 请求参数
    payload = {
        "model": "deepseek-r1",  # 使用 DeepSeek 模型
        "prompt": prompt,
        "stream": False,  # 一次性返回结果
        "options": {
            "temperature": 0.3,  # 降低随机性，确保结果稳定
            "num_predict": 200   # 限制生成长度
        }
    }

    # 打印调试信息
    print("请求 URL:", OLLAMA_URL)
    print("请求参数:", payload)

    # 发送 POST 请求
    response = requests.post(OLLAMA_URL, json=payload)
    if response.status_code != 200:
        raise Exception(f"请求失败，状态码: {response.status_code}, 响应内容: {response.text}")

    # 解析响应
    result = response.json()
    output = result.get("response", "").strip()

    # 提取三元组
    triples = []
    for line in output.split("\n"):
        if line.startswith("(") and line.endswith(")"):
            triples.append(line.strip())
    
    return triples

# 示例文本
text = """
苹果公司由史蒂夫·乔布斯于1976年创立，总部位于加利福尼亚州库比蒂诺。
乔布斯是苹果公司的首席执行官，他推出了iPhone等革命性产品。
"""

# 提取三元组
try:
    triples = extract_triples(text)
    print("提取的三元组：")
    for triple in triples:
        print(triple)
except Exception as e:
    print("发生错误:", e)