Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 88 additions & 62 deletions backend/routers/transcribe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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())
Expand All @@ -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

Expand Down Expand Up @@ -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(
Expand Down