From fc9541aa82e94b449ec5f76d68540bbb7637ac9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?th=E1=BB=8Bnh?= Date: Fri, 19 Sep 2025 13:00:30 +0700 Subject: [PATCH] Fix the pusher loops on flushing the last pieces of audio and transcripts after the functions ended --- backend/routers/transcribe.py | 150 ++++++++++++++++++++-------------- 1 file changed, 88 insertions(+), 62 deletions(-) diff --git a/backend/routers/transcribe.py b/backend/routers/transcribe.py index 0c44e57ab81..41282b40447 100644 --- a/backend/routers/transcribe.py +++ b/backend/routers/transcribe.py @@ -2,17 +2,19 @@ import json import os import struct +import time import uuid -from datetime import datetime, time, timedelta, timezone +from datetime import datetime, timedelta, timezone from enum import Enum from typing import Dict, List, Optional, Set, Tuple -import opuslib -import webrtcvad +import opuslib # type: ignore +import webrtcvad # type: ignore from fastapi import APIRouter, Depends from fastapi.websockets import WebSocket, WebSocketDisconnect -from pydub import AudioSegment +from pydub import AudioSegment # type: ignore from starlette.websockets import WebSocketState +from websockets.exceptions import ConnectionClosed import database.conversations as conversations_db import database.users as user_db @@ -50,7 +52,6 @@ from utils.other.task import safe_create_task from utils.pusher import connect_to_trigger_pusher from utils.speaker_identification import detect_speaker_from_text -from utils.stt.streaming import * from utils.stt.streaming import ( STTService, get_stt_service_for_language, @@ -75,7 +76,7 @@ async def _listen( codec: str = 'pcm8', channels: int = 1, include_speech_profile: bool = True, - stt_service: STTService = None, + stt_service: Optional[STTService] = None, including_combined_segments: bool = False, conversation_timeout: int = 120, ): @@ -293,8 +294,8 @@ async def _trigger_create_conversation_with_delay(delay_seconds: int, finished_a except asyncio.CancelledError: pass - async def _create_conversation(conversation: dict): - conversation = Conversation(**conversation) + async def _create_conversation(conversation_data: dict): + conversation = Conversation(**conversation_data) if conversation.status != ConversationStatus.processing: _send_message_event(ConversationEvent(event_type="memory_processing_started", memory=conversation)) conversations_db.update_conversation_status(uid, conversation.id, ConversationStatus.processing) @@ -499,7 +500,7 @@ async def create_conversation_on_segment_received_task(finished_at: datetime): speech_profile_duration = 0 realtime_segment_buffers = [] - realtime_photo_buffers = [] + realtime_photo_buffers: list[ConversationPhoto] = [] def stream_transcript(segments): nonlocal realtime_segment_buffers @@ -612,34 +613,39 @@ def transcript_send(segments, conversation_id): in_progress_conversation_id = conversation_id segment_buffers.extend(segments) - async def transcript_consume(): - nonlocal websocket_active + async def _transcript_flush(auto_reconnect: bool = True): nonlocal segment_buffers nonlocal in_progress_conversation_id nonlocal pusher_ws nonlocal pusher_connected - while websocket_active or len(segment_buffers) > 0: - await asyncio.sleep(1) - if pusher_connected and pusher_ws and len(segment_buffers) > 0: - try: - # 102|data - data = bytearray() - data.extend(struct.pack("I", 102)) - data.extend( - bytes( - json.dumps({"segments": segment_buffers, "memory_id": in_progress_conversation_id}), - "utf-8", - ) + if pusher_connected and pusher_ws and len(segment_buffers) > 0: + try: + # 102|data + data = bytearray() + data.extend(struct.pack("I", 102)) + data.extend( + bytes( + json.dumps({"segments": segment_buffers, "memory_id": in_progress_conversation_id}), + "utf-8", ) - segment_buffers = [] # reset - await pusher_ws.send(data) - except websockets.exceptions.ConnectionClosed as e: - print(f"Pusher transcripts Connection closed: {e}", uid, session_id) - pusher_connected = False - except Exception as e: - print(f"Pusher transcripts failed: {e}", uid, session_id) - if pusher_connected is False: - await connect() + ) + segment_buffers = [] # reset + await pusher_ws.send(data) + except ConnectionClosed as e: + print(f"Pusher transcripts Connection closed: {e}", uid, session_id) + pusher_connected = False + except Exception as e: + print(f"Pusher transcripts failed: {e}", uid, session_id) + if auto_reconnect and pusher_connected is False: + await connect() + + async def transcript_consume(): + nonlocal websocket_active + nonlocal segment_buffers + while websocket_active: + await asyncio.sleep(1) + if len(segment_buffers) > 0: + await _transcript_flush(auto_reconnect=True) # Audio bytes audio_buffers = bytearray() @@ -649,28 +655,39 @@ def audio_bytes_send(audio_bytes): nonlocal audio_buffers audio_buffers.extend(audio_bytes) + async def _audio_bytes_flush(auto_reconnect: bool = True): + nonlocal audio_buffers + nonlocal pusher_ws + nonlocal pusher_connected + if pusher_connected and pusher_ws and len(audio_buffers) > 0: + try: + # 101|data + data = bytearray() + data.extend(struct.pack("I", 101)) + data.extend(audio_buffers.copy()) + audio_buffers = bytearray() # reset + await pusher_ws.send(data) + except ConnectionClosed as e: + print(f"Pusher audio_bytes Connection closed: {e}", uid, session_id) + pusher_connected = False + except Exception as e: + print(f"Pusher audio_bytes failed: {e}", uid, session_id) + if auto_reconnect and pusher_connected is False: + await connect() + async def audio_bytes_consume(): nonlocal websocket_active nonlocal audio_buffers nonlocal pusher_ws nonlocal pusher_connected - while websocket_active or len(audio_buffers) > 0: + while websocket_active: await asyncio.sleep(1) - if pusher_connected and pusher_ws and len(audio_buffers) > 0: - try: - # 101|data - data = bytearray() - data.extend(struct.pack("I", 101)) - data.extend(audio_buffers.copy()) - audio_buffers = bytearray() # reset - await pusher_ws.send(data) - except websockets.exceptions.ConnectionClosed as e: - print(f"Pusher audio_bytes Connection closed: {e}", uid, session_id) - pusher_connected = False - except Exception as e: - print(f"Pusher audio_bytes failed: {e}", uid, session_id) - if pusher_connected is False: - await connect() + if len(audio_buffers) > 0: + await _audio_bytes_flush(auto_reconnect=True) + + async def _flush(): + await _audio_bytes_flush(auto_reconnect=False) + await _transcript_flush(auto_reconnect=False) async def connect(): nonlocal pusher_connected @@ -694,12 +711,13 @@ async def _connect(): nonlocal pusher_connected try: - pusher_ws = await connect_to_trigger_pusher(uid, sample_rate) + pusher_ws = await connect_to_trigger_pusher(uid, sample_rate, retries=5) pusher_connected = True except Exception as e: print(f"Exception in connect: {e}") async def close(code: int = 1000): + await _flush() if pusher_ws: await pusher_ws.close(code) @@ -727,9 +745,15 @@ async def close(code: int = 1000): translation_service = TranslationService() async def translate(segments: List[TranscriptSegment], conversation_id: str): + if not translation_language: + return + try: translated_segments = [] for segment in segments: + if not segment or not segment.id: + continue + segment_text = segment.text.strip() if not segment_text: continue @@ -749,14 +773,14 @@ async def translate(segments: List[TranscriptSegment], conversation_id: str): # Create/Update Translation object translation = Translation(lang=translation_language, text=translated_text) - existing_translation_index = next( - (i for i, t in enumerate(segment.translations) if t.lang == language), None - ) - - if existing_translation_index is not None: - segment.translations[existing_translation_index] = translation - else: - segment.translations.append(translation) + if segment.translations is not None: + existing_translation_index = next( + (i for i, t in enumerate(segment.translations) if t.lang == language), None + ) + if existing_translation_index is not None: + segment.translations[existing_translation_index] = translation + else: + segment.translations.append(translation) translated_segments.append(segment) @@ -886,9 +910,11 @@ async def stream_transcript_process(): ) suggested_segments.add(segment.id) - image_chunks = {} # A temporary in-memory cache for image chunks + image_chunks = {str: any} # A temporary in-memory cache for image chunks - async def process_photo(uid: str, image_b64: str, temp_id: str, send_event_func, photo_buffer: list): + async def process_photo( + uid: str, image_b64: str, temp_id: str, send_event_func, photo_buffer: list[ConversationPhoto] + ): from utils.llm.openglass import describe_image photo_id = str(uuid.uuid4()) @@ -907,14 +933,14 @@ async def process_photo(uid: str, image_b64: str, temp_id: str, send_event_func, await send_event_func(PhotoDescribedEvent(photo_id=photo_id, description=description, discarded=discarded)) async def handle_image_chunk( - uid: str, chunk_data: dict, image_chunks_cache: dict, send_event_func, photo_buffer: list + uid: str, chunk_data: dict, image_chunks_cache: dict, send_event_func, photo_buffer: list[ConversationPhoto] ): temp_id = chunk_data.get('id') index = chunk_data.get('index') total = chunk_data.get('total') data = chunk_data.get('data') - if not all([temp_id, isinstance(index, int), isinstance(total, int), data]): + if not temp_id or not isinstance(index, int) or not isinstance(total, int) or not data: print(f"Invalid image chunk received: {chunk_data}", uid, session_id) return @@ -1110,7 +1136,7 @@ async def listen_handler( codec: str = 'pcm8', channels: int = 1, include_speech_profile: bool = True, - stt_service: STTService = None, + stt_service: Optional[STTService] = None, conversation_timeout: int = 120, ): await _listen(