109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
import pytest
|
|
from unittest.mock import Mock, AsyncMock, MagicMock
|
|
|
|
from memgpt.config import MemGPTConfig, AgentConfig
|
|
from memgpt.server.websocket_interface import SyncWebSocketInterface
|
|
import memgpt.presets as presets
|
|
import memgpt.utils as utils
|
|
import memgpt.system as system
|
|
from memgpt.persistence_manager import LocalStateManager
|
|
|
|
|
|
# def test_websockets():
|
|
# # Create the websocket interface
|
|
# ws_interface = WebSocketInterface()
|
|
|
|
# # Create a dummy persistence manager
|
|
# persistence_manager = InMemoryStateManager()
|
|
|
|
# # Create an agent and hook it up to the WebSocket interface
|
|
# memgpt_agent = presets.use_preset(
|
|
# presets.DEFAULT_PRESET,
|
|
# None, # no agent config to provide
|
|
# "gpt-4-1106-preview",
|
|
# personas.get_persona_text("sam_pov"),
|
|
# humans.get_human_text("chad"),
|
|
# ws_interface,
|
|
# persistence_manager,
|
|
# )
|
|
|
|
# user_message = system.package_user_message("Hello, is anyone there?")
|
|
|
|
# # This should trigger calls to interface user_message and others
|
|
# memgpt_agent.step(user_message=user_message)
|
|
|
|
# # This should trigger the web socket to send over a
|
|
# ws_interface.print_messages(memgpt_agent.messages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dummy():
|
|
assert True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_websockets():
|
|
# Mock a WebSocket connection
|
|
mock_websocket = AsyncMock()
|
|
# mock_websocket = Mock()
|
|
|
|
# Create the WebSocket interface with the mocked WebSocket
|
|
ws_interface = SyncWebSocketInterface()
|
|
|
|
# Register the mock websocket as a client
|
|
ws_interface.register_client(mock_websocket)
|
|
|
|
# Create an agent and hook it up to the WebSocket interface
|
|
config = MemGPTConfig()
|
|
|
|
# Mock the persistence manager
|
|
# create agents with defaults
|
|
agent_config = AgentConfig(persona="sam_pov", human="basic", model="gpt-4-1106-preview")
|
|
persistence_manager = LocalStateManager(agent_config=agent_config)
|
|
|
|
memgpt_agent = presets.use_preset(
|
|
presets.DEFAULT_PRESET,
|
|
config, # no agent config to provide
|
|
"gpt-4-1106-preview",
|
|
utils.get_persona_text("sam_pov"),
|
|
utils.get_human_text("basic"),
|
|
ws_interface,
|
|
persistence_manager,
|
|
)
|
|
|
|
# Mock the user message packaging
|
|
user_message = system.package_user_message("Hello, is anyone there?")
|
|
|
|
# Mock the agent's step method
|
|
# agent_step = AsyncMock()
|
|
# memgpt_agent.step = agent_step
|
|
|
|
# Call the step method, which should trigger interface methods
|
|
ret = memgpt_agent.step(user_message=user_message, first_message=True, skip_verify=True)
|
|
print("ret\n")
|
|
print(ret)
|
|
|
|
# Print what the WebSocket received
|
|
print("client\n")
|
|
for call in mock_websocket.send.mock_calls:
|
|
# print(call)
|
|
_, args, kwargs = call
|
|
# args will be a tuple of positional arguments sent to the send method
|
|
# kwargs will be a dictionary of keyword arguments sent to the send method
|
|
print(f"Sent data: {args[0] if args else None}")
|
|
# If you're using keyword arguments, you can print them out as well:
|
|
# print(f"Sent data with kwargs: {kwargs}")
|
|
|
|
# This is required for the Sync wrapper version
|
|
ws_interface.close()
|
|
|
|
# Assertions to ensure the step method was called
|
|
# agent_step.assert_called_once()
|
|
|
|
# Assertions to ensure the WebSocket interface methods are called
|
|
# You would need to implement the logic to verify that methods like ws_interface.user_message are called
|
|
# This will require you to have some mechanism within your WebSocketInterface to track these calls
|
|
|
|
|
|
# await test_websockets()
|