Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: client tests for docker compose server (#1189)

Browse Source
This commit is contained in:
Sarah Wooders
2024-03-26 18:58:00 -07:00
committed by GitHub
parent c8297f53e4
commit c0bd66c957
15 changed files with 77 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
.github/workflows/tests.yml vendored
View File

@@ -1,7 +1,7 @@
name: Run All pytest Tests
env:
PGVECTOR_TEST_DB_URL: ${{ secrets.PGVECTOR_TEST_DB_URL }}
MEMGPT_PGURI: ${{ secrets.MEMGPT_PGURI }}
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
on:
@@ -21,6 +21,7 @@ jobs:
- name: Build and run container
run: bash db/run_postgres.sh
- name: "Setup Python, Poetry and Dependencies"
uses: packetcoders/action-setup-cache-python-poetry@main
with:
@@ -28,9 +29,25 @@ jobs:
poetry-version: "1.7.1"
install-args: "--all-extras"
- name: Initialize credentials
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
run: |
poetry run memgpt quickstart --backend openai
- name: Run docker compose server
env:
MEMGPT_PG_DB: memgpt
MEMGPT_PG_USER: memgpt
MEMGPT_PG_PASSWORD: memgpt
MEMGPT_SERVER_PASS: test_server_token
MEMGPT_CONFIG_PATH: configs/server_config.yaml
run: docker compose up &
- name: Run server tests
env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt
MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token
run: |
@@ -38,7 +55,7 @@ jobs:
- name: Run tests with pytest
env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt
MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token
PYTHONPATH: ${{ github.workspace }}:${{ env.PYTHONPATH }}
@@ -47,7 +64,7 @@ jobs:
- name: Run storage tests
env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt
MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token
run: |

6
compose.yaml
View File

@@ -13,12 +13,16 @@ services:
- "5432:5432"
memgpt_server:
image: memgpt/memgpt-server:0.3.6
image: memgpt/memgpt-server:0.3.7
depends_on:
- pgvector_db
environment:
- POSTGRES_URI=postgresql://${MEMGPT_PG_USER}:${MEMGPT_PG_PASSWORD}@pgvector_db:5432/${MEMGPT_PG_DB}
- MEMGPT_SERVER_PASS=${MEMGPT_SERVER_PASS} # memgpt server password
- MEMGPT_PG_DB=${MEMGPT_PG_DB}
- MEMGPT_PG_USER=${MEMGPT_PG_USER}
- MEMGPT_PG_PASSWORD=${MEMGPT_PG_PASSWORD}
- MEMGPT_PG_URL=pgvector_db
volumes:
- ./configs/server_config.yaml:/root/.memgpt/config # config file
- ~/.memgpt/credentials:/root/.memgpt/credentials # credentials file

2
db/run_postgres.sh
View File

@@ -4,7 +4,7 @@ docker build -f db/Dockerfile.simple -t pg-test .
# run container
docker run -d --rm \
--name memgpt-db-test \
-p 5432:5432 \
-p 8888:5432 \
-e POSTGRES_PASSWORD=password \
-v memgpt_db_test:/var/lib/postgresql/data \
pg-test:latest

4
docs/storage.md
View File

@@ -38,13 +38,13 @@ To run the Postgres backend, you will need a URI to a Postgres database that sup
- Add the following line to your shell profile (e.g., `~/.bashrc`, `~/.zshrc`):
```sh
export PGVECTOR_TEST_DB_URL=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
export MEMGPT_PGURI=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
```
- Or create a `.env` file in the root project directory with:
```sh
PGVECTOR_TEST_DB_URL=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
MEMGPT_PGURI=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
```
4. Run the script from the root project directory:

1
init.sql
View File

@@ -6,7 +6,6 @@
\set db_password `([ -r /var/run/secrets/memgpt-password ] && cat /var/run/secrets/memgpt-password) || echo "${POSTGRES_PASSWORD:-memgpt}"`
\set db_name `([ -r /var/run/secrets/memgpt-db ] && cat /var/run/secrets/memgpt-db) || echo "${POSTGRES_DB:-memgpt}"`
-- CREATE USER :"db_user"
-- WITH PASSWORD :'db_password'
-- NOCREATEDB

7
memgpt/agent_store/db.py
View File

@@ -439,15 +439,18 @@ class PostgresStorageConnector(SQLStorageConnector):
self.db_model = get_db_model(config, self.table_name, table_type, user_id, agent_id)
# construct URI from enviornment variables
if os.getenv("MEMGPT_PGURI"):
self.uri = os.getenv("MEMGPT_PGURI")
else:
db = os.getenv("MEMGPT_PG_DB", "memgpt")
user = os.getenv("MEMGPT_PG_USER", "memgpt")
password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt")
port = os.getenv("MEMGPT_PG_PORT", "5432")
url = os.getenv("MEMGPT_PG_URL", "localhost")
uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}"
self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}"
# create engine
self.engine = create_engine(uri)
self.engine = create_engine(self.uri)
for c in self.db_model.__table__.columns:
if c.name == "embedding":

3
memgpt/metadata.py
View File

@@ -305,6 +305,9 @@ class MetadataStore:
# TODO: get DB URI or path
if config.metadata_storage_type == "postgres":
# construct URI from enviornment variables
if os.getenv("MEMGPT_PGURI"):
self.uri = os.getenv("MEMGPT_PGURI")
else:
db = os.getenv("MEMGPT_PG_DB", "memgpt")
user = os.getenv("MEMGPT_PG_USER", "memgpt")
password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt")

