mirror of
https://github.com/cpacker/MemGPT.git
synced 2025-06-03 04:30:22 +00:00
61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
import asyncio
|
|
import json
|
|
import traceback
|
|
from typing import AsyncGenerator, Generator, Union
|
|
|
|
from memgpt.constants import JSON_ENSURE_ASCII
|
|
|
|
SSE_FINISH_MSG = "[DONE]" # mimic openai
|
|
SSE_ARTIFICIAL_DELAY = 0.1
|
|
|
|
|
|
def sse_formatter(data: Union[dict, str]) -> str:
|
|
"""Prefix with 'data: ', and always include double newlines"""
|
|
assert type(data) in [dict, str], f"Expected type dict or str, got type {type(data)}"
|
|
data_str = json.dumps(data, ensure_ascii=JSON_ENSURE_ASCII) if isinstance(data, dict) else data
|
|
return f"data: {data_str}\n\n"
|
|
|
|
|
|
async def sse_generator(generator: Generator[dict, None, None]) -> Generator[str, None, None]:
|
|
"""Generator that returns 'data: dict' formatted items, e.g.:
|
|
|
|
data: {"id":"chatcmpl-9E0PdSZ2IBzAGlQ3SEWHJ5YwzucSP","object":"chat.completion.chunk","created":1713125205,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"}"}}]},"logprobs":null,"finish_reason":null}]}
|
|
|
|
data: {"id":"chatcmpl-9E0PdSZ2IBzAGlQ3SEWHJ5YwzucSP","object":"chat.completion.chunk","created":1713125205,"model":"gpt-4-0613","system_fingerprint":null,"choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}]}
|
|
|
|
data: [DONE]
|
|
|
|
"""
|
|
try:
|
|
for msg in generator:
|
|
yield sse_formatter(msg)
|
|
if SSE_ARTIFICIAL_DELAY:
|
|
await asyncio.sleep(SSE_ARTIFICIAL_DELAY) # Sleep to prevent a tight loop, adjust time as needed
|
|
except Exception as e:
|
|
yield sse_formatter({"error": f"{str(e)}"})
|
|
yield sse_formatter(SSE_FINISH_MSG) # Signal that the stream is complete
|
|
|
|
|
|
async def sse_async_generator(generator: AsyncGenerator, finish_message=True):
|
|
"""
|
|
Wraps a generator for use in Server-Sent Events (SSE), handling errors and ensuring a completion message.
|
|
|
|
Args:
|
|
- generator: An asynchronous generator yielding data chunks.
|
|
|
|
Yields:
|
|
- Formatted Server-Sent Event strings.
|
|
"""
|
|
try:
|
|
async for chunk in generator:
|
|
# yield f"data: {json.dumps(chunk)}\n\n"
|
|
yield sse_formatter(chunk)
|
|
except Exception as e:
|
|
print("stream decoder hit error:", e)
|
|
print(traceback.print_stack())
|
|
yield sse_formatter({"error": "stream decoder encountered an error"})
|
|
finally:
|
|
# yield "data: [DONE]\n\n"
|
|
if finish_message:
|
|
yield sse_formatter(SSE_FINISH_MSG) # Signal that the stream is complete
|