rtcllm/rtcllm/rtc.py

222 lines
6.0 KiB
Python
Raw Normal View History

2024-08-27 18:39:25 +08:00
import asyncio
2024-09-02 18:43:58 +08:00
import random
2024-08-28 18:26:47 +08:00
import json
from functools import partial
2024-08-27 18:39:25 +08:00
from appPublic.dictObject import DictObject
from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription, RTCIceCandidate
2024-09-03 15:45:34 +08:00
from aiortc.sdp import candidate_from_sdp, candidate_to_sdp
2024-08-26 18:11:08 +08:00
from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder, MediaRelay
2024-08-27 18:39:25 +08:00
# from websockets.asyncio.client import connect
from websockets.asyncio.client import connect
import websockets
2024-09-02 18:43:58 +08:00
from stt import asr
2024-08-28 18:26:47 +08:00
from vad import AudioTrackVad
2024-09-02 18:43:58 +08:00
from aav import MyAudioTrack, MyVideoTrack
videos = ['./1.mp4', './2.mp4']
2024-08-28 18:26:47 +08:00
2024-08-27 18:39:25 +08:00
class RTCLLM:
def __init__(self, ws_url, iceServers):
self.ws_url = ws_url
self.iceServers = iceServers
self.peers = DictObject()
2024-09-02 18:43:58 +08:00
self.vid = 0
2024-08-28 18:26:47 +08:00
self.info = {
'id':'rtcllm_agent',
'name':'rtcllm_agent'
}
self.dc = None
2024-09-02 18:43:58 +08:00
self.loop = asyncio.get_event_loop()
2024-08-27 18:39:25 +08:00
2024-09-03 15:45:34 +08:00
async def ws_send(self, s):
self.ws.send(s)
2024-08-27 18:39:25 +08:00
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:
2024-08-28 18:26:47 +08:00
await self.login()
2024-08-27 18:39:25 +08:00
while True:
msg = await self.ws.recv()
data = DictObject(**json.loads(msg))
2024-08-28 18:26:47 +08:00
d = data.data
print(f'ws recv(): {d.type=}')
func = self.handlers.get(d.type)
if func:
f = partial(func, self)
await f(d)
2024-08-27 18:39:25 +08:00
self.ws.close()
async def login(self):
2024-09-03 15:45:34 +08:00
await self.ws_send(json.dumps({
2024-08-27 18:39:25 +08:00
'type':'login',
2024-08-28 18:26:47 +08:00
'info':self.info
2024-08-27 18:39:25 +08:00
}))
2024-08-28 18:26:47 +08:00
async def on_icecandidate(self, pc, candidate):
2024-09-03 15:45:34 +08:00
print('******************** on_icecandidate()', self, pc, candidate)
2024-08-28 18:26:47 +08:00
if candidate:
2024-09-03 15:45:34 +08:00
await self.ws_send(json.dumps({
2024-08-28 18:26:47 +08:00
"type":"candidate",
"candidate":candidate,
}))
2024-08-27 18:39:25 +08:00
async def save_onlineList(self, data):
2024-08-28 18:26:47 +08:00
print(f'{self}, {type(self)}')
2024-08-27 18:39:25 +08:00
self.onlineList = data.onlineList
2024-09-02 18:43:58 +08:00
async def vad_voiceend(self, peer, audio):
txt = await asr(audio)
# await peer.dc.send(txt)
2024-08-28 18:26:47 +08:00
2024-08-27 18:39:25 +08:00
async def auto_accept_call(self, data):
2024-08-28 18:26:47 +08:00
opts = DictObject(iceServers=self.iceServers)
pc = RTCPeerConnection(opts)
2024-08-27 18:39:25 +08:00
self.peers[data['from'].id] = DictObject(**{
'info':data['from'],
'pc':pc
})
2024-09-03 15:45:34 +08:00
await self.ws_send(json.dumps({'type':'callAccepted', 'to':data['from']}))
2024-08-27 18:39:25 +08:00
2024-08-28 18:26:47 +08:00
async def pc_track(self, peerid, track):
peer = self.peers[peerid]
pc = peer.pc
if track.kind == 'audio':
2024-09-02 18:43:58 +08:00
f = partial(self.vad_voiceend, peer)
2024-09-03 15:45:34 +08:00
vadtrack = AudioTrackVad(track, stage=0, onvoiceend=f)
2024-08-28 18:26:47 +08:00
peer.vadtrack = vadtrack
vadtrack.start_vad()
2024-09-02 18:43:58 +08:00
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)
2024-08-28 18:26:47 +08:00
async def pc_connectionState_changed(self, peerid):
peer = self.peers[peerid]
pc = peer.pc
2024-09-03 15:45:34 +08:00
print('conn_state={pc.connectionState} ...........')
2024-09-02 18:43:58 +08:00
peer.audiotrack = MyAudioTrack()
peer.videotrack = MyVideoTrack()
pc.addTrack(peer.audiotrack)
pc.addTrack(peer.videotrack)
self.play_random_media(peer.videotrack, peer.audiotrack)
2024-08-28 18:26:47 +08:00
if pc.connectionState == 'connected':
2024-09-02 10:36:07 +08:00
peer.dc = pc.createDataChannel(peer.info.name)
2024-08-28 18:26:47 +08:00
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
2024-09-02 18:43:58 +08:00
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)
2024-08-27 18:39:25 +08:00
async def response_offer(self, data):
pc = self.get_pc(data)
2024-08-28 18:26:47 +08:00
if pc is None:
print(f'{self.peers=}, {data=}')
return
2024-09-02 10:36:07 +08:00
pc.on("connectionstatechange", partial(self.pc_connectionState_changed, data['from'].id))
2024-08-28 18:26:47 +08:00
pc.on('track', partial(self.pc_track, data['from'].id))
pc.on('icecandidate', partial(self.on_icecandidate, pc))
2024-09-02 18:43:58 +08:00
# self.play_video(data)
2024-09-02 10:36:07 +08:00
2024-08-28 18:26:47 +08:00
offer = RTCSessionDescription(** data.offer)
2024-08-27 18:39:25 +08:00
await pc.setRemoteDescription(offer)
answer = await pc.createAnswer()
await pc.setLocalDescription(answer)
2024-09-03 15:45:34 +08:00
await self.ws_send(json.dumps({
2024-08-27 18:39:25 +08:00
'type':'answer',
2024-08-28 18:26:47 +08:00
'answer':{'type':pc.localDescription.type, 'sdp':pc.localDescription.sdp},
2024-08-27 18:39:25 +08:00
'to':data['from']
}))
2024-09-03 15:45:34 +08:00
offer = await pc.createOffer()
await pc.setLocalDescription(offer)
await self.ws_send(json.dumps({
'type':'offer',
'offer': {'type':pc.localDescription.type, 'sdp':pc.localDescription.dsp},
'to':data['from']
}))
2024-08-27 18:39:25 +08:00
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)
2024-08-28 18:26:47 +08:00
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]
2024-09-03 15:45:34 +08:00
"""
2024-08-28 18:26:47 +08:00
rtc_candidate = RTCIceCandidate(
ip=ip,
port=port,
protocol=protocol,
priority=priority,
foundation=foundation,
component=component,
type=type,
sdpMid=candidate['sdpMid'],
sdpMLineIndex=candidate['sdpMLineIndex']
)
2024-09-03 15:45:34 +08:00
"""
rtc_candidate = candidate_from_sdp(candidate['candidate'].split(":", 1)[1])
rtc_candidate.sdpMid = candidate['sdpMid']
rtc_candidate.sdpMLineIndex = candidate['sdpMLineIndex']
2024-08-28 18:26:47 +08:00
await pc.addIceCandidate(rtc_candidate)
print('addIceCandidate ok')
2024-08-27 18:39:25 +08:00
handlers = {
'onlineList':save_onlineList,
'callRequest':auto_accept_call,
'offer':response_offer,
'answer':accept_answer,
'iceCandidate':accept_iceCandidate,
}
async def main():
2024-08-28 18:26:47 +08:00
agent = RTCLLM(ws_url='wss://sage.opencomputing.cn/wss/ws/rtc_signaling.ws',
2024-08-27 18:39:25 +08:00
iceServers=[{
'urls':'stun:stun.opencomputing.cn'},{
'urls':'turn:stun.opencomputing.cn',
'username':'turn',
'credential':'server'
}])
2024-08-28 18:26:47 +08:00
print('running ...')
2024-08-27 18:39:25 +08:00
await agent.run()
if __name__ == '__main__':
asyncio.run(main())