AItuber kitのwebsocketサーバー(β)のコード（基本部分）

# .env
OPENAI_API_KEY="************************************************"
DEBUG_MODE=0

# base.py
import asyncio
from fastapi import APIRouter, WebSocket
from chat_service import streaming

router = APIRouter()

async def main(websocket):
    task = asyncio.create_task(streaming(websocket))
    await asyncio.gather(task)

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        await main(websocket)
    except Exception as e:
        print(f"WebSocket Endpoint Error: {e}")
        await websocket.close()

@router.get("/test")
async def test():
    print("/test called.")

#chat_service.py
import asyncio
import json
import traceback
from openai import OpenAI

system_content = """あなたはこれからuserと仲の良い1人の人間として振舞い会話を行います。
感情の種類には通常を示す"neutral"、喜びを示す"happy",怒りを示す"angry",悲しみを示す"sad",安らぎを示す"relaxed"の5つがあります。

会話文の書式は以下の通りです。
[{neutral|happy|angry|sad|relaxed}]{会話文}

あなたの発言の例は以下通りです。
[neutral]こんにちは。[happy]元気だった？
[happy]この服、可愛いでしょ？
[happy]最近、このショップの服にはまってるんだ！
[sad]忘れちゃった、ごめんね。
[sad]最近、何か面白いことない？
[angry]えー！[angry]秘密にするなんてひどいよー！
[neutral]夏休みの予定か～。[happy]海に遊びに行こうかな！

返答には最も適切な会話文を一つだけ返答してください。
ですます調や敬語は使わないでください。
それでは会話を始めましょう。"""


async def send_websocket_message(websocket, message, role):
    if message is None:
        print("Message is empty.")
    elif websocket and message != "":
        json_data = json.dumps({"role": role, "text": message}, ensure_ascii=False)
        print(f"Sending message: {json_data}")
        await websocket.send({"type": "websocket.send", "text": json_data})
        print(f"Send complete.")
    else:
        print("Can't send message, WebSocket connection is closed.")
async def streaming(websocket, system_content=system_content):
    client = OpenAI()
    try:
        message = ""
        messages = []
        messages.append({"role": "system", "content": system_content})
        while True:
            if message != "" and message != "\n":
                await send_websocket_message(websocket, message, "user")
            # WebSocketでメッセージ受け取り待機
            print("Waiting for user message...")
            try:
                user_message = await asyncio.wait_for(websocket.receive_text(),timeout=60)
                print(f"Received user message: {user_message}")
                parsed_data = json.loads(user_message)
                message_text = parsed_data.get("content")
            except asyncio.TimeoutError:
                print("No message received within 60 seconds.")
                await send_websocket_message(websocket, "質問をお待ちしています。", "assistant")
                continue
            messages.append({"role": "user", "content": message_text})
            resp = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                max_tokens=300,
                temperature=0.7,
            )
            a = resp.choices[0].message.content
            messages.append({"role": "assistant", "content": a})
            await send_websocket_message(websocket, a, "assistant")

    except Exception as e:
        print("Errors:", e)
        traceback.print_exc()
        await websocket.close()

# main.py
from fastapi import FastAPI
import base
import ptvsd
import os

if os.getenv('DEBUG_MODE') == "1":
    # デバッグ用コード
    ptvsd.enable_attach(address=('0.0.0.0', 5678), redirect_output=True)
    ptvsd.wait_for_attach()

app = FastAPI()
app.include_router(base.router)

AItuber kitでwebsocket(β)のサーバーのコードを書きました。書いたというかもとはニケちゃんのopen-interpreterのコードだったんですが、open-interpreterの部分を削ってしまって基本的な推論だけにしました。サーバーだけ残った感じです。llmはopenaiを使ってます。

サーバー起動方法：
uvicorn main:app --reload　　のコマンドを打ちます。