migrate Organization data model from skyvern/forge/sdk/models.py to skyvern/forge/sdk/schemas/organizations.py (#1343)

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):

skyvern/forge/sdk/db/client.py

@@ -49,7 +49,8 @@ from skyvern.forge.sdk.db.utils import (
     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

skyvern/forge/sdk/db/utils.py

@@ -22,7 +22,8 @@ from skyvern.forge.sdk.db.models import (
     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,

skyvern/forge/sdk/models.py

@@ -3,9 +3,8 @@ from __future__ import annotations
 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 @@ class Step(BaseModel):
         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

skyvern/forge/sdk/routes/agent_protocol.py

@@ -34,10 +34,11 @@ from skyvern.forge.sdk.core.permissions.permission_checker_factory import Permis
 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

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):

skyvern/forge/sdk/services/org_auth_service.py

@@ -12,7 +12,8 @@ from skyvern.config import settings
 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

skyvern/forge/sdk/services/org_auth_token_service.py

@@ -4,7 +4,7 @@ import structlog
 
 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)

skyvern/forge/sdk/workflow/context_manager.py

@@ -5,7 +5,7 @@ import structlog
 
 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

skyvern/forge/sdk/workflow/service.py

@@ -22,7 +22,8 @@ from skyvern.forge.sdk.core import skyvern_context
 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,