99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
import json
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
|
|
from custom.beaver_llm import BeaverLLM, latest_user_text
|
|
from livekit.agents import ChatContext
|
|
|
|
|
|
def test_latest_user_text_uses_most_recent_user_message() -> None:
|
|
ctx = ChatContext.empty()
|
|
ctx.add_message(role="user", content="first")
|
|
ctx.add_message(role="assistant", content="ignored")
|
|
ctx.add_message(role="user", content=["second", "line"])
|
|
|
|
assert latest_user_text(ctx) == "second\nline"
|
|
|
|
|
|
async def test_beaver_llm_sends_latest_user_text_and_returns_reply(
|
|
unused_tcp_port: int,
|
|
) -> None:
|
|
received: list[dict[str, object]] = []
|
|
|
|
async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
|
|
async for message in ws:
|
|
assert message.type == aiohttp.WSMsgType.TEXT
|
|
frame = json.loads(message.data)
|
|
received.append(frame)
|
|
|
|
if frame["type"] == "connect":
|
|
await ws.send_json(
|
|
{
|
|
"type": "connected",
|
|
"channel_id": "terminal-dev",
|
|
"session_id": "terminal-dev:local:livekit-room",
|
|
}
|
|
)
|
|
elif frame["type"] == "message":
|
|
await ws.send_json(
|
|
{
|
|
"type": "ack",
|
|
"message_id": frame["message_id"],
|
|
"session_id": "terminal-dev:local:livekit-room",
|
|
"accepted": True,
|
|
}
|
|
)
|
|
await ws.send_json(
|
|
{
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"message_id": frame["message_id"],
|
|
"run_id": "run-1",
|
|
"text": "beaver reply",
|
|
"finish_reason": "stop",
|
|
}
|
|
)
|
|
|
|
return ws
|
|
|
|
app = web.Application()
|
|
app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
|
|
runner = web.AppRunner(app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
|
|
await site.start()
|
|
|
|
beaver_llm = BeaverLLM(
|
|
url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
|
|
peer_id="livekit-room",
|
|
device_name="livekit-custom-agent",
|
|
)
|
|
ctx = ChatContext.empty()
|
|
ctx.add_message(role="system", content="ignored instructions")
|
|
ctx.add_message(role="user", content="hello beaver")
|
|
|
|
try:
|
|
collected = await beaver_llm.chat(chat_ctx=ctx).collect()
|
|
finally:
|
|
await beaver_llm.aclose()
|
|
await runner.cleanup()
|
|
|
|
assert collected.text == "beaver reply"
|
|
assert received == [
|
|
{
|
|
"type": "connect",
|
|
"peer_id": "livekit-room",
|
|
"device_name": "livekit-custom-agent",
|
|
"capabilities": ["text"],
|
|
},
|
|
{
|
|
"type": "message",
|
|
"message_id": "livekit-room-000001",
|
|
"text": "hello beaver",
|
|
},
|
|
]
|