feat: support dump audio by session

Author: sunxilin

interface/ten_ai_base/asr.py

@@ -3,13 +3,35 @@
 # Licensed under the Apache License, Version 2.0.
 # See the LICENSE file for more information.
 #
+"""ASR base class: connection, buffering, metrics and result helpers."""
 from abc import abstractmethod
+import asyncio
+from functools import wraps
+import json
+import os
 from typing import Any, final
 import uuid
 
-from .struct import ASRResult
-from .types import ASRBufferConfig, ASRBufferConfigModeDiscard, ASRBufferConfigModeKeep
+from ten_runtime import (
+    AsyncExtension,
+    AsyncTenEnv,
+    Cmd,
+    Data,
+    AudioFrame,
+    StatusCode,
+    CmdResult,
+)
 
+from .const import (
+    DATA_IN_ASR_FINALIZE,
+    DATA_OUT_ASR_FINALIZE_END,
+    DATA_OUT_METRICS,
+    PROPERTY_KEY_DUMP,
+    PROPERTY_KEY_DUMP_PATH,
+    PROPERTY_KEY_METADATA,
+    PROPERTY_KEY_SESSION_ID,
+)
+from .dumper import Dumper
 from .message import (
     ModuleError,
     ModuleErrorCode,
@@ -18,28 +40,16 @@
     ModuleMetrics,
     ModuleType,
 )
+from .struct import ASRResult
 from .timeline import AudioTimeline
-from .const import (
-    DATA_IN_ASR_FINALIZE,
-    DATA_OUT_ASR_FINALIZE_END,
-    DATA_OUT_METRICS,
-    PROPERTY_KEY_METADATA,
-    PROPERTY_KEY_SESSION_ID,
-)
-from ten_runtime import (
-    AsyncExtension,
-    AsyncTenEnv,
-    Cmd,
-    Data,
-    AudioFrame,
-    StatusCode,
-    CmdResult,
-)
-import asyncio
-import json
+from .types import ASRBufferConfig, ASRBufferConfigModeDiscard, ASRBufferConfigModeKeep
+
 
+class AsyncASRBaseExtension(
+    AsyncExtension
+):  # pylint: disable=too-many-instance-attributes, too-many-public-methods
+    """Asynchronous base class for ASR modules."""
 
-class AsyncASRBaseExtension(AsyncExtension):
     def __init__(self, name: str):
         super().__init__(name)
 
@@ -52,6 +62,7 @@ def __init__(self, name: str):
         self.buffered_frames_size = 0
         self.audio_frames_queue = asyncio.Queue[AudioFrame]()
         self.audio_timeline = AudioTimeline()
+        self.dumper: Dumper | None = None
         self.audio_actual_send_metrics_task: asyncio.Task[None] | None = None
         self.uuid = self._get_uuid()  # Unique identifier for the current final turn
 
@@ -68,6 +79,18 @@ async def on_init(self, ten_env: AsyncTenEnv) -> None:
         self.ten_env = ten_env
         asyncio.create_task(self._audio_frame_consumer())
 
+        enable_dump, err = await ten_env.get_property_bool(PROPERTY_KEY_DUMP)
+        if err:
+            ten_env.log_info(f"dump not set, disable dump: {err}")
+        elif enable_dump:
+            dump_path, err = await ten_env.get_property_string(PROPERTY_KEY_DUMP_PATH)
+            if err:
+                ten_env.log_warn(f"dump_path not set, use current directory: {err}")
+                dump_path = "."
+
+            dump_file_path = os.path.join(dump_path, self.dump_file_name())
+            self.dumper = Dumper(dump_file_path, None)
+
     async def on_start(self, ten_env: AsyncTenEnv) -> None:
         ten_env.log_info("on_start")
 
@@ -127,6 +150,10 @@ def vendor(self) -> str:
         """Get the name of the ASR vendor."""
         raise NotImplementedError("This method should be implemented in subclasses.")
 
+    def dump_file_name(self) -> str:
+        """Return the base dump filename."""
+        return f"{self.name}_out.pcm"
+
     @abstractmethod
     async def start_connection(self) -> None:
         """Start the connection to the ASR service."""
