# # # Agora Real Time Engagement # Created by Wei Hu in 2024-08. # Copyright (c) 2024 Agora IO. All rights reserved. # # import asyncio import base64 import io import json from enum import Enum import traceback import time import numpy as np from datetime import datetime from typing import Iterable from pydub import AudioSegment from ten import ( AudioFrame, AsyncTenEnv, Cmd, StatusCode, CmdResult, Data, ) from ten.audio_frame import AudioFrameDataFmt from ten_ai_base.const import CMD_PROPERTY_RESULT, CMD_TOOL_CALL from dataclasses import dataclass from ten_ai_base.config import BaseConfig from ten_ai_base.chat_memory import ( ChatMemory, EVENT_MEMORY_EXPIRED, EVENT_MEMORY_APPENDED, ) from ten_ai_base.usage import ( LLMUsage, LLMCompletionTokensDetails, LLMPromptTokensDetails, ) from ten_ai_base.types import ( LLMToolMetadata, LLMToolResult, LLMChatCompletionContentPartParam, ) from ten_ai_base.llm import AsyncLLMBaseExtension from .realtime.connection import RealtimeApiConnection from .realtime.struct import ( AudioFormats, ItemCreate, SessionCreated, ItemCreated, UserMessageItemParam, AssistantMessageItemParam, ItemInputAudioTranscriptionCompleted, ItemInputAudioTranscriptionFailed, ResponseCreated, ResponseDone, ResponseAudioTranscriptDelta, ResponseTextDelta, ResponseAudioTranscriptDone, ResponseTextDone, ResponseOutputItemDone, ResponseOutputItemAdded, ResponseAudioDelta, ResponseAudioDone, InputAudioBufferSpeechStarted, InputAudioBufferSpeechStopped, ResponseFunctionCallArgumentsDone, ErrorMessage, ItemDelete, SessionUpdate, SessionUpdateParams, InputAudioTranscription, ContentType, FunctionCallOutputItemParam, ResponseCreate, ) CMD_IN_FLUSH = "flush" CMD_IN_ON_USER_JOINED = "on_user_joined" CMD_IN_ON_USER_LEFT = "on_user_left" CMD_OUT_FLUSH = "flush" class Role(str, Enum): User = "user" Assistant = "assistant" @dataclass class GLMRealtimeConfig(BaseConfig): base_uri: str = "wss://open.bigmodel.cn" api_key: str = "" path: str = "/api/paas/v4/realtime" prompt: str = "" temperature: float = 0.5 max_tokens: int = 1024 server_vad: bool = True audio_out: bool = True input_transcript: bool = True sample_rate: int = 24000 stream_id: int = 0 dump: bool = False max_history: int = 20 enable_storage: bool = False greeting: str = "" language: str = "en-US" def build_ctx(self) -> dict: return { } class GLMRealtimeExtension(AsyncLLMBaseExtension): def __init__(self, name: str): super().__init__(name) self.ten_env: AsyncTenEnv = None self.conn = None self.session = None self.session_id = None self.config: GLMRealtimeConfig = None self.stopped: bool = False self.connected: bool = False self.buffer: bytearray = b"" self.memory: ChatMemory = None self.total_usage: LLMUsage = LLMUsage() self.users_count = 0 self.stream_id: int = 0 self.remote_stream_id: int = 0 self.channel_name: str = "" self.audio_len_threshold: int = 5120 self.completion_times = [] self.connect_times = [] self.first_token_times = [] self.transcript: str = "" self.ctx: dict = {} self.input_end = time.time() self.input_audio_queue = asyncio.Queue() async def on_init(self, ten_env: AsyncTenEnv) -> None: await super().on_init(ten_env) ten_env.log_debug("on_init") async def on_start(self, ten_env: AsyncTenEnv) -> None: await super().on_start(ten_env) ten_env.log_debug("on_start") self.ten_env = ten_env self.loop = asyncio.get_event_loop() self.loop.create_task(self._on_process_audio()) self.config = await GLMRealtimeConfig.create_async(ten_env=ten_env) ten_env.log_info(f"config: {self.config}") if not self.config.api_key: ten_env.log_error("api_key is required") return try: self.memory = ChatMemory(self.config.max_history) if self.config.enable_storage: [result, _] = await ten_env.send_cmd(Cmd.create("retrieve")) if result.get_status_code() == StatusCode.OK: try: history = json.loads(result.get_property_string("response")) for i in history: self.memory.put(i) ten_env.log_info(f"on retrieve context {history}") except Exception as e: ten_env.log_error(f"Failed to handle retrieve result {e}") else: ten_env.log_warn("Failed to retrieve content") self.memory.on(EVENT_MEMORY_EXPIRED, self._on_memory_expired) self.memory.on(EVENT_MEMORY_APPENDED, self._on_memory_appended) self.ctx = self.config.build_ctx() self.conn = RealtimeApiConnection( ten_env=ten_env, base_uri=self.config.base_uri, path=self.config.path, api_key=self.config.api_key, ) ten_env.log_info("Finish init client") self.loop.create_task(self._loop()) except Exception as e: traceback.print_exc() self.ten_env.log_error(f"Failed to init client {e}") async def on_stop(self, ten_env: AsyncTenEnv) -> None: await super().on_stop(ten_env) ten_env.log_info("on_stop") self.input_audio_queue.put_nowait(None) self.stopped = True async def on_audio_frame(self, _: AsyncTenEnv, audio_frame: AudioFrame) -> None: try: stream_id = audio_frame.get_property_int("stream_id") if self.channel_name == "": self.channel_name = audio_frame.get_property_string("channel") if self.remote_stream_id == 0: self.remote_stream_id = stream_id frame_buf = audio_frame.get_buf() self.input_audio_queue.put_nowait(frame_buf) if not self.config.server_vad: self.input_end = time.time() except Exception as e: traceback.print_exc() self.ten_env.log_error(f"GLMV2VExtension on audio frame failed {e}") async def on_cmd(self, ten_env: AsyncTenEnv, cmd: Cmd) -> None: cmd_name = cmd.get_name() ten_env.log_debug("on_cmd name {}".format(cmd_name)) status = StatusCode.OK detail = "success" if cmd_name == CMD_IN_FLUSH: # Will only flush if it is client side vad await self._flush() await ten_env.send_cmd(Cmd.create(CMD_OUT_FLUSH)) ten_env.log_info("on flush") elif cmd_name == CMD_IN_ON_USER_JOINED: self.users_count += 1 # Send greeting when first user joined if self.users_count == 1: await self._greeting() elif cmd_name == CMD_IN_ON_USER_LEFT: self.users_count -= 1 else: # Register tool await super().on_cmd(ten_env, cmd) return cmd_result = CmdResult.create(status) cmd_result.set_property_string("detail", detail) await ten_env.return_result(cmd_result, cmd) # Not support for now async def on_data(self, ten_env: AsyncTenEnv, data: Data) -> None: pass async def _on_process_audio(self) -> None: while True: try: audio_frame = await self.input_audio_queue.get() if audio_frame is None: break self._dump_audio_if_need(audio_frame, Role.User) if self.connected: wav_buff = self.convert_to_wav_in_memory(audio_frame) await self.conn.send_audio_data(wav_buff) except Exception as e: traceback.print_exc() self.ten_env.log_error(f"Error processing audio frame {e}") async def _loop(self): def get_time_ms() -> int: current_time = datetime.now() return current_time.microsecond // 1000 try: start_time = time.time() await self.conn.connect() self.connect_times.append(time.time() - start_time) item_id = "" # For truncate response_id = "" # content_index = 0 relative_start_ms = get_time_ms() flushed = set() self.ten_env.log_info("Client loop started") async for message in self.conn.listen(): try: # self.ten_env.log_info(f"Received message: {message.type}") match message: case SessionCreated(): self.ten_env.log_info( f"Session is created: {message.session}" ) self.session_id = message.session.id self.session = message.session await self._update_session() history = self.memory.get() for h in history: if h["role"] == "user": await self.conn.send_request( ItemCreate( item=UserMessageItemParam( content=[ { "type": ContentType.InputText, "text": h["content"], } ] ) ) ) elif h["role"] == "assistant": await self.conn.send_request( ItemCreate( item=AssistantMessageItemParam( content=[ { "type": ContentType.InputText, "text": h["content"], } ] ) ) ) self.ten_env.log_info(f"Finish send history {history}") self.memory.clear() if not self.connected: self.connected = True await self._greeting() case ItemInputAudioTranscriptionCompleted(): self.ten_env.log_info( f"On request transcript {message.transcript}" ) self._send_transcript(message.transcript, Role.User, True) self.memory.put( { "role": "user", "content": message.transcript, # "id": message.item_id, } ) case ItemInputAudioTranscriptionFailed(): self.ten_env.log_warn( f"On request transcript failed {message.item_id} {message.error}" ) case ItemCreated(): self.ten_env.log_info(f"On item created {message.item}") case ResponseCreated(): response_id = message.response.id self.ten_env.log_info(f"On response created {response_id}") case ResponseDone(): msg_resp_id = message.response.id status = message.response.status if msg_resp_id == response_id: response_id = "" self.ten_env.log_info( f"On response done {msg_resp_id} {status} {message.response.usage}" ) # workaround as GLM does not have responseAudioTranscriptDone self.transcript = "" self._send_transcript("", Role.Assistant, True) if message.response.usage: pass # await self._update_usage(message.response.usage) case ResponseAudioTranscriptDelta(): self.ten_env.log_info( f"On response transcript delta {message.output_index} {message.content_index} {message.delta}" ) if message.response_id in flushed: self.ten_env.log_warn( f"On flushed transcript delta {message.output_index} {message.content_index} {message.delta}" ) continue self._send_transcript(message.delta, Role.Assistant, False) case ResponseTextDelta(): self.ten_env.log_info( f"On response text delta {message.output_index} {message.content_index} {message.delta}" ) # if message.response_id in flushed: # self.ten_env.log_warn( # f"On flushed text delta {message.output_index} {message.content_index} {message.delta}" # ) # continue # if item_id != message.item_id: # item_id = message.item_id # self.first_token_times.append( # time.time() - self.input_end # ) self._send_transcript(message.delta, Role.Assistant, False) case ResponseAudioTranscriptDone(): # this is not triggering by GLM self.ten_env.log_info( f"On response transcript done {message.output_index} {message.content_index} {message.transcript}" ) if message.response_id in flushed: self.ten_env.log_warn( "On flushed transcript done" ) continue self.memory.put( { "role": "assistant", "content": message.transcript, # "id": message.item_id, } ) self.transcript = "" self._send_transcript("", Role.Assistant, True) case ResponseTextDone(): self.ten_env.log_info( f"On response text done {message.output_index} {message.content_index} {message.text}" ) # if message.response_id in flushed: # self.ten_env.log_warn( # f"On flushed text done {message.response_id}" # ) # continue self.completion_times.append(time.time() - self.input_end) self.transcript = "" self._send_transcript("", Role.Assistant, True) case ResponseOutputItemDone(): self.ten_env.log_info(f"Output item done {message.item}") case ResponseOutputItemAdded(): self.ten_env.log_info( f"Output item added {message.output_index} {message.item}" ) case ResponseAudioDelta(): # if message.response_id in flushed: # self.ten_env.log_warn( # f"On flushed audio delta {message.response_id} {message.item_id} {message.content_index}" # ) # continue # if item_id != message.item_id: # item_id = message.item_id # self.first_token_times.append( # time.time() - self.input_end # ) # content_index = message.content_index await self._on_audio_delta(message.delta) case ResponseAudioDone(): self.completion_times.append(time.time() - self.input_end) case InputAudioBufferSpeechStarted(): self.ten_env.log_info( f"On server listening, in response {response_id}, last item {item_id}" ) # Tuncate the on-going audio stream # end_ms = get_time_ms() - relative_start_ms # if item_id: # truncate = ItemTruncate( # item_id=item_id, # content_index=content_index, # audio_end_ms=end_ms, # ) # await self.conn.send_request(truncate) if self.config.server_vad: await self._flush() if response_id and self.transcript: transcript = self.transcript + "[interrupted]" self._send_transcript(transcript, Role.Assistant, True) self.transcript = "" # memory leak, change to lru later flushed.add(response_id) item_id = "" case InputAudioBufferSpeechStopped(): # Only for server vad self.input_end = time.time() relative_start_ms = get_time_ms() - message.audio_end_ms self.ten_env.log_info( f"On server stop listening, {message.audio_end_ms}, relative {relative_start_ms}" ) case ResponseFunctionCallArgumentsDone(): # tool_call_id = message.call_id name = message.name arguments = message.arguments self.ten_env.log_info(f"need to call func {name}") self.loop.create_task( self._handle_tool_call(name, arguments) ) case ErrorMessage(): self.ten_env.log_error( f"Error message received: {message.error}" ) case _: self.ten_env.log_debug(f"Not handled message {message}") except Exception as e: traceback.print_exc() self.ten_env.log_error(f"Error processing message: {message} {e}") self.ten_env.log_info("Client loop finished") except Exception as e: traceback.print_exc() self.ten_env.log_error(f"Failed to handle loop {e}") # clear so that new session can be triggered self.connected = False self.remote_stream_id = 0 if not self.stopped: await self.conn.close() await asyncio.sleep(0.5) self.ten_env.log_info("Reconnect") self.conn = RealtimeApiConnection( ten_env=self.ten_env, base_uri=self.config.base_uri, path=self.config.path, api_key=self.config.api_key, ) self.loop.create_task(self._loop()) async def _on_memory_expired(self, message: dict) -> None: self.ten_env.log_info(f"Memory expired: {message}") item_id = message.get("item_id") if item_id: await self.conn.send_request(ItemDelete(item_id=item_id)) async def _on_memory_appended(self, message: dict) -> None: self.ten_env.log_info(f"Memory appended: {message}") if not self.config.enable_storage: return role = message.get("role") stream_id = self.remote_stream_id if role == Role.User else 0 try: d = Data.create("append") d.set_property_string("text", message.get("content")) d.set_property_string("role", role) d.set_property_int("stream_id", stream_id) asyncio.create_task(self.ten_env.send_data(d)) except Exception as e: self.ten_env.log_error(f"Error send append_context data {message} {e}") # Direction: IN def convert_to_wav_in_memory(self, buff: bytearray) -> bytes: """ Converts the accumulated PCM data to WAV format in-memory. Returns the WAV data as bytes. """ # Convert PCM data to numpy array of int16 type pcm_data = np.frombuffer(buff, dtype=np.int16) # Use pydub to create an AudioSegment audio_segment = AudioSegment( pcm_data.tobytes(), frame_rate=24000, sample_width=2, channels=1 ) # Create an in-memory stream to store the WAV file memory_stream = io.BytesIO() # Export the AudioSegment to the in-memory stream as WAV audio_segment.export(memory_stream, format="wav") # Return the WAV data as bytes wav_bytes = memory_stream.getvalue() return wav_bytes async def _update_session(self) -> None: tools = [] def tool_dict(tool: LLMToolMetadata): t = { "type": "function", "name": tool.name, "description": tool.description, "parameters": { "type": "object", "properties": {}, "required": [], "additionalProperties": False, }, } for param in tool.parameters: t["parameters"]["properties"][param.name] = { "type": param.type, "description": param.description, } if param.required: t["parameters"]["required"].append(param.name) return t if self.available_tools: tool_prompt = "You have several tools that you can get help from:\n" for t in self.available_tools: tool_prompt += f"- ***{t.name}***: {t.description}" self.ctx["tools"] = tool_prompt tools = [tool_dict(t) for t in self.available_tools] prompt = self._replace(self.config.prompt) self.ten_env.log_info(f"update session {prompt} {tools}") su = SessionUpdate( session=SessionUpdateParams( instructions=prompt, input_audio_format=AudioFormats.WAV24, output_audio_format=AudioFormats.PCM, tools=tools, ) ) if self.config.audio_out: # su.session.voice = self.config.voice pass else: su.session.modalities = ["text"] if self.config.input_transcript: su.session.input_audio_transcription = InputAudioTranscription( model="whisper-1" ) await self.conn.send_request(su) async def on_tools_update(self, _: AsyncTenEnv, tool: LLMToolMetadata) -> None: """Called when a new tool is registered. Implement this method to process the new tool.""" self.ten_env.log_info(f"on tools update {tool}") # await self._update_session() def _replace(self, prompt: str) -> str: result = prompt for token, value in self.ctx.items(): result = result.replace("{" + token + "}", value) return result # Direction: OUT async def _on_audio_delta(self, delta: bytes) -> None: audio_data = base64.b64decode(delta) self.ten_env.log_debug( f"on_audio_delta audio_data len {len(audio_data)} samples {len(audio_data) // 2}" ) self._dump_audio_if_need(audio_data, Role.Assistant) f = AudioFrame.create("pcm_frame") f.set_sample_rate(self.config.sample_rate) f.set_bytes_per_sample(2) f.set_number_of_channels(1) f.set_data_fmt(AudioFrameDataFmt.INTERLEAVE) f.set_samples_per_channel(len(audio_data) // 2) f.alloc_buf(len(audio_data)) buff = f.lock_buf() buff[:] = audio_data f.unlock_buf(buff) await self.ten_env.send_audio_frame(f) def _send_transcript(self, content: str, role: Role, is_final: bool) -> None: def is_punctuation(char): if char in [",", "，", ".", "。", "?", "？", "!", "！"]: return True return False def parse_sentences(sentence_fragment, content): sentences = [] current_sentence = sentence_fragment for char in content: current_sentence += char if is_punctuation(char): # Check if the current sentence contains non-punctuation characters stripped_sentence = current_sentence if any(c.isalnum() for c in stripped_sentence): sentences.append(stripped_sentence) current_sentence = "" # Reset for the next sentence remain = current_sentence # Any remaining characters form the incomplete sentence return sentences, remain def send_data( ten_env: AsyncTenEnv, sentence: str, stream_id: int, role: str, is_final: bool, ): try: d = Data.create("text_data") d.set_property_string("text", sentence) d.set_property_bool("end_of_segment", is_final) d.set_property_string("role", role) d.set_property_int("stream_id", stream_id) ten_env.log_info( f"send transcript text [{sentence}] stream_id {stream_id} is_final {is_final} end_of_segment {is_final} role {role}" ) asyncio.create_task(ten_env.send_data(d)) except Exception as e: ten_env.log_error( f"Error send text data {role}: {sentence} {is_final} {e}" ) stream_id = self.remote_stream_id if role == Role.User else 0 try: if role == Role.Assistant and not is_final: sentences, self.transcript = parse_sentences(self.transcript, content) for s in sentences: send_data(self.ten_env, s, stream_id, role, is_final) else: send_data(self.ten_env, content, stream_id, role, is_final) except Exception as e: self.ten_env.log_error( f"Error send text data {role}: {content} {is_final} {e}" ) def _dump_audio_if_need(self, buf: bytearray, role: Role) -> None: if not self.config.dump: return with open("{}_{}.pcm".format(role, self.channel_name), "ab") as dump_file: dump_file.write(buf) async def _handle_tool_call( self, name: str, arguments: str ) -> None: self.ten_env.log_info(f"_handle_tool_call {name} {arguments}") cmd: Cmd = Cmd.create(CMD_TOOL_CALL) cmd.set_property_string("name", name) cmd.set_property_from_json("arguments", arguments) [result, _] = await self.ten_env.send_cmd(cmd) tool_response = ItemCreate( item=FunctionCallOutputItemParam( output='{"success":false}', ) ) if result.get_status_code() == StatusCode.OK: tool_result: LLMToolResult = json.loads( result.get_property_to_json(CMD_PROPERTY_RESULT) ) result_content = tool_result["content"] tool_response.item.output = json.dumps( self._convert_to_content_parts(result_content) ) self.ten_env.log_info(f"tool_result: {tool_result}") else: self.ten_env.log_error("Tool call failed") await self.conn.send_request(tool_response) await self.conn.send_request(ResponseCreate()) self.ten_env.log_info(f"_remote_tool_call finish {name} {arguments}") def _greeting_text(self) -> str: text = "Hi, there." if self.config.language == "zh-CN": text = "你好。" elif self.config.language == "ja-JP": text = "こんにちは" elif self.config.language == "ko-KR": text = "안녕하세요" return text def _convert_tool_params_to_dict(self, tool: LLMToolMetadata): json_dict = {"type": "object", "properties": {}, "required": []} for param in tool.parameters: json_dict["properties"][param.name] = { "type": param.type, "description": param.description, } if param.required: json_dict["required"].append(param.name) return json_dict def _convert_to_content_parts( self, content: Iterable[LLMChatCompletionContentPartParam] ): content_parts = [] if isinstance(content, str): content_parts.append({"type": "text", "text": content}) else: for part in content: # Only text content is supported currently for v2v model if part["type"] == "text": content_parts.append(part) return content_parts async def _greeting(self) -> None: if self.connected and self.users_count == 1: # somehow it's not working text = self._greeting_text() if self.config.greeting: text = "Say '" + self.config.greeting + "' to me." self.ten_env.log_info(f"send greeting {text}") # await self.conn.send_request( # ItemCreate( # item=UserMessageItemParam( # content=[{"type": ContentType.InputText, "text": text}] # ) # ) # ) # await self.conn.send_request(ResponseCreate()) async def _flush(self) -> None: try: c = Cmd.create("flush") await self.ten_env.send_cmd(c) except Exception: self.ten_env.log_error("Error flush") async def _update_usage(self, usage: dict) -> None: self.total_usage.completion_tokens += usage.get("output_tokens") or 0 self.total_usage.prompt_tokens += usage.get("input_tokens") or 0 self.total_usage.total_tokens += usage.get("total_tokens") or 0 if not self.total_usage.completion_tokens_details: self.total_usage.completion_tokens_details = LLMCompletionTokensDetails() if not self.total_usage.prompt_tokens_details: self.total_usage.prompt_tokens_details = LLMPromptTokensDetails() if usage.get("output_token_details"): self.total_usage.completion_tokens_details.accepted_prediction_tokens += ( usage["output_token_details"].get("text_tokens") ) self.total_usage.completion_tokens_details.audio_tokens += usage[ "output_token_details" ].get("audio_tokens") if usage.get("input_token_details:"): self.total_usage.prompt_tokens_details.audio_tokens += usage[ "input_token_details" ].get("audio_tokens") self.total_usage.prompt_tokens_details.cached_tokens += usage[ "input_token_details" ].get("cached_tokens") self.total_usage.prompt_tokens_details.text_tokens += usage[ "input_token_details" ].get("text_tokens") self.ten_env.log_info(f"total usage: {self.total_usage}") data = Data.create("llm_stat") data.set_property_from_json("usage", json.dumps(self.total_usage.model_dump())) if self.connect_times and self.completion_times and self.first_token_times: data.set_property_from_json( "latency", json.dumps( { "connection_latency_95": np.percentile(self.connect_times, 95), "completion_latency_95": np.percentile( self.completion_times, 95 ), "first_token_latency_95": np.percentile( self.first_token_times, 95 ), "connection_latency_99": np.percentile(self.connect_times, 99), "completion_latency_99": np.percentile( self.completion_times, 99 ), "first_token_latency_99": np.percentile( self.first_token_times, 99 ), } ), ) asyncio.create_task(self.ten_env.send_data(data)) async def on_call_chat_completion(self, async_ten_env, **kargs): raise NotImplementedError async def on_data_chat_completion(self, async_ten_env, **kargs): raise NotImplementedError