Update app.py

app.py

@@ -162,7 +162,16 @@ def messages():
         def generate():
             nonlocal total_output_tokens
 
-            start_event = create_event("message_start", {
+            def send_event(event_type, data):
+                event = create_event(event_type, data)
+                logger.info(f"Sending {event_type} event", extra={
+                    'event_type': event_type,
+                    'data': {'content': event}
+                })
+                yield event
+
+            # Send initial events
+            yield from send_event("message_start", {
                 "type": "message_start",
                 "message": {
                     "id": msg_id,
@@ -175,25 +184,8 @@ def messages():
                     "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
                 },
             })
-            logger.info("Sending message_start event", extra={
-                'event_type': 'message_start',
-                'data': {'content': start_event}
-            })
-            yield start_event
-
-            block_start_event = create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
-            logger.info("Sending content_block_start event", extra={
-                'event_type': 'content_block_start',
-                'data': {'content': block_start_event}
-            })
-            yield block_start_event
-
-            ping_event = create_event("ping", {"type": "ping"})
-            logger.info("Sending ping event", extra={
-                'event_type': 'ping',
-                'data': {'content': ping_event}
-            })
-            yield ping_event
+            yield from send_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
+            yield from send_event("ping", {"type": "ping"})
 
             def on_query_progress(data):
                 nonlocal total_output_tokens, response_text
@@ -204,23 +196,22 @@ def messages():
                         response_text.append(chunk)
                         chunk_tokens = calculate_tokens(chunk)
                         total_output_tokens += chunk_tokens
-                        delta_event = create_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-                        logger.info("Sending content_block_delta event", extra={
-                            'event_type': 'content_block_delta',
+                        logger.info("Received chunk", extra={
+                            'event_type': 'chunk_received',
                             'data': {
                                 'chunk': chunk,
                                 'tokens': chunk_tokens,
-                                'total_tokens': total_output_tokens,
-                                'content': delta_event
+                                'total_tokens': total_output_tokens
                             }
                         })
-                        yield delta_event
 
                 if data.get('final', False):
+                    logger.info("Final response received", extra={
+                        'event_type': 'response_complete',
+                        'data': {
+                            'total_tokens': total_output_tokens
+                        }
+                    })
                     response_event.set()
 
             def on_connect():
@@ -255,13 +246,17 @@ def messages():
                 
                 while not response_event.is_set():
                     sio.sleep(0.1)
+                    while response_text:
+                        chunk = response_text.pop(0)
+                        yield from send_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
                 
             except Exception as e:
-                logger.error(f"Error during socket connection: {str(e)}", extra={
-                    'event_type': 'connection_error',
-                    'data': {'error': str(e)}
-                })
-                yield create_event("content_block_delta", {
+                logger.error(f"Error during socket connection: {str(e)}", exc_info=True)
+                yield from send_event("content_block_delta", {
                     "type": "content_block_delta",
                     "index": 0,
                     "delta": {"type": "text_delta", "text": f"Error during socket connection: {str(e)}"},
@@ -270,38 +265,19 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
-            stop_event = create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-            logger.info("Sending content_block_stop event", extra={
-                'event_type': 'content_block_stop',
-                'data': {'content': stop_event}
-            })
-            yield stop_event
-
-            message_delta_event = create_event("message_delta", {
+            # Send final events
+            yield from send_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+            yield from send_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                 "usage": {"output_tokens": total_output_tokens},
             })
-            logger.info("Sending message_delta event", extra={
-                'event_type': 'message_delta',
-                'data': {'content': message_delta_event}
-            })
-            yield message_delta_event
-
-            message_stop_event = create_event("message_stop", {"type": "message_stop"})
-            logger.info("Sending message_stop event", extra={
-                'event_type': 'message_stop',
-                'data': {'content': message_stop_event}
-            })
-            yield message_stop_event
+            yield from send_event("message_stop", {"type": "message_stop"})
 
         return Response(generate(), content_type='text/event-stream')
 
     except Exception as e:
-        logger.error(f"Request error: {str(e)}", extra={
-            'event_type': 'request_error',
-            'data': {'error': str(e)}
-        })
+        logger.error(f"Request error: {str(e)}", exc_info=True)
         log_request(request.remote_addr, request.path, 400)
         return jsonify({"error": str(e)}), 400