From 759187354633f37085badcd224e76cd795fb2af2 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 6 Dec 2024 17:15:11 -0800 Subject: [PATCH] migrate Organization data model from skyvern/forge/sdk/models.py to skyvern/forge/sdk/schemas/organizations.py (#1343) --- scripts/create_api_key.py | 2 +- skyvern/forge/agent.py | 3 +- skyvern/forge/agent_functions.py | 3 +- skyvern/forge/app.py | 2 +- .../core/permissions/permission_checkers.py | 2 +- skyvern/forge/sdk/db/client.py | 3 +- skyvern/forge/sdk/db/utils.py | 3 +- skyvern/forge/sdk/models.py | 29 +---------------- skyvern/forge/sdk/routes/agent_protocol.py | 3 +- skyvern/forge/sdk/schemas/organizations.py | 32 +++++++++++++++++-- .../forge/sdk/services/org_auth_service.py | 3 +- .../sdk/services/org_auth_token_service.py | 2 +- skyvern/forge/sdk/workflow/context_manager.py | 2 +- skyvern/forge/sdk/workflow/service.py | 3 +- 14 files changed, 50 insertions(+), 42 deletions(-) diff --git a/scripts/create_api_key.py b/scripts/create_api_key.py index 17eafb6de..783ada993 100644 --- a/scripts/create_api_key.py +++ b/scripts/create_api_key.py @@ -5,7 +5,7 @@ from skyvern.forge.app import DATABASE from skyvern.forge.sdk.core import security -from skyvern.forge.sdk.models import OrganizationAuthToken, OrganizationAuthTokenType +from skyvern.forge.sdk.schemas.organizations import OrganizationAuthToken, OrganizationAuthTokenType API_KEY_LIFETIME = timedelta(weeks=5200) diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 2da5b0a65..aa9d5c5b0 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -50,7 +50,8 @@ from skyvern.forge.sdk.core.security import generate_skyvern_signature from skyvern.forge.sdk.core.validators import prepend_scheme_and_validate_url from skyvern.forge.sdk.db.enums import TaskType -from skyvern.forge.sdk.models import Organization, Step, StepStatus +from skyvern.forge.sdk.models import Step, StepStatus +from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import Task, TaskRequest, TaskResponse, TaskStatus from skyvern.forge.sdk.workflow.context_manager import WorkflowRunContext from skyvern.forge.sdk.workflow.models.block import ActionBlock, BaseTaskBlock, ValidationBlock diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py index 0b7a329a4..5e2f1e1d0 100644 --- a/skyvern/forge/agent_functions.py +++ b/skyvern/forge/agent_functions.py @@ -14,7 +14,8 @@ from skyvern.forge.async_operations import AsyncOperation from skyvern.forge.prompts import prompt_engine from skyvern.forge.sdk.api.llm.exceptions import LLMProviderError -from skyvern.forge.sdk.models import Organization, Step, StepStatus +from skyvern.forge.sdk.models import Step, StepStatus +from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus from skyvern.webeye.browser_factory import BrowserState from skyvern.webeye.scraper.scraper import ELEMENT_NODE_ATTRIBUTES, CleanupElementTreeFunc, json_to_html diff --git a/skyvern/forge/app.py b/skyvern/forge/app.py index 630aa700d..49f5e45a3 100644 --- a/skyvern/forge/app.py +++ b/skyvern/forge/app.py @@ -11,7 +11,7 @@ from skyvern.forge.sdk.cache.factory import CacheFactory from skyvern.forge.sdk.db.client import AgentDB from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider -from skyvern.forge.sdk.models import Organization +from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.forge.sdk.workflow.context_manager import WorkflowContextManager from skyvern.forge.sdk.workflow.service import WorkflowService diff --git a/skyvern/forge/sdk/core/permissions/permission_checkers.py b/skyvern/forge/sdk/core/permissions/permission_checkers.py index 2049ee3ad..01b9a7481 100644 --- a/skyvern/forge/sdk/core/permissions/permission_checkers.py +++ b/skyvern/forge/sdk/core/permissions/permission_checkers.py @@ -1,6 +1,6 @@ import abc -from skyvern.forge.sdk.models import Organization +from skyvern.forge.sdk.schemas.organizations import Organization class PermissionChecker(abc.ABC): diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 7cad9dce7..7d4dfeed3 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -49,7 +49,8 @@ convert_to_workflow_run_output_parameter, convert_to_workflow_run_parameter, ) -from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus +from skyvern.forge.sdk.models import Step, StepStatus +from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken from skyvern.forge.sdk.schemas.task_generations import TaskGeneration from skyvern.forge.sdk.schemas.tasks import OrderBy, ProxyLocation, SortDirection, Task, TaskStatus from skyvern.forge.sdk.schemas.totp_codes import TOTPCode diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index 7c11c26ed..4e182e295 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -22,7 +22,8 @@ WorkflowRunOutputParameterModel, WorkflowRunParameterModel, ) -from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus +from skyvern.forge.sdk.models import Step, StepStatus +from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus from skyvern.forge.sdk.workflow.models.parameter import ( AWSSecretParameter, diff --git a/skyvern/forge/sdk/models.py b/skyvern/forge/sdk/models.py index e40b0fcc6..642ad08f9 100644 --- a/skyvern/forge/sdk/models.py +++ b/skyvern/forge/sdk/models.py @@ -3,9 +3,8 @@ from datetime import datetime from enum import StrEnum -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel -from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType from skyvern.webeye.actions.actions import ActionType from skyvern.webeye.actions.models import AgentStepOutput @@ -130,32 +129,6 @@ def is_terminated(self) -> bool: return False -class Organization(BaseModel): - model_config = ConfigDict(from_attributes=True) - - organization_id: str - organization_name: str - webhook_callback_url: str | None = None - max_steps_per_run: int | None = None - max_retries_per_step: int | None = None - domain: str | None = None - bw_organization_id: str | None = None - bw_collection_ids: list[str] | None = None - - created_at: datetime - modified_at: datetime - - -class OrganizationAuthToken(BaseModel): - id: str - organization_id: str - token_type: OrganizationAuthTokenType - token: str - valid: bool - created_at: datetime - modified_at: datetime - - class TokenPayload(BaseModel): sub: str exp: int diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index db37a2d36..a5981274c 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -34,10 +34,11 @@ from skyvern.forge.sdk.core.security import generate_skyvern_signature from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType from skyvern.forge.sdk.executor.factory import AsyncExecutorFactory -from skyvern.forge.sdk.models import Organization, Step +from skyvern.forge.sdk.models import Step from skyvern.forge.sdk.schemas.organizations import ( GetOrganizationAPIKeysResponse, GetOrganizationsResponse, + Organization, OrganizationUpdate, ) from skyvern.forge.sdk.schemas.task_generations import GenerateTaskRequest, TaskGeneration, TaskGenerationBase diff --git a/skyvern/forge/sdk/schemas/organizations.py b/skyvern/forge/sdk/schemas/organizations.py index 8d145d5b4..2b4494248 100644 --- a/skyvern/forge/sdk/schemas/organizations.py +++ b/skyvern/forge/sdk/schemas/organizations.py @@ -1,6 +1,34 @@ -from pydantic import BaseModel +from datetime import datetime -from skyvern.forge.sdk.models import Organization, OrganizationAuthToken +from pydantic import BaseModel, ConfigDict + +from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType + + +class Organization(BaseModel): + model_config = ConfigDict(from_attributes=True) + + organization_id: str + organization_name: str + webhook_callback_url: str | None = None + max_steps_per_run: int | None = None + max_retries_per_step: int | None = None + domain: str | None = None + bw_organization_id: str | None = None + bw_collection_ids: list[str] | None = None + + created_at: datetime + modified_at: datetime + + +class OrganizationAuthToken(BaseModel): + id: str + organization_id: str + token_type: OrganizationAuthTokenType + token: str + valid: bool + created_at: datetime + modified_at: datetime class GetOrganizationsResponse(BaseModel): diff --git a/skyvern/forge/sdk/services/org_auth_service.py b/skyvern/forge/sdk/services/org_auth_service.py index 1eda1a7ea..edd57d4e7 100644 --- a/skyvern/forge/sdk/services/org_auth_service.py +++ b/skyvern/forge/sdk/services/org_auth_service.py @@ -12,7 +12,8 @@ from skyvern.forge import app from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.db.client import AgentDB -from skyvern.forge.sdk.models import Organization, OrganizationAuthTokenType, TokenPayload +from skyvern.forge.sdk.models import TokenPayload +from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthTokenType AUTHENTICATION_TTL = 60 * 60 # one hour CACHE_SIZE = 128 diff --git a/skyvern/forge/sdk/services/org_auth_token_service.py b/skyvern/forge/sdk/services/org_auth_token_service.py index f7faab181..ad05b331d 100644 --- a/skyvern/forge/sdk/services/org_auth_token_service.py +++ b/skyvern/forge/sdk/services/org_auth_token_service.py @@ -4,7 +4,7 @@ from skyvern.forge.app import DATABASE from skyvern.forge.sdk.core import security -from skyvern.forge.sdk.models import OrganizationAuthToken, OrganizationAuthTokenType +from skyvern.forge.sdk.schemas.organizations import OrganizationAuthToken, OrganizationAuthTokenType LOG = structlog.get_logger() API_KEY_LIFETIME = timedelta(weeks=5200) diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 5c0429724..cf19e3eb3 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -5,7 +5,7 @@ from skyvern.exceptions import BitwardenBaseError, WorkflowRunContextNotInitialized from skyvern.forge.sdk.api.aws import AsyncAWSClient -from skyvern.forge.sdk.models import Organization +from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import TaskStatus from skyvern.forge.sdk.services.bitwarden import BitwardenConstants, BitwardenService from skyvern.forge.sdk.workflow.exceptions import OutputParameterKeyCollisionError diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index c407f19dd..7575a28a9 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -22,7 +22,8 @@ from skyvern.forge.sdk.core.security import generate_skyvern_signature from skyvern.forge.sdk.core.skyvern_context import SkyvernContext from skyvern.forge.sdk.db.enums import TaskType -from skyvern.forge.sdk.models import Organization, Step +from skyvern.forge.sdk.models import Step +from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task from skyvern.forge.sdk.workflow.exceptions import ( ContextParameterSourceNotDefined,