From c0bd66c957f4fbbdf427bd72368a0380bae2df20 Mon Sep 17 00:00:00 2001 From: Sarah Wooders Date: Tue, 26 Mar 2024 18:58:00 -0700 Subject: [PATCH] feat: client tests for `docker compose` server (#1189) --- .github/workflows/tests.yml | 25 +++++++++++++++++++++---- compose.yaml | 6 +++++- db/run_postgres.sh | 2 +- docs/storage.md | 4 ++-- init.sql | 1 - memgpt/agent_store/db.py | 17 ++++++++++------- memgpt/metadata.py | 15 +++++++++------ scripts/pack_docker.sh | 4 ++++ tests/clear_postgres_db.py | 2 +- tests/test_admin_client.py | 4 ++-- tests/test_client.py | 12 +++++++----- tests/test_load_archival.py | 8 ++++---- tests/test_metadata_store.py | 12 ++++++------ tests/test_server.py | 4 ++-- tests/test_storage.py | 6 +++--- 15 files changed, 77 insertions(+), 45 deletions(-) create mode 100644 scripts/pack_docker.sh diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 7bbbf23f..cfb97663 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -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: | diff --git a/compose.yaml b/compose.yaml index 8aa4298a..b334b304 100644 --- a/compose.yaml +++ b/compose.yaml @@ -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 diff --git a/db/run_postgres.sh b/db/run_postgres.sh index 77477741..65021938 100644 --- a/db/run_postgres.sh +++ b/db/run_postgres.sh @@ -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 diff --git a/docs/storage.md b/docs/storage.md index 69f75607..0eae7026 100644 --- a/docs/storage.md +++ b/docs/storage.md @@ -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: diff --git a/init.sql b/init.sql index 04758779..c1244ff0 100644 --- a/init.sql +++ b/init.sql @@ -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 diff --git a/memgpt/agent_store/db.py b/memgpt/agent_store/db.py index 0826c14d..cafdd916 100644 --- a/memgpt/agent_store/db.py +++ b/memgpt/agent_store/db.py @@ -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 - 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}" + 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") + 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": diff --git a/memgpt/metadata.py b/memgpt/metadata.py index 25f0ca4e..77e36268 100644 --- a/memgpt/metadata.py +++ b/memgpt/metadata.py @@ -305,12 +305,15 @@ class MetadataStore: # TODO: get DB URI or path if config.metadata_storage_type == "postgres": # construct URI from enviornment variables - 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") - self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}" + 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") + self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}" elif config.metadata_storage_type == "sqlite": path = os.path.join(config.metadata_storage_path, "sqlite.db") self.uri = f"sqlite:///{path}" diff --git a/scripts/pack_docker.sh b/scripts/pack_docker.sh new file mode 100644 index 00000000..348f1474 --- /dev/null +++ b/scripts/pack_docker.sh @@ -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} diff --git a/tests/clear_postgres_db.py b/tests/clear_postgres_db.py index 457beea7..7d10afe3 100644 --- a/tests/clear_postgres_db.py +++ b/tests/clear_postgres_db.py @@ -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", ) diff --git a/tests/test_admin_client.py b/tests/test_admin_client.py index ab1ac3a5..a75e8099 100644 --- a/tests/test_admin_client.py +++ b/tests/test_admin_client.py @@ -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( diff --git a/tests/test_client.py b/tests/test_client.py index 9b3764c9..4b87da30 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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 diff --git a/tests/test_load_archival.py b/tests/test_load_archival.py index bb96162a..7fa56cdc 100644 --- a/tests/test_load_archival.py +++ b/tests/test_load_archival.py @@ -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") diff --git a/tests/test_metadata_store.py b/tests/test_metadata_store.py index 1e051dff..35a0d65b 100644 --- a/tests/test_metadata_store.py +++ b/tests/test_metadata_store.py @@ -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 diff --git a/tests/test_server.py b/tests/test_server.py index 1f441e33..4704849d 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -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( diff --git a/tests/test_storage.py b/tests/test_storage.py index 79c24670..d592792c 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -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":