mirror of
https://github.com/cpacker/MemGPT.git
synced 2025-06-03 04:30:22 +00:00
85faf5f474
Co-authored-by: Charles Packer <packercharles@gmail.com> Co-authored-by: Shubham Naik <shubham.naik10@gmail.com> Co-authored-by: Shubham Naik <shub@memgpt.ai>
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
import websockets
|
|
|
|
from letta.server.constants import WS_DEFAULT_PORT
|
|
from letta.server.ws_api.server import WebSocketServer
|
|
from letta.utils import json_dumps
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dummy():
|
|
assert True
|
|
|
|
|
|
@pytest.mark.skip(reason="websockets is temporarily unsupported in 0.2.12")
|
|
@pytest.mark.asyncio
|
|
async def test_websocket_server():
|
|
# host = "127.0.0.1"
|
|
host = "localhost"
|
|
server = WebSocketServer(host=host)
|
|
server_task = asyncio.create_task(server.run()) # Create a task for the server
|
|
|
|
# the agent config we want to ask the server to instantiate with
|
|
# test_config = AgentConfig(
|
|
# persona="sam_pov",
|
|
# human="cs_phd",
|
|
# preset="memgpt_chat",
|
|
# model_endpoint=
|
|
# )
|
|
test_config = {}
|
|
|
|
uri = f"ws://{host}:{WS_DEFAULT_PORT}"
|
|
try:
|
|
async with websockets.connect(uri) as websocket:
|
|
# Initialize the server with a test config
|
|
print("Sending config to server...")
|
|
await websocket.send(json_dumps({"type": "initialize", "config": test_config}))
|
|
# Wait for the response
|
|
response = await websocket.recv()
|
|
print(f"Response from the agent: {response}")
|
|
|
|
await asyncio.sleep(1) # just in case
|
|
|
|
# Send a message to the agent
|
|
print("Sending message to server...")
|
|
await websocket.send(json_dumps({"type": "message", "content": "Hello, Agent!"}))
|
|
# Wait for the response
|
|
# NOTE: we should be waiting for multiple responses
|
|
response = await websocket.recv()
|
|
print(f"Response from the agent: {response}")
|
|
except (OSError, ConnectionRefusedError) as e:
|
|
print(f"Was unable to connect: {e}")
|
|
finally:
|
|
server_task.cancel() # Cancel the server task after the test
|