task v2 proxy location update (#1671)

This commit is contained in:
Shuchang Zheng 2025-01-29 04:31:03 +08:00 committed by GitHub
parent 003f3d0284
commit c7f56f30b3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -49,7 +49,7 @@ from skyvern.forge.sdk.artifact.models import ArtifactType
from skyvern.forge.sdk.core.validators import prepend_scheme_and_validate_url from skyvern.forge.sdk.core.validators import prepend_scheme_and_validate_url
from skyvern.forge.sdk.db.enums import TaskType from skyvern.forge.sdk.db.enums import TaskType
from skyvern.forge.sdk.schemas.observers import ObserverTaskStatus from skyvern.forge.sdk.schemas.observers import ObserverTaskStatus
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskOutput, TaskStatus from skyvern.forge.sdk.schemas.tasks import Task, TaskOutput, TaskStatus
from skyvern.forge.sdk.workflow.context_manager import BlockMetadata, WorkflowRunContext from skyvern.forge.sdk.workflow.context_manager import BlockMetadata, WorkflowRunContext
from skyvern.forge.sdk.workflow.exceptions import ( from skyvern.forge.sdk.workflow.exceptions import (
FailedToFormatJinjaStyleParameter, FailedToFormatJinjaStyleParameter,
@ -2106,12 +2106,15 @@ class TaskV2Block(Block):
organization = await app.DATABASE.get_organization(organization_id) organization = await app.DATABASE.get_organization(organization_id)
if not organization: if not organization:
raise ValueError(f"Organization not found {organization_id}") raise ValueError(f"Organization not found {organization_id}")
workflow_run = await app.DATABASE.get_workflow_run(workflow_run_id, organization_id)
if not workflow_run:
raise ValueError(f"WorkflowRun not found {workflow_run_id} when running TaskV2Block")
observer_task = await observer_service.initialize_observer_task( observer_task = await observer_service.initialize_observer_task(
organization, organization,
user_prompt=self.prompt, user_prompt=self.prompt,
user_url=self.url, user_url=self.url,
parent_workflow_run_id=workflow_run_id, parent_workflow_run_id=workflow_run_id,
proxy_location=ProxyLocation.NONE, proxy_location=workflow_run.proxy_location,
) )
await app.DATABASE.update_observer_cruise( await app.DATABASE.update_observer_cruise(
observer_task.observer_cruise_id, status=ObserverTaskStatus.queued, organization_id=organization_id observer_task.observer_cruise_id, status=ObserverTaskStatus.queued, organization_id=organization_id