save code

main.py

@@ -18,7 +18,7 @@ from pydantic import BaseModel
 
 from config import settings
 from core.pipeline import ChatPipeline
-from core.state_machine import Arbitrator
+from core.state_machine import Arbitrator, SessionState
 from services.asr import ASRService
 from services.avatar import AvatarService
 from services.llm import LLMService
@@ -52,6 +52,7 @@ audio_bus = AudioBus(sample_rate=48000)
 pcs: set[RTCPeerConnection] = set()
 subtitle_clients: set[WebSocket] = set()
 animation_clients: set[WebSocket] = set()
+MAX_BUFFERED_AUDIO_SAMPLES = 16000 * 45
 
 
 @dataclass
@@ -73,7 +74,9 @@ class RuntimeState:
     last_animation_mode: str = settings.avatar_driver_mode
     last_animation_frame_count: int = 0
     busy: bool = False
+    pending_voice_turn: bool = False
     buffer_16k: list[np.ndarray] = field(default_factory=list)
+    buffered_16k_samples: int = 0
 
 
 runtime = RuntimeState()
@@ -155,6 +158,19 @@ def _runtime_payload() -> dict:
     }
 
 
+def _clear_buffered_audio() -> None:
+    runtime.buffer_16k.clear()
+    runtime.buffered_16k_samples = 0
+
+
+def _append_buffered_audio(chunk: np.ndarray) -> None:
+    runtime.buffer_16k.append(chunk)
+    runtime.buffered_16k_samples += int(chunk.shape[0])
+    while runtime.buffered_16k_samples > MAX_BUFFERED_AUDIO_SAMPLES and runtime.buffer_16k:
+        dropped = runtime.buffer_16k.pop(0)
+        runtime.buffered_16k_samples -= int(dropped.shape[0])
+
+
 def _resample_mono(audio: np.ndarray, from_sr: int, to_sr: int) -> np.ndarray:
     if from_sr == to_sr:
         return audio.astype(np.float32, copy=False)
@@ -202,15 +218,16 @@ async def _enqueue_audio_and_animation(
 
 
 async def _run_pipeline_from_buffer() -> None:
-    if runtime.busy or not runtime.buffer_16k:
+    if runtime.busy or not runtime.buffer_16k or not runtime.pending_voice_turn:
         return
 
     runtime.busy = True
+    runtime.pending_voice_turn = False
     start_t = time.perf_counter()
     try:
         runtime.last_input_mode = "voice"
         audio_16k = np.concatenate(runtime.buffer_16k, axis=0)
-        runtime.buffer_16k.clear()
+        _clear_buffered_audio()
 
         async def on_user_text_now(user_text: str) -> None:
             runtime.last_user_text = user_text
@@ -245,6 +262,8 @@ async def _run_pipeline_from_buffer() -> None:
             await _broadcast_animation(await avatar_service.build_controls_from_audio(audio, sr, text=result["reply_text"]))
     finally:
         runtime.busy = False
+        if runtime.pending_voice_turn and runtime.buffer_16k:
+            asyncio.create_task(_run_pipeline_from_buffer())
 
 
 async def _consume_user_audio(track) -> None:
@@ -256,7 +275,7 @@ async def _consume_user_audio(track) -> None:
         layout = getattr(frame, "layout", None)
         channels = len(layout.channels) if layout and layout.channels else 1
         mono_16k = _resample_to_16k_mono(pcm, sample_rate, channels)
-        runtime.buffer_16k.append(mono_16k)
+        _append_buffered_audio(mono_16k)
         vad_buffer = np.concatenate([vad_buffer, mono_16k], axis=0)
 
         while vad_buffer.shape[0] >= 512:
@@ -269,7 +288,7 @@ async def _consume_user_audio(track) -> None:
                 event = None
             if event and "start" in event:
                 barge_t0 = time.perf_counter()
-                was_avatar_speaking = str(arbitrator.state) == "avatar_speaking"
+                was_avatar_speaking = arbitrator.state == SessionState.AVATAR_SPEAKING
                 await arbitrator.on_speech_start()
                 await audio_bus.clear()
                 await _broadcast_animation(await avatar_service.build_reset_payload(reason="speech-start"))
@@ -280,6 +299,7 @@ async def _consume_user_audio(track) -> None:
             if event and "end" in event:
                 await arbitrator.on_speech_end()
                 runtime.vad_end_count += 1
+                runtime.pending_voice_turn = True
                 asyncio.create_task(_run_pipeline_from_buffer())
 
 
@@ -400,6 +420,8 @@ async def chat_text(req: TextChatRequest):
 @app.post("/chat/reset")
 async def chat_reset():
     llm_service.clear_history()
+    _clear_buffered_audio()
+    runtime.pending_voice_turn = False
     runtime.last_user_text = ""
     runtime.last_reply_text = ""
     runtime.pipeline_runs = 0
@@ -415,6 +437,7 @@ async def chat_reset():
     runtime.vad_start_count = 0
     runtime.vad_end_count = 0
     runtime.last_animation_frame_count = 0
+    await audio_bus.clear()
     await _broadcast_animation(await avatar_service.build_reset_payload(reason="chat-reset"))
     return {"ok": True}