rtcllm/rtcllm/rtc.py
2024-09-02 18:43:58 +08:00

206 lines
5.5 KiB
Python

import asyncio
import random
import json
from functools import partial
from appPublic.dictObject import DictObject
from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription, RTCIceCandidate
from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder, MediaRelay
# from websockets.asyncio.client import connect
from websockets.asyncio.client import connect
import websockets
from stt import asr
from vad import AudioTrackVad
from aav import MyAudioTrack, MyVideoTrack
videos = ['./1.mp4', './2.mp4']
class RTCLLM:
def __init__(self, ws_url, iceServers):
self.ws_url = ws_url
self.iceServers = iceServers
self.peers = DictObject()
self.vid = 0
self.info = {
'id':'rtcllm_agent',
'name':'rtcllm_agent'
}
self.dc = None
self.loop = asyncio.get_event_loop()
def get_pc(self, data):
return self.peers[data['from'].id].pc
async def run(self):
async with connect(self.ws_url) as self.ws:
await self.login()
while True:
msg = await self.ws.recv()
data = DictObject(**json.loads(msg))
d = data.data
print(f'ws recv(): {d.type=}')
func = self.handlers.get(d.type)
if func:
f = partial(func, self)
await f(d)
self.ws.close()
async def login(self):
await self.ws.send(json.dumps({
'type':'login',
'info':self.info
}))
async def on_icecandidate(self, pc, candidate):
print('on_icecandidate()', self, pc, candidate)
if candidate:
await self.ws.send(json.dumps({
"type":"candidate",
"candidate":candidate,
}))
async def save_onlineList(self, data):
print(f'{self}, {type(self)}')
self.onlineList = data.onlineList
async def vad_voiceend(self, peer, audio):
txt = await asr(audio)
# await peer.dc.send(txt)
async def auto_accept_call(self, data):
opts = DictObject(iceServers=self.iceServers)
pc = RTCPeerConnection(opts)
self.peers[data['from'].id] = DictObject(**{
'info':data['from'],
'pc':pc
})
await self.ws.send(json.dumps({'type':'callAccepted', 'to':data['from']}))
async def pc_track(self, peerid, track):
peer = self.peers[peerid]
pc = peer.pc
if track.kind == 'audio':
f = partial(self.vad_voiceend, peer)
vadtrack = AudioTrackVad(track, stage=3, onvoiceend=f)
peer.vadtrack = vadtrack
vadtrack.start_vad()
def play_random_media(self, vt, at):
i = random.randint(0,1)
player = MediaPlayer(videos[i])
vt.set_source(player)
at.set_source(player)
f = partial(self.play_random_media, vt, at)
self.loop.call_later(180, f)
async def pc_connectionState_changed(self, peerid):
peer = self.peers[peerid]
pc = peer.pc
peer.audiotrack = MyAudioTrack()
peer.videotrack = MyVideoTrack()
pc.addTrack(peer.audiotrack)
pc.addTrack(peer.videotrack)
self.play_random_media(peer.videotrack, peer.audiotrack)
if pc.connectionState == 'connected':
peer.dc = pc.createDataChannel(peer.info.name)
return
if pc.connectionState == 'closed':
await pc.close()
if peer.dc:
await peer.dc.close()
peers = {
k:v for k,v in self.peers.items() if k != peerid
}
self.peers = peers
if len([k for k in self.peers.keys()]) == 0:
await self.ws.close()
def play_video(self, peerid):
print('play video ........................')
pc = self.peers[peerid].pc
"""
player = MediaPlayer(videos[0])
if player:
pc.addTrack(player.audio)
pc.addTrack(player.video)
"""
player = MediaPlayer(videos[1])
if player:
pc.addTrack(player.audio)
pc.addTrack(player.video)
async def response_offer(self, data):
pc = self.get_pc(data)
if pc is None:
print(f'{self.peers=}, {data=}')
return
pc.on("connectionstatechange", partial(self.pc_connectionState_changed, data['from'].id))
pc.on('track', partial(self.pc_track, data['from'].id))
pc.on('icecandidate', partial(self.on_icecandidate, pc))
# self.play_video(data)
offer = RTCSessionDescription(** data.offer)
await pc.setRemoteDescription(offer)
answer = await pc.createAnswer()
await pc.setLocalDescription(answer)
await self.ws.send(json.dumps({
'type':'answer',
'answer':{'type':pc.localDescription.type, 'sdp':pc.localDescription.sdp},
'to':data['from']
}))
async def accept_answer(self, data):
pc = self.get_pc(data)
answer = RTCSessionDescription(data.answer);
pc.setRemoteDescription(answer)
async def accept_iceCandidate(self, data):
pc = self.get_pc(data)
candidate = data.candidate
ip = candidate['candidate'].split(' ')[4]
port = candidate['candidate'].split(' ')[5]
protocol = candidate['candidate'].split(' ')[7]
priority = candidate['candidate'].split(' ')[3]
foundation = candidate['candidate'].split(' ')[0]
component = candidate['candidate'].split(' ')[1]
type = candidate['candidate'].split(' ')[7]
rtc_candidate = RTCIceCandidate(
ip=ip,
port=port,
protocol=protocol,
priority=priority,
foundation=foundation,
component=component,
type=type,
sdpMid=candidate['sdpMid'],
sdpMLineIndex=candidate['sdpMLineIndex']
)
await pc.addIceCandidate(rtc_candidate)
print('addIceCandidate ok')
handlers = {
'onlineList':save_onlineList,
'callRequest':auto_accept_call,
'offer':response_offer,
'answer':accept_answer,
'iceCandidate':accept_iceCandidate,
}
async def main():
agent = RTCLLM(ws_url='wss://sage.opencomputing.cn/wss/ws/rtc_signaling.ws',
iceServers=[{
'urls':'stun:stun.opencomputing.cn'},{
'urls':'turn:stun.opencomputing.cn',
'username':'turn',
'credential':'server'
}])
print('running ...')
await agent.run()
if __name__ == '__main__':
asyncio.run(main())