4
scripts/pack_docker.sh Normal file
View File

@@ -0,0 +1,4 @@
export MEMGPT_VERSION=$(memgpt version)
#docker build -t memgpt/memgpt-server:${MEMGPT_VERSION} --platform linux/x86_64 .
docker buildx build --platform=linux/amd64,linux/arm64,linux/x86_64 -t memgpt/memgpt-server:${MEMGPT_VERSION} .
docker push memgpt/memgpt-server:${MEMGPT_VERSION}

2
tests/clear_postgres_db.py
View File

@@ -5,7 +5,7 @@ from sqlalchemy import create_engine, MetaData
def main():
uri = os.environ.get(
"PGVECTOR_TEST_DB_URL",
"MEMGPT_PGURI",
"postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt",
)

4
tests/test_admin_client.py
View File

@@ -32,8 +32,8 @@ def run_server():
load_dotenv()
# Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL")
assert db_url, "Missing PGVECTOR_TEST_DB_URL"
db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig(

12
tests/test_client.py
View File

@@ -43,8 +43,8 @@ def run_server():
load_dotenv()
# Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL")
assert db_url, "Missing PGVECTOR_TEST_DB_URL"
db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig(
@@ -97,6 +97,7 @@ def run_server():
config.save()
credentials.save()
print("Starting server...")
start_server(debug=True)
@@ -104,22 +105,23 @@ def run_server():
@pytest.fixture(
params=[
{"base_url": local_service_url},
# {"base_url": docker_compose_url}, # TODO: add when docker compose added to tests
{"base_url": docker_compose_url}, # TODO: add when docker compose added to tests
# {"base_url": None} # TODO: add when implemented
],
scope="module",
)
# @pytest.fixture(params=[{"base_url": test_base_url}], scope="module")
def client(request):
print("CLIENT", request.param["base_url"])
if request.param["base_url"]:
if request.param["base_url"] == local_service_url:
# start server
print("Starting server...")
print("Starting server thread")
thread = threading.Thread(target=run_server, daemon=True)
thread.start()
time.sleep(5)
admin = Admin(local_service_url, test_server_token)
admin = Admin(request.param["base_url"], test_server_token)
response = admin.create_user(test_user_id) # Adjust as per your client's method
user_id = response.user_id
token = response.api_key

8
tests/test_load_archival.py
View File

@@ -46,10 +46,10 @@ def test_load_directory(
wipe_config()
# setup config
if metadata_storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"):
if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI")
return
TEST_MEMGPT_CONFIG.metadata_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL")
TEST_MEMGPT_CONFIG.metadata_storage_uri = os.getenv("MEMGPT_PGURI")
TEST_MEMGPT_CONFIG.metadata_storage_type = "postgres"
elif metadata_storage_connector == "sqlite":
print("testing sqlite metadata")
@@ -57,10 +57,10 @@ def test_load_directory(
else:
raise NotImplementedError(f"Storage type {metadata_storage_connector} not implemented")
if passage_storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"):
if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI")
return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL")
TEST_MEMGPT_CONFIG.archival_storage_uri = os.getenv("MEMGPT_PGURI")
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
elif passage_storage_connector == "chroma":
print("testing chroma passage storage")

12
tests/test_metadata_store.py
View File

@@ -16,11 +16,11 @@ from memgpt.models.pydantic_models import HumanModel, PersonaModel
@pytest.mark.parametrize("storage_connector", ["sqlite"])
def test_storage(storage_connector):
if storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"):
if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI")
return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"]
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
TEST_MEMGPT_CONFIG.recall_storage_type = "postgres"
if storage_connector == "sqlite":
@@ -49,8 +49,8 @@ def test_storage(storage_connector):
user_id=user_1.id,
name="agent_1",
preset=DEFAULT_PRESET,
persona=get_persona_text(DEFAULT_PERSONA),
human=get_human_text(DEFAULT_HUMAN),
persona=DEFAULT_PERSONA,
human=DEFAULT_HUMAN,
llm_config=TEST_MEMGPT_CONFIG.default_llm_config,
embedding_config=TEST_MEMGPT_CONFIG.default_embedding_config,
)
@@ -73,7 +73,7 @@ def test_storage(storage_connector):
from memgpt.presets.presets import add_default_presets
add_default_presets(user_1.id, ms)
preset_obj = ms.get_preset(preset_name=DEFAULT_PRESET, user_id=user_1.id)
preset_obj = ms.get_preset(name=DEFAULT_PRESET, user_id=user_1.id)
from memgpt.interface import CLIInterface as interface # for printing to terminal
# Overwrite fields in the preset if they were specified

4
tests/test_server.py
View File

@@ -20,8 +20,8 @@ def server():
wipe_memgpt_home()
# Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL")
assert db_url, "Missing PGVECTOR_TEST_DB_URL"
db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig(

6
tests/test_storage.py
View File

@@ -122,11 +122,11 @@ def test_storage(
# del globals()['Message']
if storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"):
if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI")
return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"]
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
TEST_MEMGPT_CONFIG.recall_storage_type = "postgres"
if storage_connector == "lancedb":