201 lines
6.4 KiB
Python
201 lines
6.4 KiB
Python
#
|
|
# Copyright (c) 2025, Daily
|
|
#
|
|
# SPDX-License-Identifier: BSD 2-Clause License
|
|
#
|
|
import asyncio
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any, Dict, List
|
|
|
|
import aiohttp
|
|
import uvicorn
|
|
from dotenv import load_dotenv
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pipecat.transports.services.helpers.daily_rest import DailyRESTHelper, DailyRoomParams
|
|
|
|
# Load environment variables
|
|
load_dotenv(override=True)
|
|
|
|
NUMBER_OF_ROOMS = 1
|
|
|
|
|
|
class RoomPool:
|
|
"""Manages a pool of pre-created rooms for quick allocation."""
|
|
|
|
def __init__(self, daily_rest_helper: DailyRESTHelper):
|
|
self.daily_rest_helper = daily_rest_helper
|
|
self.pool: List[Dict[str, str]] = []
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def fill_pool(self, count: int):
|
|
"""Fills the pool with `count` new rooms."""
|
|
for _ in range(count):
|
|
await self.add_room()
|
|
|
|
async def add_room(self):
|
|
"""Creates a new room and adds it to the pool."""
|
|
try:
|
|
room = await self.daily_rest_helper.create_room(DailyRoomParams())
|
|
if not room.url:
|
|
raise HTTPException(status_code=500, detail="Failed to create room")
|
|
|
|
user_token = await self.daily_rest_helper.get_token(room.url)
|
|
if not user_token:
|
|
raise HTTPException(status_code=500, detail="Failed to get user token")
|
|
|
|
bot_token = await self.daily_rest_helper.get_token(room.url)
|
|
if not bot_token:
|
|
raise HTTPException(status_code=500, detail="Failed to get bot token")
|
|
|
|
async with self.lock:
|
|
self.pool.append(
|
|
{"room_url": room.url, "user_token": user_token, "bot_token": bot_token}
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Error adding room to pool: {e}")
|
|
|
|
async def get_room(self) -> Dict[str, str]:
|
|
"""Retrieves a room from the pool and requests a new one to maintain the size."""
|
|
async with self.lock:
|
|
if not self.pool:
|
|
raise HTTPException(status_code=503, detail="No available rooms")
|
|
|
|
room = self.pool.pop(0) # Get first available room
|
|
|
|
# Start a background task to replenish the pool
|
|
asyncio.create_task(self.add_room())
|
|
|
|
return room
|
|
|
|
async def delete_room(self, room_url: str):
|
|
"""Deletes a room when it is not needed anymore"""
|
|
await self.daily_rest_helper.delete_room_by_url(room_url)
|
|
|
|
async def cleanup(self):
|
|
for rooms in self.pool:
|
|
room_url = rooms["room_url"]
|
|
await self.delete_room(room_url)
|
|
|
|
|
|
class BotManager:
|
|
"""Manages bot subprocesses asynchronously."""
|
|
|
|
def __init__(self):
|
|
self.bot_procs: Dict[int, asyncio.subprocess.Process] = {}
|
|
self.room_mappings: Dict[int, str] = {} # Maps process ID to room URL
|
|
|
|
async def start_bot(self, room_url: str, token: str) -> int:
|
|
bot_file = "single_bot"
|
|
command = f"python3 -m {bot_file} -u {room_url} -t {token}"
|
|
|
|
try:
|
|
proc = await asyncio.create_subprocess_shell(
|
|
command,
|
|
cwd=os.path.dirname(os.path.abspath(__file__)),
|
|
)
|
|
if proc.pid is None:
|
|
raise HTTPException(status_code=500, detail="Failed to get subprocess PID")
|
|
|
|
self.bot_procs[proc.pid] = proc
|
|
self.room_mappings[proc.pid] = room_url
|
|
# Monitor the process and delete the room when it exits
|
|
asyncio.create_task(self._monitor_process(proc.pid))
|
|
|
|
return proc.pid
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Failed to start subprocess: {e}")
|
|
|
|
async def _monitor_process(self, pid: int):
|
|
"""Monitors a bot process and deletes the associated room when it exits."""
|
|
proc = self.bot_procs.get(pid)
|
|
if proc:
|
|
await proc.wait() # Wait for the process to exit
|
|
room_url = self.room_mappings.pop(pid, None)
|
|
|
|
if room_url:
|
|
await room_pool.delete_room(room_url)
|
|
print(f"Deleted room: {room_url}")
|
|
|
|
del self.bot_procs[pid]
|
|
|
|
async def cleanup(self):
|
|
"""Terminates all running bot processes and deletes associated rooms."""
|
|
for pid, proc in list(self.bot_procs.items()):
|
|
try:
|
|
proc.terminate()
|
|
await asyncio.wait_for(proc.wait(), timeout=5)
|
|
|
|
room_url = self.room_mappings.pop(pid, None)
|
|
if room_url:
|
|
await room_pool.delete_room(room_url) # Delete room when process terminates
|
|
print(f"Deleted room: {room_url}")
|
|
|
|
except asyncio.TimeoutError:
|
|
print(f"Process {pid} did not terminate in time.")
|
|
except Exception as e:
|
|
print(f"Error terminating process {pid}: {e}")
|
|
|
|
# Clear remaining mappings
|
|
self.bot_procs.clear()
|
|
self.room_mappings.clear()
|
|
|
|
|
|
# Global instances
|
|
bot_manager = BotManager()
|
|
room_pool: RoomPool # Will be initialized in lifespan
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Handles FastAPI startup and shutdown."""
|
|
global room_pool
|
|
aiohttp_session = aiohttp.ClientSession()
|
|
daily_rest_helper = DailyRESTHelper(
|
|
daily_api_key=os.getenv("DAILY_API_KEY", ""),
|
|
daily_api_url=os.getenv("DAILY_API_URL", "https://api.daily.co/v1"),
|
|
aiohttp_session=aiohttp_session,
|
|
)
|
|
|
|
room_pool = RoomPool(daily_rest_helper)
|
|
await room_pool.fill_pool(NUMBER_OF_ROOMS) # Fill pool on startup
|
|
|
|
yield # Run app
|
|
|
|
await bot_manager.cleanup()
|
|
await room_pool.cleanup()
|
|
await aiohttp_session.close()
|
|
|
|
|
|
# Initialize FastAPI app with lifespan manager
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
# Configure CORS to allow requests from any origin
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@app.post("/connect")
|
|
async def bot_connect(request: Request) -> Dict[Any, Any]:
|
|
try:
|
|
room = await room_pool.get_room()
|
|
await bot_manager.start_bot(room["room_url"], room["bot_token"])
|
|
except HTTPException as e:
|
|
return {"error": str(e)}
|
|
|
|
return {
|
|
"room_url": room["room_url"],
|
|
"token": room["user_token"],
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=7860)
|