from openai import OpenAI
from datetime import datetime
# Initialize Stepfun client
STEPFUN_KEY = ""
client = OpenAI(
base_url="https://api.stepfun.ai/v1",
api_key=STEPFUN_KEY
)
# Data structures
class Chat:
id: str
user_id: str
class Message:
chat_id: str
role : str # system, user, assistant, tool
content: str # user input or model output
created_at: datetime
# Fetch messages from the database
# messages_from_db = orm.order_by("created_at","asc").first(5)
messages_from_db = [
{
"chat_id":"chat_1",
"role":"system",
"content":"You are the Stepfun assistant",
"created_at": "2024-01-01 10:01:00"
},
{
"chat_id":"chat_1",
"role":"user",
"content":"How is the weather today?",
"created_at": "2024-01-01 10:02:00"
},
{
"chat_id":"chat_1",
"role":"assistant",
"content":"Sorry, I can’t answer weather questions.",
"created_at": "2024-01-01 10:03:00"
},
{
"chat_id":"chat_1",
"role":"user",
"content":"Is Beijing a good place to travel?",
"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]
# Call the completion API
stream = client.chat.completions.create(
model="step-3.7-flash",
messages=messages_for_chat,
stream=True,
)
# Render streaming output
for chunk in stream:
if chunk.choices[0].delta.content is not None:
print(chunk.choices[0].delta.content, end="")