@@ -163,6 +190,26 @@ def input_audio_sample_width(self) -> int:
         """
         return 2
 
+    # Automatically wrap subclass start_connection to update dumper session first
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        orig = cls.__dict__.get("start_connection")
+        if orig is None:
+            return
+
+        # Only wrap coroutine functions
+        @wraps(orig)
+        async def wrapped(self: AsyncASRBaseExtension, *args, **kw):
+            try:
+                if self.dumper is not None:
+                    await self.dumper.update_session()
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                if self.ten_env is not None:
+                    self.ten_env.log_error(f"auto update_session failed: {e}")
+            return await orig(self, *args, **kw)
+
+        setattr(cls, "start_connection", wrapped)
+
     def buffer_strategy(self) -> ASRBufferConfig:
         """
         Get the buffer strategy for audio frames when not connected
@@ -442,7 +489,7 @@ async def _audio_frame_consumer(self) -> None:
                 await self._handle_audio_frame(self.ten_env, audio_frame)
             except asyncio.CancelledError:
                 break
-            except Exception as e:
+            except Exception as e:  # pylint: disable=broad-exception-caught
                 self.ten_env.log_error(f"Error consuming audio frame: {e}")
 
     async def _send_audio_actual_send_metrics_task(self) -> None:

interface/ten_ai_base/const.py

@@ -35,3 +35,5 @@
 
 PROPERTY_KEY_METADATA = "metadata"
 PROPERTY_KEY_SESSION_ID = "session_id"
+PROPERTY_KEY_DUMP = "dump"
+PROPERTY_KEY_DUMP_PATH = "dump_path"

interface/ten_ai_base/dumper.py

@@ -3,34 +3,75 @@
 # Licensed under the Apache License, Version 2.0.
 # See the LICENSE file for more information.
 #
+"""Simple async dumper for writing bytes to session-rotated files."""
 
+import asyncio
+from pathlib import Path
 import aiofiles
-import os
 
 
 class Dumper:
-    def __init__(self, dump_file_path: str):
-        self.dump_file_path: str = dump_file_path
+    """Asynchronous dumper that writes to files with session-based names."""
+
+    def __init__(
+        self,
+        base_dump_file_path: str,
+        session_name: str | None = None,
+        delimiter: str = "_",
+    ):
+        self.base_dump_file_path: str = base_dump_file_path
+        self.session_name: str | None = session_name
+        self.delimiter: str = delimiter
         self._file: aiofiles.threadpool.binary.AsyncBufferedIOBase | None = None
+        self._lock: asyncio.Lock = asyncio.Lock()
+
+    @property
+    def dump_file_path(self) -> str:
+        base_path = Path(self.base_dump_file_path)
+        if not self.session_name:
+            return str(base_path)
+
+        stem = base_path.stem
+        suffix = base_path.suffix
+        new_name = f"{stem}{self.delimiter}{self.session_name}{suffix}"
+        return str(base_path.with_name(new_name))
+
+    async def close(self) -> None:
+        async with self._lock:
+            if self._file:
+                await self._file.close()
+                self._file = None
+
+    async def update_session(self) -> None:
+        """Rotate to a new session file and open it if needed.
 
-    async def start(self):
-        if self._file:
-            return
+        The session name is generated from current event-loop time
+        to ensure uniqueness and ordering. If the generated name equals
+        the current one, this function is a no-op.
+        """
+        async with self._lock:
+            # Generate a new session name based on timestamp
+            current_time = asyncio.get_event_loop().time()
+            new_session_name = f"{current_time:.6f}"
 
-        os.makedirs(os.path.dirname(self.dump_file_path), exist_ok=True)
+            if new_session_name == self.session_name and self._file is not None:
+                # Already opened on the same session
+                return
 
-        self._file = await aiofiles.open(self.dump_file_path, mode="wb")
+            # Close previous file if any
+            if self._file is not None:
+                await self._file.close()
+                self._file = None
 
-    async def stop(self):
-        if self._file:
-            await self._file.close()
-            self._file = None
+            # Update session and open the new file
+            self.session_name = new_session_name
+            Path(self.dump_file_path).parent.mkdir(parents=True, exist_ok=True)
+            self._file = await aiofiles.open(self.dump_file_path, mode="wb")
 
-    async def push_bytes(self, data: bytes):
-        if not self._file:
-            raise RuntimeError(
-                "Dumper for {} is not opened. Please start the Dumper first.".format(
-                    self.dump_file_path
+    async def push_bytes(self, data: bytes) -> int:
+        async with self._lock:
+            if not self._file:
+                raise RuntimeError(
+                    f"Dumper for {self.dump_file_path} is not opened. Call update_session() first."
                 )
-            )
-        _ = await self._file.write(data)
+            return await self._file.write(data)