From f146247fe9ba55ab7049f9c72ce164a84a7f1806 Mon Sep 17 00:00:00 2001 From: Kerem Yilmaz Date: Thu, 18 Jul 2024 18:09:28 -0700 Subject: [PATCH] Support get organizations and get api keys APIs in the open source (#623) --- skyvern/forge/agent.py | 14 ---------- skyvern/forge/sdk/db/client.py | 3 --- skyvern/forge/sdk/routes/agent_protocol.py | 30 +++++++++++++++++++++- skyvern/forge/sdk/schemas/organizations.py | 10 ++++++++ skyvern/forge/sdk/schemas/tasks.py | 2 -- 5 files changed, 39 insertions(+), 20 deletions(-) diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index f5875855..5a74efae 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -232,20 +232,6 @@ class ForgeAgent: # TODO: shall we send task response here? return step, None, None - context = skyvern_context.current() - override_max_steps_per_run = context.max_steps_override if context else None - max_steps_per_run = ( - override_max_steps_per_run - or task.max_steps_per_run - or organization.max_steps_per_run - or SettingsManager.get_settings().MAX_STEPS_PER_RUN - ) - if max_steps_per_run and task.max_steps_per_run != max_steps_per_run: - await app.DATABASE.update_task( - task_id=task.task_id, - organization_id=organization.organization_id, - max_steps_per_run=max_steps_per_run, - ) next_step: Step | None = None detailed_output: DetailedAgentStepOutput | None = None num_files_before = 0 diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 62460c05..ede377f4 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -376,7 +376,6 @@ class AgentDB: extracted_information: dict[str, Any] | list | str | None = None, failure_reason: str | None = None, errors: list[dict[str, Any]] | None = None, - max_steps_per_run: int | None = None, organization_id: str | None = None, ) -> Task: if status is None and extracted_information is None and failure_reason is None and errors is None: @@ -398,8 +397,6 @@ class AgentDB: task.failure_reason = failure_reason if errors is not None: task.errors = errors - if max_steps_per_run is not None: - task.max_steps_per_run = max_steps_per_run await session.commit() updated_task = await self.get_task(task_id, organization_id=organization_id) if not updated_task: diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index b4edd42c..66abfc79 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -29,9 +29,14 @@ from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.core.permissions.permission_checker_factory import PermissionCheckerFactory 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.schemas.organizations import OrganizationUpdate +from skyvern.forge.sdk.schemas.organizations import ( + GetOrganizationAPIKeysResponse, + GetOrganizationsResponse, + OrganizationUpdate, +) from skyvern.forge.sdk.schemas.task_generations import GenerateTaskRequest, TaskGeneration, TaskGenerationBase from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, Task, TaskRequest, TaskResponse, TaskStatus from skyvern.forge.sdk.services import org_auth_service @@ -795,6 +800,29 @@ async def update_organization( ) +@base_router.get("/organizations/", include_in_schema=False) +@base_router.get("/organizations") +async def get_organizations( + current_org: Organization = Depends(org_auth_service.get_current_org), +) -> GetOrganizationsResponse: + return GetOrganizationsResponse(organizations=[current_org]) + + +@base_router.get("/organizations/{organization_id}/apikeys/", include_in_schema=False) +@base_router.get("/organizations/{organization_id}/apikeys") +async def get_org_api_keys( + organization_id: str, + current_org: Organization = Depends(org_auth_service.get_current_org), +) -> GetOrganizationAPIKeysResponse: + if organization_id != current_org.organization_id: + raise HTTPException(status_code=403, detail="You do not have permission to access this organization") + api_keys = [] + org_auth_token = await app.DATABASE.get_valid_org_auth_token(organization_id, OrganizationAuthTokenType.api) + if org_auth_token: + api_keys.append(org_auth_token) + return GetOrganizationAPIKeysResponse(api_keys=api_keys) + + async def validate_file_size(file: UploadFile) -> UploadFile: try: file.file.seek(0, 2) # Move the pointer to the end of the file diff --git a/skyvern/forge/sdk/schemas/organizations.py b/skyvern/forge/sdk/schemas/organizations.py index ae67a53f..5c16802f 100644 --- a/skyvern/forge/sdk/schemas/organizations.py +++ b/skyvern/forge/sdk/schemas/organizations.py @@ -1,5 +1,15 @@ from pydantic import BaseModel +from skyvern.forge.sdk.models import Organization, OrganizationAuthToken + + +class GetOrganizationsResponse(BaseModel): + organizations: list[Organization] + + +class GetOrganizationAPIKeysResponse(BaseModel): + api_keys: list[OrganizationAuthToken] + class OrganizationUpdate(BaseModel): organization_name: str | None = None diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index 4c704d76..c4dcd584 100644 --- a/skyvern/forge/sdk/schemas/tasks.py +++ b/skyvern/forge/sdk/schemas/tasks.py @@ -221,7 +221,6 @@ class Task(TaskRequest): screenshot_url=screenshot_url, recording_url=recording_url, errors=self.errors, - max_steps_per_run=self.max_steps_per_run, ) @@ -237,7 +236,6 @@ class TaskResponse(BaseModel): recording_url: str | None = None failure_reason: str | None = None errors: list[dict[str, Any]] = [] - max_steps_per_run: int | None = None class TaskOutput(BaseModel):