Skyvern/skyvern/langchain/remote.py
2025-02-10 16:03:57 +08:00

65 lines
2.8 KiB
Python

from typing import Any, Dict
from langchain.tools import tool
from pydantic import BaseModel, Field
from skyvern.agent.parameter import TaskV1Request, TaskV2Request
from skyvern.agent.remote import RemoteAgent
from skyvern.forge.sdk.schemas.observers import ObserverTask
from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
class RunTaskV1Schema(BaseModel):
api_key: str = Field(
description="The API key of the Skyvern API. You can get the API key from the Skyvern dashboard."
)
endpoint: str = Field(
description="The endpoint of the Skyvern API. Don't add any path to the endpoint. Default is https://api.skyvern.com"
)
task: TaskV1Request
class RunTaskV2Schema(BaseModel):
api_key: str = Field(
description="The API key of the Skyvern API. You can get the API key from the Skyvern dashboard."
)
endpoint: str = Field(
description="The endpoint of the Skyvern API. Don't add any path to the endpoint. Default is https://api.skyvern.com"
)
task: TaskV2Request
class GetTaskSchema(BaseModel):
api_key: str = Field(
description="The API key of the Skyvern API. You can get the API key from the Skyvern dashboard."
)
endpoint: str = Field(
description="The endpoint of the Skyvern API. Don't add any path to the endpoint. Default is https://api.skyvern.com"
)
task_id: str
@tool("run-remote-skyvern-simple-task", args_schema=RunTaskV1Schema)
async def run_task_v1(
task: Dict[str, Any], api_key: str, endpoint: str = "https://api.skyvern.com"
) -> CreateTaskResponse:
"""Use remote Skyvern to run a v1 task. v1 task is usually used for the simple tasks."""
return await RemoteAgent(api_key, endpoint).run_task_v1(TaskV1Request.model_validate(task))
@tool("get-remote-skyvern-simple-task", args_schema=GetTaskSchema)
async def get_task_v1(task_id: str, api_key: str, endpoint: str = "https://api.skyvern.com") -> TaskResponse:
"""Use remote Skyvern to get a v1 task information. v1 task is usually used for the simple tasks."""
return await RemoteAgent(api_key, endpoint).get_task_v1(task_id)
@tool("run-skyvern-complicated-task", args_schema=RunTaskV2Schema)
async def run_task_v2(task: Dict[str, Any], api_key: str, endpoint: str = "https://api.skyvern.com") -> ObserverTask:
"""Use remote Skyvern to run a v2 task. v2 task is usually used for the complicated tasks."""
return await RemoteAgent(api_key, endpoint).run_task_v2(TaskV2Request.model_validate(task))
@tool("get-remote-skyvern-complicated-task", args_schema=GetTaskSchema)
async def get_task_v2(task_id: str, api_key: str, endpoint: str = "https://api.skyvern.com") -> ObserverTask:
"""Use remote Skyvern to get a v2 task information. v2 task is usually used for the complicated tasks."""
return await RemoteAgent(api_key, endpoint).get_task_v2(task_id)