diff --git a/chatgpt/stream_v1.py b/chatgpt/stream_v1.py index 43dcb85..ad45b47 100644 --- a/chatgpt/stream_v1.py +++ b/chatgpt/stream_v1.py @@ -11,45 +11,56 @@ async def transform_delta_stream(input_stream) -> AsyncGenerator[str, None]: current_o = None current_p = None - async for line in input_stream: - line = line.decode("utf-8").strip() - if not line: - continue - - if line.startswith("event: "): - next_is_delta = True - continue - - if line.startswith("data: ") and next_is_delta: - data = line[6:] - - if data == "[DONE]": - yield line + try: + async for line in input_stream: + line = line.decode("utf-8").strip() + if not line: continue - try: - json_data = json.loads(data) + if line.startswith("event: "): + next_is_delta = True + continue - if 'c' in json_data: - current_c = json_data['c'] - if 'v' in json_data and isinstance(json_data['v'], dict): - current_documents[current_c] = json_data['v'] - yield f'data: {json.dumps(current_documents[current_c])}' + if line.startswith("data: ") and next_is_delta: + data = line[6:] + + if data == "[DONE]": + yield line continue - if 'v' in json_data: - if 'p' in json_data: - current_p = json_data['p'] - if 'o' in json_data: - current_o = json_data['o'] - if current_c is not None and current_c in current_documents: - apply_patch(current_documents, current_p, current_o, json_data['v'], current_c) + try: + json_data = json.loads(data) - yield f'data: {json.dumps(current_documents[current_c])}' - except: + if 'c' in json_data: + current_c = json_data['c'] + if 'v' in json_data and isinstance(json_data['v'], dict): + current_documents[current_c] = json_data['v'] + yield f'data: {json.dumps(current_documents[current_c])}' + continue + + if 'v' in json_data: + if 'p' in json_data: + current_p = json_data['p'] + if 'o' in json_data: + current_o = json_data['o'] + if current_c is not None and current_c in current_documents: + apply_patch(current_documents, current_p, current_o, json_data['v'], current_c) + + yield f'data: {json.dumps(current_documents[current_c])}' + except Exception as e: + logging.error(f"Error processing JSON data: {e}") + yield line + else: yield line - else: - yield line + except GeneratorExit: + raise + finally: + if hasattr(input_stream, 'aclose') and callable(input_stream.aclose): + try: + await input_stream.aclose() + except Exception as e: + logging.error(f"Error closing input_stream: {e}") + def apply_patch(document, p, o, v, c): if o == 'patch': diff --git a/version.txt b/version.txt index b18246b..e9d84b6 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -1.9.0-beta4 \ No newline at end of file +1.9.0-beta5 \ No newline at end of file