llamaindex support (#1758)

skyvern/agent/parameter.py

@@ -1,3 +1,5 @@
+from pydantic import BaseModel, Field
+
 from skyvern.forge.sdk.schemas.observers import ObserverTaskRequest
 from skyvern.forge.sdk.schemas.tasks import TaskRequest
 
@@ -8,3 +10,36 @@ class TaskV1Request(TaskRequest):
 
 class TaskV2Request(ObserverTaskRequest):
     max_iterations: int = 10
+
+
+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",
+        default="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",
+        default="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",
+        default="https://api.skyvern.com",
+    )
+    task_id: str

skyvern/langchain/remote.py

@@ -1,44 +1,13 @@
 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.parameter import GetTaskSchema, RunTaskV1Schema, RunTaskV2Schema, 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"

skyvern/llama_index/local.py

@@ -0,0 +1,18 @@
+from llama_index.core.tools import FunctionTool
+
+from skyvern.agent.local import Agent
+from skyvern.agent.parameter import TaskV1Request, TaskV2Request
+
+run_task_v1 = FunctionTool.from_defaults(
+    async_fn=lambda **kwargs: Agent().run_task_v1(TaskV1Request(**kwargs)),
+    name="run-local-skyvern-simple-task",
+    description="Use local Skyvern to run a v1 task. v1 task is usually used for the simple tasks.",
+    fn_schema=TaskV1Request,
+)
+
+run_task_v2 = FunctionTool.from_defaults(
+    async_fn=lambda **kwargs: Agent().run_task_v2(TaskV2Request(**kwargs)),
+    name="run-local-skyvern-complicated-task",
+    description="Use local Skyvern to run a v2 task. v2 task is usually used for the complicated tasks.",
+    fn_schema=TaskV2Request,
+)

skyvern/llama_index/remote.py

@@ -0,0 +1,43 @@
+from llama_index.core.tools import FunctionTool
+
+from skyvern.agent.parameter import GetTaskSchema, RunTaskV1Schema, RunTaskV2Schema, TaskV1Request, TaskV2Request
+from skyvern.agent.remote import RemoteAgent
+
+run_task_v1 = FunctionTool.from_defaults(
+    async_fn=lambda task, api_key, endpoint="https://api.skyvern.com": RemoteAgent(api_key, endpoint).run_task_v1(
+        TaskV1Request.model_validate(task)
+    ),
+    name="run-remote-skyvern-simple-task",
+    description="Use remote Skyvern to run a v1 task. v1 task is usually used for the simple tasks.",
+    fn_schema=RunTaskV1Schema,
+)
+
+
+get_task_v1 = FunctionTool.from_defaults(
+    async_fn=lambda task_id, api_key, endpoint="https://api.skyvern.com": RemoteAgent(api_key, endpoint).get_task_v1(
+        task_id
+    ),
+    name="get-remote-skyvern-simple-task",
+    description="Use remote Skyvern to get a v1 task information. v1 task is usually used for the simple tasks.",
+    fn_schema=GetTaskSchema,
+)
+
+
+run_task_v2 = FunctionTool.from_defaults(
+    async_fn=lambda task, api_key, endpoint="https://api.skyvern.com": RemoteAgent(api_key, endpoint).run_task_v2(
+        TaskV2Request.model_validate(task)
+    ),
+    name="run-remote-skyvern-complicated-task",
+    description="Use remote Skyvern to run a v2 task. v2 task is usually used for the complicated tasks.",
+    fn_schema=RunTaskV2Schema,
+)
+
+
+get_task_v2 = FunctionTool.from_defaults(
+    async_fn=lambda task_id, api_key, endpoint="https://api.skyvern.com": RemoteAgent(api_key, endpoint).get_task_v2(
+        task_id
+    ),
+    name="get-remote-skyvern-complicated-task",
+    description="Use remote Skyvern to get a v2 task information. v2 task is usually used for the complicated tasks.",
+    fn_schema=GetTaskSchema,
+)