migrate Organization data model from skyvern/forge/sdk/models.py to s…

…kyvern/forge/sdk/schemas/organizations.py (#1343)
Skyvern-AI · Dec 7, 2024 · 7591873 · 7591873
1 parent 6151527
commit 7591873
Show file tree

Hide file tree

Showing 14 changed files with 50 additions and 42 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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,