2024-08-29 20:17:17 +08:00
|
|
|
import json
|
|
|
|
from collections.abc import Generator
|
|
|
|
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
|
|
|
|
|
|
class BaseBackwardsInvocation:
|
|
|
|
@classmethod
|
2024-08-29 20:50:36 +08:00
|
|
|
def convert_to_event_stream(cls, response: Generator[BaseModel | dict | str, None, None] | BaseModel | dict):
|
2024-08-29 20:17:17 +08:00
|
|
|
if isinstance(response, Generator):
|
|
|
|
for chunk in response:
|
|
|
|
if isinstance(chunk, BaseModel):
|
|
|
|
yield chunk.model_dump_json().encode() + b'\n\n'
|
2024-08-29 21:14:23 +08:00
|
|
|
elif isinstance(chunk, str):
|
2024-08-29 20:50:36 +08:00
|
|
|
yield f"event: {chunk}\n\n".encode()
|
2024-08-29 20:17:17 +08:00
|
|
|
else:
|
|
|
|
yield json.dumps(chunk).encode() + b'\n\n'
|
|
|
|
else:
|
|
|
|
if isinstance(response, BaseModel):
|
|
|
|
yield response.model_dump_json().encode() + b'\n\n'
|
|
|
|
else:
|
2024-08-29 20:50:36 +08:00
|
|
|
yield json.dumps(response).encode() + b'\n\n'
|