add SkyvernClient (#1943)
This commit is contained in:
@@ -6,20 +6,18 @@ from pydantic import Field
|
||||
from skyvern_langchain.schema import CreateTaskInput, GetTaskInput
|
||||
from skyvern_langchain.settings import settings
|
||||
|
||||
from skyvern.agent import Agent
|
||||
from skyvern.agent import SkyvernAgent
|
||||
from skyvern.forge import app
|
||||
from skyvern.forge.prompts import prompt_engine
|
||||
from skyvern.forge.sdk.schemas.observers import ObserverTask, ObserverTaskRequest
|
||||
from skyvern.forge.sdk.schemas.task_generations import TaskGenerationBase
|
||||
from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Request
|
||||
from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskRequest, TaskResponse
|
||||
|
||||
agent = Agent()
|
||||
|
||||
|
||||
class SkyvernTaskBaseTool(BaseTool):
|
||||
engine: Literal["TaskV1", "TaskV2"] = Field(default=settings.engine)
|
||||
run_task_timeout_seconds: int = Field(default=settings.run_task_timeout_seconds)
|
||||
agent: Agent = agent
|
||||
agent: SkyvernAgent = SkyvernAgent()
|
||||
|
||||
def _run(self, *args: Any, **kwargs: Any) -> None:
|
||||
raise NotImplementedError("skyvern task tool does not support sync")
|
||||
@@ -36,7 +34,7 @@ class RunTask(SkyvernTaskBaseTool):
|
||||
description: str = """Use Skyvern agent to run a task. This function won't return until the task is finished."""
|
||||
args_schema: Type[BaseModel] = CreateTaskInput
|
||||
|
||||
async def _arun(self, user_prompt: str, url: str | None = None) -> TaskResponse | ObserverTask:
|
||||
async def _arun(self, user_prompt: str, url: str | None = None) -> TaskResponse | TaskV2:
|
||||
if self.engine == "TaskV1":
|
||||
return await self._arun_task_v1(user_prompt=user_prompt, url=url)
|
||||
else:
|
||||
@@ -50,8 +48,8 @@ class RunTask(SkyvernTaskBaseTool):
|
||||
|
||||
return await self.agent.run_task(task_request=task_request, timeout_seconds=self.run_task_timeout_seconds)
|
||||
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> ObserverTask:
|
||||
task_request = ObserverTaskRequest(user_prompt=user_prompt, url=url)
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> TaskV2:
|
||||
task_request = TaskV2Request(user_prompt=user_prompt, url=url)
|
||||
return await self.agent.run_observer_task_v_2(
|
||||
task_request=task_request, timeout_seconds=self.run_task_timeout_seconds
|
||||
)
|
||||
@@ -62,7 +60,7 @@ class DispatchTask(SkyvernTaskBaseTool):
|
||||
description: str = """Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background."""
|
||||
args_schema: Type[BaseModel] = CreateTaskInput
|
||||
|
||||
async def _arun(self, user_prompt: str, url: str | None = None) -> CreateTaskResponse | ObserverTask:
|
||||
async def _arun(self, user_prompt: str, url: str | None = None) -> CreateTaskResponse | TaskV2:
|
||||
if self.engine == "TaskV1":
|
||||
return await self._arun_task_v1(user_prompt=user_prompt, url=url)
|
||||
else:
|
||||
@@ -76,8 +74,8 @@ class DispatchTask(SkyvernTaskBaseTool):
|
||||
|
||||
return await self.agent.create_task(task_request=task_request)
|
||||
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> ObserverTask:
|
||||
task_request = ObserverTaskRequest(user_prompt=user_prompt, url=url)
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> TaskV2:
|
||||
task_request = TaskV2Request(user_prompt=user_prompt, url=url)
|
||||
return await self.agent.observer_task_v_2(task_request=task_request)
|
||||
|
||||
|
||||
@@ -86,7 +84,7 @@ class GetTask(SkyvernTaskBaseTool):
|
||||
description: str = """Use Skyvern agent to get a task."""
|
||||
args_schema: Type[BaseModel] = GetTaskInput
|
||||
|
||||
async def _arun(self, task_id: str) -> TaskResponse | ObserverTask | None:
|
||||
async def _arun(self, task_id: str) -> TaskResponse | TaskV2 | None:
|
||||
if self.engine == "TaskV1":
|
||||
return await self._arun_task_v1(task_id=task_id)
|
||||
else:
|
||||
@@ -95,5 +93,5 @@ class GetTask(SkyvernTaskBaseTool):
|
||||
async def _arun_task_v1(self, task_id: str) -> TaskResponse | None:
|
||||
return await self.agent.get_task(task_id=task_id)
|
||||
|
||||
async def _arun_task_v2(self, task_id: str) -> ObserverTask | None:
|
||||
async def _arun_task_v2(self, task_id: str) -> TaskV2 | None:
|
||||
return await self.agent.get_observer_task_v_2(task_id=task_id)
|
||||
|
||||
@@ -7,7 +7,7 @@ from skyvern_langchain.schema import CreateTaskInput, GetTaskInput
|
||||
from skyvern_langchain.settings import settings
|
||||
|
||||
from skyvern.client import AsyncSkyvern
|
||||
from skyvern.forge.sdk.schemas.observers import ObserverTaskRequest
|
||||
from skyvern.forge.sdk.schemas.task_v2 import TaskV2Request
|
||||
from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskRequest, TaskResponse
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ class RunTask(SkyvernTaskBaseTool):
|
||||
)
|
||||
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> TaskResponse:
|
||||
task_request = ObserverTaskRequest(url=url, user_prompt=user_prompt)
|
||||
task_request = TaskV2Request(url=url, user_prompt=user_prompt)
|
||||
return await self.get_client().agent.run_observer_task_v_2(
|
||||
timeout_seconds=self.run_task_timeout_seconds,
|
||||
user_prompt=task_request.user_prompt,
|
||||
@@ -106,7 +106,7 @@ class DispatchTask(SkyvernTaskBaseTool):
|
||||
)
|
||||
|
||||
async def _arun_task_v2(self, user_prompt: str, url: str | None = None) -> Dict[str, Any | None]:
|
||||
task_request = ObserverTaskRequest(url=url, user_prompt=user_prompt)
|
||||
task_request = TaskV2Request(url=url, user_prompt=user_prompt)
|
||||
return await self.get_client().agent.observer_task_v_2(
|
||||
user_prompt=task_request.user_prompt,
|
||||
url=task_request.url,
|
||||
|
||||
@@ -4,20 +4,18 @@ from llama_index.core.tools import FunctionTool
|
||||
from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
|
||||
from skyvern_llamaindex.settings import settings
|
||||
|
||||
from skyvern.agent import Agent
|
||||
from skyvern.agent import SkyvernAgent
|
||||
from skyvern.forge import app
|
||||
from skyvern.forge.prompts import prompt_engine
|
||||
from skyvern.forge.sdk.schemas.observers import ObserverTask, ObserverTaskRequest
|
||||
from skyvern.forge.sdk.schemas.task_generations import TaskGenerationBase
|
||||
from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Request
|
||||
from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskRequest, TaskResponse
|
||||
|
||||
default_agent = Agent()
|
||||
|
||||
|
||||
class SkyvernTool:
|
||||
def __init__(self, agent: Optional[Agent] = None):
|
||||
def __init__(self, agent: Optional[SkyvernAgent] = None):
|
||||
if agent is None:
|
||||
agent = default_agent
|
||||
agent = SkyvernAgent()
|
||||
self.agent = agent
|
||||
|
||||
def run_task(self) -> FunctionTool:
|
||||
@@ -43,12 +41,12 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
agent: Optional[Agent] = None,
|
||||
agent: SkyvernAgent | None = None,
|
||||
engine: Literal["TaskV1", "TaskV2"] = settings.engine,
|
||||
run_task_timeout_seconds: int = settings.run_task_timeout_seconds,
|
||||
) -> None:
|
||||
if agent is None:
|
||||
agent = Agent()
|
||||
agent = SkyvernAgent()
|
||||
self.agent = agent
|
||||
self.engine = engine
|
||||
self.run_task_timeout_seconds = run_task_timeout_seconds
|
||||
@@ -59,7 +57,7 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
llm_response = await app.LLM_API_HANDLER(prompt=llm_prompt, prompt_name="generate-task")
|
||||
return TaskGenerationBase.model_validate(llm_response)
|
||||
|
||||
async def run_task(self, user_prompt: str, url: Optional[str] = None) -> TaskResponse | ObserverTask:
|
||||
async def run_task(self, user_prompt: str, url: Optional[str] = None) -> TaskResponse | TaskV2:
|
||||
"""
|
||||
Use Skyvern agent to run a task. This function won't return until the task is finished.
|
||||
|
||||
@@ -73,7 +71,7 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
else:
|
||||
return await self.run_task_v2(user_prompt=user_prompt, url=url)
|
||||
|
||||
async def dispatch_task(self, user_prompt: str, url: Optional[str] = None) -> CreateTaskResponse | ObserverTask:
|
||||
async def dispatch_task(self, user_prompt: str, url: Optional[str] = None) -> CreateTaskResponse | TaskV2:
|
||||
"""
|
||||
Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.
|
||||
|
||||
@@ -87,7 +85,7 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
else:
|
||||
return await self.dispatch_task_v2(user_prompt=user_prompt, url=url)
|
||||
|
||||
async def get_task(self, task_id: str) -> TaskResponse | ObserverTask | None:
|
||||
async def get_task(self, task_id: str) -> TaskResponse | TaskV2 | None:
|
||||
"""
|
||||
Use Skyvern agent to get a task.
|
||||
|
||||
@@ -119,15 +117,15 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
async def get_task_v1(self, task_id: str) -> TaskResponse | None:
|
||||
return await self.agent.get_task(task_id=task_id)
|
||||
|
||||
async def run_task_v2(self, user_prompt: str, url: Optional[str] = None) -> ObserverTask:
|
||||
task_request = ObserverTaskRequest(user_prompt=user_prompt, url=url)
|
||||
async def run_task_v2(self, user_prompt: str, url: Optional[str] = None) -> TaskV2:
|
||||
task_request = TaskV2Request(user_prompt=user_prompt, url=url)
|
||||
return await self.agent.run_observer_task_v_2(
|
||||
task_request=task_request, timeout_seconds=self.run_task_timeout_seconds
|
||||
)
|
||||
|
||||
async def dispatch_task_v2(self, user_prompt: str, url: Optional[str] = None) -> ObserverTask:
|
||||
task_request = ObserverTaskRequest(user_prompt=user_prompt, url=url)
|
||||
async def dispatch_task_v2(self, user_prompt: str, url: Optional[str] = None) -> TaskV2:
|
||||
task_request = TaskV2Request(user_prompt=user_prompt, url=url)
|
||||
return await self.agent.observer_task_v_2(task_request=task_request)
|
||||
|
||||
async def get_task_v2(self, task_id: str) -> ObserverTask | None:
|
||||
async def get_task_v2(self, task_id: str) -> TaskV2 | None:
|
||||
return await self.agent.get_observer_task_v_2(task_id=task_id)
|
||||
|
||||
@@ -7,7 +7,7 @@ from pydantic import BaseModel
|
||||
from skyvern_llamaindex.settings import settings
|
||||
|
||||
from skyvern.client import AsyncSkyvern
|
||||
from skyvern.forge.sdk.schemas.observers import ObserverTaskRequest
|
||||
from skyvern.forge.sdk.schemas.task_v2 import TaskV2Request
|
||||
from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskRequest, TaskResponse
|
||||
|
||||
|
||||
@@ -153,7 +153,7 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
return await self.client.agent.get_task(task_id=task_id)
|
||||
|
||||
async def run_task_v2(self, user_prompt: str, url: Optional[str] = None) -> Dict[str, Any | None]:
|
||||
task_request = ObserverTaskRequest(url=url, user_prompt=user_prompt)
|
||||
task_request = TaskV2Request(url=url, user_prompt=user_prompt)
|
||||
return await self.client.agent.run_observer_task_v_2(
|
||||
timeout_seconds=self.run_task_timeout_seconds,
|
||||
user_prompt=task_request.user_prompt,
|
||||
@@ -162,7 +162,7 @@ class SkyvernTaskToolSpec(BaseToolSpec):
|
||||
)
|
||||
|
||||
async def dispatch_task_v2(self, user_prompt: str, url: Optional[str] = None) -> Dict[str, Any | None]:
|
||||
task_request = ObserverTaskRequest(url=url, user_prompt=user_prompt)
|
||||
task_request = TaskV2Request(url=url, user_prompt=user_prompt)
|
||||
return await self.client.agent.observer_task_v_2(
|
||||
user_prompt=task_request.user_prompt,
|
||||
url=task_request.url,
|
||||
|
||||
Reference in New Issue
Block a user