Files
03Rag/rag/llm.py
T
heyong.fu a17c65c4bc feat: rag
2026-05-06 11:35:10 +08:00

56 lines
1.9 KiB
Python

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)