import os

# Install SDK:  pip install 'volcengine-python-sdk[ark]'
# from volcenginesdkarkruntime import Ark

# client = Ark(
#     # The base URL for model invocation
#     base_url="https://ark.cn-beijing.volces.com/api/v3/chat/completions",
#     api_key=os.getenv("ARK_API_KEY", "79b39c58-56db-4d8a-a8f8-84b95fca08db"),
# )

# completion = client.chat.completions.create(
#     # Replace with Model ID
#     model="doubao-seed-1-6-lite-251015",
#     messages=[
#         {
#             "role": "system",
#             "content": "请将下面内容进行结构化处理：火山方舟是火山引擎推出的大模型服务平台，提供模型训练、推理、评测、精调等全方位功能与服务，并重点支撑大模型生态。 火山方舟通过稳定可靠的安全互信方案，保障模型提供方的模型安全与模型使用者的信息安全，加速大模型能力渗透到千行百业，助力模型提供方和使用者实现商业新增长。",
#         },
#     ],
# )

# print(completion.choices[0].message.content)

# 使用豆包来向量化文本

import requests

VOLC_EMBEDDINGS_API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
VOLC_API_KEY = "79b39c58-56db-4d8a-a8f8-84b95fca08db"


def get_doubao_llm(prompt):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {VOLC_API_KEY}",
    }
    params = {
        "model": "doubao-seed-1-6-lite-251015",
        "messages": [
            {"role": "system", "content": f"{prompt}"},
        ],
    }
    response = requests.post(VOLC_EMBEDDINGS_API_URL, json=params, headers=headers)
    if response.status_code == 200:
        data = response.json()
        print(data)
        message = data["choices"][0]["message"]["content"]
        return message
    else:
        raise Exception(f"Embedding API error:{response.text}")


answer = get_doubao_llm("红楼梦的作者是谁")
print(answer)