from openai import OpenAI
from datetime import datetime
 
 
 
# 初始化 阶跃星辰 Client
STEPFUN_KEY = ""
 
client = OpenAI(
  base_url="https://api.stepfun.com/v1",
  api_key=STEPFUN_KEY
)
 
 
# 定义数据结构
 
class Chat:
  id: str
  user_id: str
 
class Message:
  chat_id: str
  role : str # 可选值 system, user,assistant,tool
  content: str # 用户输入的信息 / 大模型返回的信息
  created_at: datetime
 
# 从数据库中提取数据
 
# messages_from_db = orm.order_by ("created_at","asc").first (5)
 
messages_from_db = [
  {
    "chat_id":"chat_1",
    "role":"system",
    "content":"你是阶跃星辰大模型助手",
    "created_at": "2024-01-01 10:01:00"
  },
  {
    "chat_id":"chat_1",
    "role":"user",
    "content":"今天天气怎么样？",
    "created_at": "2024-01-01 10:02:00"
  },
  {
    "chat_id":"chat_1",
    "role":"assistant",
    "content":"对不起，我不能回答你关于天气的问题。",
    "created_at": "2024-01-01 10:03:00"
  },
  {
    "chat_id":"chat_1",
    "role":"user",
    "content":"那你告诉我北京适合去旅游么？",
    "created_at": "2024-01-01 10:04:00"
  }
]
 
def clean_msg(msg):
  del msg["chat_id"]
  del msg["created_at"]
  return msg
 
messages_for_chat = [clean_msg (item) for item in messages_from_db]
 
 
 
# 调用补全接口进行补全
 
stream = client.chat.completions.create(
  model="step-1-8k",
  messages=messages_for_chat,
  stream=True,
)
 
# 对流式返回的内容进行打印 / 渲染输出
 
for chunk in stream:
  if chunk.choices[0].delta.content is not None:
    print(chunk.choices[0].delta.content, end="")