Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@sarahwooders
sarahwooders committed Sep 10, 2024
1 parent 95fc88c commit d463ce5
Showing 1 changed file with 72 additions and 36 deletions.
108 changes: 72 additions & 36 deletions locust_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,13 @@

from locust import HttpUser, between, task

from memgpt.client.client import RESTClient
from memgpt.constants import BASE_TOOLS, DEFAULT_HUMAN, DEFAULT_PERSONA
from memgpt.schemas.agent import AgentState, CreateAgent
from memgpt.schemas.memgpt_request import MemGPTRequest
from memgpt.schemas.memgpt_response import MemGPTResponse
from memgpt.schemas.memory import ChatMemory
from memgpt.schemas.message import MessageCreate, MessageRole
from memgpt.utils import get_human_text, get_persona_text


class MemGPTUser(HttpUser):
Expand All @@ -27,44 +33,74 @@ def on_start(self):
# reset to use user token as headers
self.client.headers = {"Authorization": f"Bearer {self.token}"}

## Create an agent for this user
# agent_data = {
# "name": f"Agent-{self.token[:8]}",
# "tools": BASE_TOOLS
# }
# response = self.client.post("/agents", json=agent_data)
# self.agent_id = response.json()["id"]

print("HOST", self.host)
client = RESTClient(token=self.token, base_url=self.host)
agent_state = client.create_agent(name=f"Agent-{self.token[:8]}")
self.agent_id = agent_state.id
self.memgpt_client = client
print("Created agent", self.agent_id)
# generate random name
name = "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
request = CreateAgent(

Check failure on line 38 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Arguments missing for parameters "description", "metadata_", "user_id", "message_ids", "system", "llm_config", "embedding_config" (reportCallIssue)
name=f"Agent-{name}",
tools=BASE_TOOLS,
memory=ChatMemory(human=get_human_text(DEFAULT_HUMAN), persona=get_persona_text(DEFAULT_PERSONA)),

Check failure on line 41 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Argument of type "str | None" cannot be assigned to parameter "human" of type "str" in function "__init__"
  Type "str | None" is not assignable to type "str"
    "None" is not assignable to "str" (reportArgumentType)
)

# create an agent
with self.client.post("/api/agents", json=request.model_dump(), headers=self.client.headers, catch_response=True) as response:
if response.status_code != 200:
response.failure(f"Failed to create agent: {response.text}")

Check failure on line 47 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Cannot access attribute "failure" for class "Response"
  Attribute "failure" is unknown (reportAttributeAccessIssue)

Check failure on line 47 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Cannot access attribute "failure" for class "LocustResponse"
  Attribute "failure" is unknown (reportAttributeAccessIssue)

response_json = response.json()
agent_state = AgentState(**response_json)
self.agent_id = agent_state.id
print("Created agent", self.agent_id, agent_state.name)

@task(1)
def send_message(self):
try:
response = self.memgpt_client.send_message(message="Hello, world!", agent_id=self.agent_id, role="user")
except Exception as e:
with self.client.get("/", catch_response=True) as response:
response.failure(str(e))

@task(2)
def get_agent_state(self):
try:
agent_state = self.memgpt_client.get_agent(agent_id=self.agent_id)
except Exception as e:
with self.client.get("/", catch_response=True) as response:
response.failure(str(e))

@task(3)
def get_agent_memory(self):
try:
memory = self.memgpt_client.get_in_context_memory(agent_id=self.agent_id)
except Exception as e:
with self.client.get("/", catch_response=True) as response:
response.failure(str(e))
messages = [MessageCreate(role=MessageRole("user"), text="hello")]

Check failure on line 56 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Argument missing for parameter "name" (reportCallIssue)
request = MemGPTRequest(messages=messages, stream_steps=False, stream_tokens=False, return_message_object=False)

with self.client.post(
f"/api/agents/{self.agent_id}/messages", json=request.model_dump(), headers=self.client.headers, catch_response=True
) as response:
if response.status_code != 200:
response.failure(f"Failed to send message: {response.text}")

Check failure on line 63 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Cannot access attribute "failure" for class "Response"
  Attribute "failure" is unknown (reportAttributeAccessIssue)

Check failure on line 63 in locust_test.py

View workflow job for this annotation

GitHub Actions / Pyright types check (3.11) 

Cannot access attribute "failure" for class "LocustResponse"
  Attribute "failure" is unknown (reportAttributeAccessIssue)

response = MemGPTResponse(**response.json())
print("Response", response.usage)

# @task(1)
# def send_message_stream(self):

# messages = [MessageCreate(role=MessageRole("user"), text="hello")]
# request = MemGPTRequest(messages=messages, stream_steps=True, stream_tokens=True, return_message_object=True)
# if stream_tokens or stream_steps:
# from memgpt.client.streaming import _sse_post

# request.return_message_object = False
# return _sse_post(f"{self.base_url}/api/agents/{agent_id}/messages", request.model_dump(), self.headers)
# else:
# response = requests.post(f"{self.base_url}/api/agents/{agent_id}/messages", json=request.model_dump(), headers=self.headers)
# if response.status_code != 200:
# raise ValueError(f"Failed to send message: {response.text}")
# return MemGPTResponse(**response.json())
# try:
# response = self.memgpt_client.send_message(message="Hello, world!", agent_id=self.agent_id, role="user")
# except Exception as e:
# with self.client.get("/", catch_response=True) as response:
# response.failure(str(e))

# @task(2)
# def get_agent_state(self):
# try:
# agent_state = self.memgpt_client.get_agent(agent_id=self.agent_id)
# except Exception as e:
# with self.client.get("/", catch_response=True) as response:
# response.failure(str(e))

# @task(3)
# def get_agent_memory(self):
# try:
# memory = self.memgpt_client.get_in_context_memory(agent_id=self.agent_id)
# except Exception as e:
# with self.client.get("/", catch_response=True) as response:
# response.failure(str(e))


# class AdminUser(HttpUser):
Expand Down

0 comments on commit d463ce5

Please sign in to comment.