| # | |
| # SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org> | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| import json # Import the json module to handle JSON encoding and decoding of data | |
| import time # Import the time module to work with timestamps | |
| import uuid # Import the uuid module to generate unique identifiers | |
| from typing import AsyncGenerator # Import AsyncGenerator for typing asynchronous generator functions | |
| from src.cores.sessions import session_store # Import the session_store object to manage user sessions | |
| async def event_generator(user_input: str, model: str, session_id: str) -> AsyncGenerator[str, None]: | |
| """ | |
| Asynchronous generator function that streams AI-generated responses incrementally as Server-Sent Events (SSE). | |
| Parameters: | |
| - user_input: The input text provided by the user to the AI model. | |
| - model: The identifier of the AI model to be used for generating responses. | |
| - session_id: A unique string representing the current user session. | |
| Yields: | |
| - JSON-formatted strings representing incremental chunks of the AI response, | |
| formatted as Server-Sent Events for real-time streaming to the client. | |
| """ | |
| # Retrieve the last update time and session data from the session store using the session ID | |
| last_update, session_data = session_store.get(session_id, (0, None)) | |
| # If no session data is found for the given session ID, yield an error message and stop | |
| if session_data is None: | |
| yield f"data: {json.dumps({'error': 'Session not found'})}\n\n" | |
| return | |
| # Extract the AI client object from the session data | |
| client = session_data["client"] | |
| # If the client is missing in the session data, yield an error message and stop | |
| if client is None: | |
| yield f"data: {json.dumps({'error': 'AI client not available'})}\n\n" | |
| return | |
| try: | |
| # Submit the user's input text to the AI model via the client's submit method | |
| # The 'multi' parameter wraps the text, and 'api_name' specifies the API endpoint | |
| jarvis_response = client.submit(multi={"text": user_input}, api_name="/api") | |
| except Exception as e: | |
| # If submission to the AI fails, yield an error message with the exception details and stop | |
| yield f"data: {json.dumps({'error': f'Failed to submit to AI: {str(e)}'})}\n\n" | |
| return | |
| buffer = "" # Initialize an empty string buffer to accumulate the full AI response progressively | |
| try: | |
| # Iterate over the partial responses received from the AI client submission | |
| for partial in jarvis_response: | |
| # Extract the current partial text chunk from the nested response structure | |
| text = partial[0][0][1] | |
| # Determine the new delta text by comparing with the buffer | |
| # If the new text starts with the buffer, delta is the new appended text; otherwise, delta is the entire text | |
| if text.startswith(buffer): | |
| delta = text[len(buffer):] | |
| else: | |
| delta = text | |
| buffer = text # Update the buffer with the latest full text from the AI | |
| # Skip yielding if the delta is empty (no new text) | |
| if delta == "": | |
| continue | |
| # Construct a chunk dictionary following OpenAI's streaming response format | |
| chunk = { | |
| "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", # Unique chunk ID with a UUID suffix | |
| "object": "chat.completion.chunk", # Object type indicating a chunk of chat completion | |
| "created": int(time.time()), # Timestamp of chunk creation in Unix time | |
| "model": model, # The AI model used for generating this chunk | |
| "choices": [ | |
| { | |
| "index": 0, # Index of the choice in the response (usually 0 for single response) | |
| "delta": {"content": delta}, # The incremental new text content in this chunk | |
| "finish_reason": None # No finish reason yet, stream is ongoing | |
| } | |
| ] | |
| } | |
| # Yield the chunk as a Server-Sent Event formatted string with 'data:' prefix and double newline suffix | |
| yield f"data: {json.dumps(chunk)}\n\n" | |
| # After all chunks have been streamed, append the full input-response pair to the session history | |
| session_data["history"].append({"input": user_input, "response": buffer}) | |
| # Update the session store with the new last access time and updated session data | |
| session_store[session_id] = (time.time(), session_data) | |
| # Prepare a final chunk indicating the end of the stream with finish_reason set to 'stop' | |
| done_chunk = { | |
| "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", # Unique ID for the final chunk | |
| "object": "chat.completion.chunk", # Object type for consistency | |
| "created": int(time.time()), # Timestamp of completion | |
| "model": model, # Model identifier | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "delta": {}, # Empty delta indicating no new content | |
| "finish_reason": "stop" # Signal that the stream has finished | |
| } | |
| ] | |
| } | |
| # Yield the final completion chunk to signal the client that streaming is done | |
| yield f"data: {json.dumps(done_chunk)}\n\n" | |
| except Exception as e: | |
| # If any error occurs during streaming, yield an error chunk with the exception message | |
| error_chunk = { | |
| "error": {"message": f"Streaming error: {str(e)}"} | |
| } | |
| yield f"data: {json.dumps(error_chunk)}\n\n" | |