62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
豆包（火山方舟）API 封装
|
|
使用 OpenAI 兼容接口
|
|
"""
|
|
import os
|
|
import json
|
|
import requests
|
|
|
|
|
|
class DoubaoService:
|
|
"""豆包流式对话服务"""
|
|
|
|
def __init__(self):
|
|
self.api_key = os.getenv("DOUBAO_API_KEY", "")
|
|
self.model = os.getenv("DOUBAO_MODEL", "doubao-1-5-pro-32k-250115")
|
|
self.base_url = os.getenv(
|
|
"DOUBAO_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3"
|
|
)
|
|
|
|
def stream_chat(self, messages):
|
|
"""
|
|
流式对话，逐块 yield 文本内容
|
|
:param messages: [{"role": "user"|"assistant"|"system", "content": str}]
|
|
"""
|
|
if not self.api_key:
|
|
raise RuntimeError("未配置 DOUBAO_API_KEY，请在 .env 中设置")
|
|
|
|
url = f"{self.base_url}/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
body = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
}
|
|
|
|
with requests.post(
|
|
url, headers=headers, json=body, stream=True, timeout=120
|
|
) as resp:
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(
|
|
f"豆包接口错误 {resp.status_code}: {resp.text}"
|
|
)
|
|
for line in resp.iter_lines(decode_unicode=True):
|
|
if not line:
|
|
continue
|
|
if line.startswith("data:"):
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
break
|
|
try:
|
|
obj = json.loads(data)
|
|
delta = obj.get("choices", [{}])[0].get("delta", {})
|
|
content = delta.get("content")
|
|
if content:
|
|
yield content
|
|
except json.JSONDecodeError:
|
|
continue
|