This commit is contained in:
Timothy J. Baek 2024-05-21 18:49:34 -07:00
parent 39da1c982a
commit be06b61e83
2 changed files with 4 additions and 5 deletions

View File

@ -135,13 +135,13 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
print(res)
if isinstance(res, str):
message = stream_message_template(res)
message = stream_message_template(form_data.model, res)
yield f"data: {json.dumps(message)}\n\n"
elif isinstance(res, Generator):
for message in res:
print(message)
message = stream_message_template(message)
message = stream_message_template(form_data.model, message)
yield f"data: {json.dumps(message)}\n\n"
finish_message = {

View File

@ -3,15 +3,14 @@ import time
from typing import List
from schemas import OpenAIChatMessage
from config import MODEL_ID
def stream_message_template(message: str):
def stream_message_template(model: str, message: str):
return {
"id": f"rag-{str(uuid.uuid4())}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": MODEL_ID,
"model": model,
"choices": [
{
"index": 0,