大功告成

core/pipeline.py

@@ -11,6 +11,9 @@ from services.llm import LLMService
 from services.tts import TTSService
 
 
+AudioChunkCallback = Callable[[np.ndarray, int, str, int, int], Awaitable[None]]
+
+
 class ChatPipeline:
     def __init__(
         self,
@@ -75,7 +78,7 @@ class ChatPipeline:
     async def _synthesize_sentence_first(
         self,
         reply_text: str,
-        on_audio_chunk: Optional[Callable[[np.ndarray, int], Awaitable[None]]] = None,
+        on_audio_chunk: Optional[AudioChunkCallback] = None,
         on_text_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
         should_interrupt: Optional[Callable[[], bool]] = None,
     ) -> tuple[np.ndarray, int, int, int]:
@@ -94,7 +97,7 @@ class ChatPipeline:
             if i == 0:
                 first_chunk_ms = int((time.perf_counter() - t0) * 1000)
             if on_audio_chunk is not None:
-                await on_audio_chunk(seg_audio, seg_sr)
+                await on_audio_chunk(seg_audio, seg_sr, seg, i, len(segments))
             audio_chunks.append(seg_audio)
         tts_total_ms = int((time.perf_counter() - t0) * 1000)
         audio = np.concatenate(audio_chunks) if audio_chunks else np.zeros(1, dtype=np.float32)
@@ -103,7 +106,7 @@ class ChatPipeline:
     async def process_turn(
         self,
         audio_16k: np.ndarray,
-        on_audio_chunk: Optional[Callable[[np.ndarray, int], Awaitable[None]]] = None,
+        on_audio_chunk: Optional[AudioChunkCallback] = None,
         on_user_text: Optional[Callable[[str], Awaitable[None]]] = None,
         on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None,
         on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
@@ -145,7 +148,7 @@ class ChatPipeline:
     async def process_text_turn(
         self,
         user_text: str,
-        on_audio_chunk: Optional[Callable[[np.ndarray, int], Awaitable[None]]] = None,
+        on_audio_chunk: Optional[AudioChunkCallback] = None,
         on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None,
         on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None,
     ) -> tuple[dict, np.ndarray, int]: