From aca67068f24226ed6f484cad5283fd5b15afcf90 Mon Sep 17 00:00:00 2001 From: Harish Mohan Raj Date: Tue, 10 Sep 2024 15:31:06 +0530 Subject: [PATCH] Unique property name (#165) * WIP * WIP * Ensure property name is unique * Fix tests * Add migration script * Polishing * Hardcode the error message --- docs/docs/SUMMARY.md | 1 + .../check_model_name_uniqueness_and_raise.md | 11 ++ fastagency/studio/app.py | 7 ++ fastagency/studio/helpers.py | 17 +++ .../migration.sql | 21 ++++ tests/studio/app/test_model_routes.py | 119 ++++++++++++++++-- 6 files changed, 168 insertions(+), 8 deletions(-) create mode 100644 docs/docs/en/api/fastagency/studio/helpers/check_model_name_uniqueness_and_raise.md create mode 100644 migrations/20240910072037_make_json_str_name_unique_per_user/migration.sql diff --git a/docs/docs/SUMMARY.md b/docs/docs/SUMMARY.md index ba176db0a..563046fda 100644 --- a/docs/docs/SUMMARY.md +++ b/docs/docs/SUMMARY.md @@ -168,6 +168,7 @@ search: - [ping_handler](api/fastagency/studio/faststream_app/ping_handler.md) - helpers - [add_model_to_user](api/fastagency/studio/helpers/add_model_to_user.md) + - [check_model_name_uniqueness_and_raise](api/fastagency/studio/helpers/check_model_name_uniqueness_and_raise.md) - [create_autogen](api/fastagency/studio/helpers/create_autogen.md) - [create_model](api/fastagency/studio/helpers/create_model.md) - [create_model_ref](api/fastagency/studio/helpers/create_model_ref.md) diff --git a/docs/docs/en/api/fastagency/studio/helpers/check_model_name_uniqueness_and_raise.md b/docs/docs/en/api/fastagency/studio/helpers/check_model_name_uniqueness_and_raise.md new file mode 100644 index 000000000..a426e40c5 --- /dev/null +++ b/docs/docs/en/api/fastagency/studio/helpers/check_model_name_uniqueness_and_raise.md @@ -0,0 +1,11 @@ +--- +# 0.5 - API +# 2 - Release +# 3 - Contributing +# 5 - Template Page +# 10 - Default +search: + boost: 0.5 +--- + +::: fastagency.studio.helpers.check_model_name_uniqueness_and_raise diff --git a/fastagency/studio/app.py b/fastagency/studio/app.py index 13989eb6e..5205d5d63 100644 --- a/fastagency/studio/app.py +++ b/fastagency/studio/app.py @@ -24,6 +24,7 @@ from .db.prisma import fastapi_lifespan from .helpers import ( add_model_to_user, + check_model_name_uniqueness_and_raise, create_model, get_all_models_for_user, ) @@ -146,6 +147,8 @@ async def add_model( model: dict[str, Any], background_tasks: BackgroundTasks, ) -> dict[str, Any]: + await check_model_name_uniqueness_and_raise(user_uuid, model["name"]) + return await add_model_to_user( user_uuid=user_uuid, type_name=type_name, @@ -203,6 +206,10 @@ async def update_model( validated_model = registry.validate(type_name, model_name, model) found_model = await DefaultDB.backend().find_model(model_uuid=model_uuid) + found_model_name = found_model["json_str"].get("name") + if model["name"] != found_model_name: + await check_model_name_uniqueness_and_raise(user_uuid, model["name"]) + await DefaultDB.backend().update_model( model_uuid=found_model["uuid"], user_uuid=user_uuid, diff --git a/fastagency/studio/helpers.py b/fastagency/studio/helpers.py index 73f3796a5..a3184d4cb 100644 --- a/fastagency/studio/helpers.py +++ b/fastagency/studio/helpers.py @@ -216,3 +216,20 @@ async def create_autogen( model = await get_model_by_ref(model_ref) return await model.create_autogen(model_id=model_id, user_id=user_id, **kwargs) + + +async def check_model_name_uniqueness_and_raise( + user_uuid: str, model_name: str +) -> None: + existing_models = await DefaultDB.backend().find_many_model(user_uuid=user_uuid) + + if any(model["json_str"].get("name") == model_name for model in existing_models): + raise HTTPException( + status_code=422, + detail=[ + { + "loc": ("name",), + "msg": "Name already exists. Please enter a different name", + } + ], + ) diff --git a/migrations/20240910072037_make_json_str_name_unique_per_user/migration.sql b/migrations/20240910072037_make_json_str_name_unique_per_user/migration.sql new file mode 100644 index 000000000..fc2170f8f --- /dev/null +++ b/migrations/20240910072037_make_json_str_name_unique_per_user/migration.sql @@ -0,0 +1,21 @@ +-- Update records with duplicate names for the same user +UPDATE "Model" +SET json_str = jsonb_set( + json_str::jsonb, + '{name}', + -- Append a random 5-digit suffix to the name + to_jsonb(json_str->>'name' || '_' || LPAD(FLOOR(RANDOM() * 100000)::text, 5, '0')) +) +-- Select records with duplicate names +WHERE (user_uuid, json_str->>'name') IN ( + SELECT user_uuid, json_str->>'name' + FROM "Model" + GROUP BY user_uuid, json_str->>'name' + HAVING COUNT(*) > 1 +) +-- Exclude the first occurrence of each duplicate set +AND uuid NOT IN ( + SELECT DISTINCT ON (user_uuid, json_str->>'name') uuid + FROM "Model" + ORDER BY user_uuid, json_str->>'name', created_at +); diff --git a/tests/studio/app/test_model_routes.py b/tests/studio/app/test_model_routes.py index 6d14fc3f3..8e6c5fc11 100644 --- a/tests/studio/app/test_model_routes.py +++ b/tests/studio/app/test_model_routes.py @@ -1,6 +1,6 @@ import random import uuid -from typing import Optional +from typing import Any, Optional from unittest.mock import AsyncMock, patch import pytest @@ -179,6 +179,106 @@ async def test_add_model(self, user_uuid: str) -> None: actual = response.json() assert actual == expected + @pytest.mark.asyncio + async def test_add_model_with_duplicate_name(self, user_uuid: str) -> None: + model_uuid = str(uuid.uuid4()) + name = f"model_name_{model_uuid}" + azure_oai_api_key = AzureOAIAPIKey(api_key="whatever", name=name) + response = client.post( + f"/user/{user_uuid}/models/secret/AzureOAIAPIKey/{model_uuid}", + json=azure_oai_api_key.model_dump(), + ) + + assert response.status_code == 200 + expected = { + "api_key": "whatever", # pragma: allowlist secret + "name": name, + } + actual = response.json() + assert actual == expected + + response = client.get(f"/user/{user_uuid}/models") + assert response.status_code == 200 + + existing_name = name + new_model_uuid = str(uuid.uuid4()) + new_azure_oai_api_key = AzureOAIAPIKey( + api_key="whatever", # pragma: allowlist secret + name=existing_name, + ) + response = client.post( + f"/user/{user_uuid}/models/secret/AzureOAIAPIKey/{new_model_uuid}", + json=new_azure_oai_api_key.model_dump(), + ) + assert response.status_code == 422 + expected_error_response: dict[str, list[dict[str, Any]]] = { + "detail": [ + { + "loc": ["name"], + "msg": "Name already exists. Please enter a different name", + } + ] + } + actual = response.json() + assert actual == expected_error_response + + @pytest.mark.asyncio + async def test_update_model_with_duplicate_name(self, user_uuid: str) -> None: + models = [ + {"uuid": str(uuid.uuid4()), "name": f"model_name_{i}"} for i in range(2) + ] + # Add two models + for model in models: + model_uuid = model["uuid"] + name = model["name"] + azure_oai_api_key = AzureOAIAPIKey(api_key="whatever", name=name) + response = client.post( + f"/user/{user_uuid}/models/secret/AzureOAIAPIKey/{model_uuid}", + json=azure_oai_api_key.model_dump(), + ) + assert response.status_code == 200 + expected = { + "api_key": "whatever", # pragma: allowlist secret + "name": name, + } + actual = response.json() + assert actual == expected + + # update name of the second model + new_name = f"updated_{models[1]['name']}" + model_uuid = models[1]["uuid"] + updated_model = AzureOAIAPIKey(api_key="new_key", name=new_name) + response = client.put( + f"/user/{user_uuid}/models/secret/AzureOAIAPIKey/{model_uuid}", + json=updated_model.model_dump(), + ) + assert response.status_code == 200 + expected = { + "api_key": "new_key", # pragma: allowlist secret + "name": new_name, + } + actual = response.json() + assert actual == expected + + # Try to update the second model name with the first model name (should fail) + first_model_name = models[0]["name"] + updated_model = AzureOAIAPIKey(api_key="new_key", name=first_model_name) + response = client.put( + f"/user/{user_uuid}/models/secret/AzureOAIAPIKey/{model_uuid}", + json=updated_model.model_dump(), + ) + assert response.status_code == 422 + expected_error_response: dict[str, list[dict[str, Any]]] = { + "detail": [ + { + "loc": ["name"], + "msg": "Name already exists. Please enter a different name", + } + ] + } + actual = response.json() + assert actual == expected_error_response + @pytest.mark.asyncio async def test_add_model_deployment(self, user_uuid: str) -> None: team_uuid = str(uuid.uuid4()) @@ -336,7 +436,8 @@ async def test_update_model( self, user_uuid: str, monkeypatch: pytest.MonkeyPatch ) -> None: key_uuid = str(uuid.uuid4()) - azure_oai_api_key = AzureOAIAPIKey(api_key="who cares", name="whatever") + unique_name = f"unique_name_{key_uuid}" + azure_oai_api_key = AzureOAIAPIKey(api_key="who cares", name=unique_name) type_name = "secret" model_name = "AzureOAIAPIKey" @@ -355,7 +456,7 @@ async def test_update_model( assert response.status_code == 200 expected = { "api_key": "who cares", # pragma: allowlist secret - "name": "whatever", + "name": unique_name, } actual = response.json() assert actual == expected @@ -366,8 +467,9 @@ async def test_update_model_deployment(self, user_uuid: str) -> None: deployment_uuid = str(uuid.uuid4()) gh_token_uuid = str(uuid.uuid4()) fly_token_uuid = str(uuid.uuid4()) + unique_name = f"unique_name_{deployment_uuid}" model = { - "name": "name", + "name": unique_name, "repo_name": "repo_name", "fly_app_name": "Test", "team": {"uuid": team_uuid, "type": "team", "name": "TwoAgentTeam"}, @@ -418,7 +520,7 @@ async def test_update_model_deployment(self, user_uuid: str) -> None: # Update deployment new_gh_token_uuid = str(uuid.uuid4()) model = { - "name": "name", + "name": unique_name, "repo_name": "repo_name", "fly_app_name": "Test", "team": {"uuid": team_uuid, "type": "team", "name": "TwoAgentTeam"}, @@ -439,7 +541,7 @@ async def test_update_model_deployment(self, user_uuid: str) -> None: assert response.status_code == 200 expected = { - "name": "name", + "name": unique_name, "repo_name": "repo_name", "fly_app_name": "Test", "team": { @@ -467,7 +569,8 @@ async def test_delete_model( self, user_uuid: str, monkeypatch: pytest.MonkeyPatch ) -> None: key_uuid = str(uuid.uuid4()) - azure_oai_api_key = AzureOAIAPIKey(api_key="whatever", name="whatever") + unique_name = f"unique_name_{key_uuid}" + azure_oai_api_key = AzureOAIAPIKey(api_key="whatever", name=unique_name) type_name = "secret" model_name = "AzureOAIAPIKey" @@ -483,7 +586,7 @@ async def test_delete_model( assert response.status_code == 200 expected = { "api_key": "whatever", # pragma: allowlist secret - "name": "whatever", + "name": unique_name, } actual = response.json() assert actual == expected