Skip to content

Commit

Permalink
fix: fix DB session management to avoid connection overflow error (#1758
Browse files Browse the repository at this point in the history
)
  • Loading branch information
sarahwooders authored Sep 18, 2024
1 parent ca11164 commit 45a1101
Show file tree
Hide file tree
Showing 5 changed files with 273 additions and 240 deletions.
3 changes: 3 additions & 0 deletions configs/server_config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -32,5 +32,8 @@ type = postgres
path = /root/.memgpt
uri = postgresql+pg8000://memgpt:memgpt@pgvector_db:5432/memgpt

[version]
memgpt_version = 0.4.0

[client]
anon_clientid = 00000000-0000-0000-0000-000000000000
12 changes: 7 additions & 5 deletions locust_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,18 +21,20 @@ def on_start(self):
# Create a user and get the token
self.client.headers = {"Authorization": "Bearer password"}
user_data = {"name": f"User-{''.join(random.choices(string.ascii_lowercase + string.digits, k=8))}"}
response = self.client.post("/admin/users", json=user_data)
response = self.client.post("/v1/admin/users", json=user_data)
response_json = response.json()
print(response_json)
self.user_id = response_json["id"]

# create a token
response = self.client.post("/admin/users/keys", json={"user_id": self.user_id})
response = self.client.post("/v1/admin/users/keys", json={"user_id": self.user_id})
self.token = response.json()["key"]

# reset to use user token as headers
self.client.headers = {"Authorization": f"Bearer {self.token}"}

# @task(1)
# def create_agent(self):
# generate random name
name = "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
request = CreateAgent(
Expand All @@ -42,7 +44,7 @@ def on_start(self):
)

# create an agent
with self.client.post("/api/agents", json=request.model_dump(), headers=self.client.headers, catch_response=True) as response:
with self.client.post("/v1/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}")

Expand All @@ -57,10 +59,10 @@ def send_message(self):
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
f"/v1/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}")
response.failure(f"Failed to send message {response.status_code}: {response.text}")

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

0 comments on commit 45a1101

Please sign in to comment.