skyvern-llamaindex 0.0.1__tar.gz → 0.0.2__tar.gz

{skyvern_llamaindex-0.0.1 → skyvern_llamaindex-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skyvern-llamaindex
-Version: 0.0.1
+Version: 0.0.2
 Summary: Skyvern integration for LlamaIndex
 Author: lawyzheng
 Author-email: lawy@skyvern.com
@@ -19,9 +19,9 @@ Description-Content-Type: text/markdown
   - [Installation](#installation)
   - [Usage](#usage)
     - [Run a task(sync) with skyvern agent (calling skyvern agent function directly in the tool)](#run-a-tasksync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
-    - [Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)](#run-a-taskasync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
+    - [Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)](#dispatch-a-taskasync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
     - [Run a task(sync) with skyvern client (calling skyvern OpenAPI in the tool)](#run-a-tasksync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
-    - [Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)](#run-a-taskasync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
+    - [Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)](#dispatch-a-taskasync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
 
@@ -48,14 +48,14 @@ import asyncio
 from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
-from skyvern_llamaindex.agent import SkyvernAgentToolSpec
+from skyvern_llamaindex.agent import SkyvernToolSpec
 
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_tool = SkyvernAgentToolSpec()
+skyvern_tool = SkyvernToolSpec()
 
-tools = skyvern_tool.to_tool_list(["run_task_v2"])
+tools = skyvern_tool.to_tool_list(["run_task"])
 
 agent = OpenAIAgent.from_tools(
     tools=tools,
@@ -69,8 +69,8 @@ response = agent.chat("Run the task with skyvern. The task is about 'Navigate to
 print(response)
 ```
 
-### Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
-> async task will return immediately and the task will be running in the background. You can use `get_task_v2` tool to poll the task information until the task is finished.
+### Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
+> dispatch task will return immediately and the task will be running in the background. You can use `get_task` tool to poll the task information until the task is finished.
 
 :warning: :warning: if you want to run this code block, you need to run `skyvern init --openai-api-key <your_openai_api_key>` command in your terminal to set up skyvern first.
 
@@ -80,7 +80,7 @@ from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
 from llama_index.core.tools import FunctionTool
-from skyvern_llamaindex.agent import SkyvernAgentToolSpec
+from skyvern_llamaindex.agent import SkyvernToolSpec
 
 async def sleep(seconds: int) -> str:
     await asyncio.sleep(seconds)
@@ -89,7 +89,7 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_tool = SkyvernAgentToolSpec()
+skyvern_tool = SkyvernToolSpec()
 
 sleep_tool = FunctionTool.from_defaults(
     async_fn=sleep,
@@ -97,7 +97,7 @@ sleep_tool = FunctionTool.from_defaults(
     name="sleep",
 )
 
-tools = skyvern_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
+tools = skyvern_tool.to_tool_list(["dispatch_task", "get_task"])
 tools.append(sleep_tool)
 
 agent = OpenAIAgent.from_tools(
@@ -107,7 +107,7 @@ agent = OpenAIAgent.from_tools(
     max_function_calls=10,
 )
 
-response = agent.chat("Queue a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
+response = agent.chat("Run a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
 print(response)
 
 ```
@@ -122,7 +122,7 @@ import asyncio
 from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
-from skyvern_llamaindex.client import SkyvernClientToolSpec
+from skyvern_llamaindex.client import SkyvernToolSpec
 
 
 async def sleep(seconds: int) -> str:
@@ -132,11 +132,11 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_client_tool = SkyvernClientToolSpec(
+skyvern_client_tool = SkyvernToolSpec(
     credential="<your_organization_api_key>",
 )
 
-tools = skyvern_client_tool.to_tool_list(["run_task_v2"])
+tools = skyvern_client_tool.to_tool_list(["run_task"])
 
 agent = OpenAIAgent.from_tools(
     tools=tools,
@@ -150,8 +150,8 @@ print(response)
 
 ```
 
-### Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
-> async task will return immediately and the task will be running in the background. You can use `GetSkyvernClientTaskV2Tool` tool to poll the task information until the task is finished.
+### Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
+> dispatch task will return immediately and the task will be running in the background. You can use `get_task` tool to poll the task information until the task is finished.
 
 no need to run `skyvern init` command in your terminal to set up skyvern before using this integration.
 
@@ -161,7 +161,7 @@ from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
 from llama_index.core.tools import FunctionTool
-from skyvern_llamaindex.client import SkyvernClientToolSpec
+from skyvern_llamaindex.client import SkyvernToolSpec
 
 
 async def sleep(seconds: int) -> str:
@@ -171,7 +171,7 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_client_tool = SkyvernClientToolSpec(
+skyvern_client_tool = SkyvernToolSpec(
     credential="<your_organization_api_key>",
 )
 
@@ -181,7 +181,7 @@ sleep_tool = FunctionTool.from_defaults(
     name="sleep",
 )
 
-tools = skyvern_client_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
+tools = skyvern_client_tool.to_tool_list(["dispatch_task", "get_task"])
 tools.append(sleep_tool)
 
 agent = OpenAIAgent.from_tools(
@@ -191,7 +191,7 @@ agent = OpenAIAgent.from_tools(
     max_function_calls=10,
 )
 
-response = agent.chat("Queue a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
+response = agent.chat("Run a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
 print(response)
 
 ```

{skyvern_llamaindex-0.0.1 → skyvern_llamaindex-0.0.2}/README.md

@@ -6,9 +6,9 @@
   - [Installation](#installation)
   - [Usage](#usage)
     - [Run a task(sync) with skyvern agent (calling skyvern agent function directly in the tool)](#run-a-tasksync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
-    - [Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)](#run-a-taskasync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
+    - [Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)](#dispatch-a-taskasync-with-skyvern-agent-calling-skyvern-agent-function-directly-in-the-tool)
     - [Run a task(sync) with skyvern client (calling skyvern OpenAPI in the tool)](#run-a-tasksync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
-    - [Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)](#run-a-taskasync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
+    - [Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)](#dispatch-a-taskasync-with-skyvern-client-calling-skyvern-openapi-in-the-tool)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
 
@@ -35,14 +35,14 @@ import asyncio
 from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
-from skyvern_llamaindex.agent import SkyvernAgentToolSpec
+from skyvern_llamaindex.agent import SkyvernToolSpec
 
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_tool = SkyvernAgentToolSpec()
+skyvern_tool = SkyvernToolSpec()
 
-tools = skyvern_tool.to_tool_list(["run_task_v2"])
+tools = skyvern_tool.to_tool_list(["run_task"])
 
 agent = OpenAIAgent.from_tools(
     tools=tools,
@@ -56,8 +56,8 @@ response = agent.chat("Run the task with skyvern. The task is about 'Navigate to
 print(response)
 ```
 
-### Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
-> async task will return immediately and the task will be running in the background. You can use `get_task_v2` tool to poll the task information until the task is finished.
+### Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
+> dispatch task will return immediately and the task will be running in the background. You can use `get_task` tool to poll the task information until the task is finished.
 
 :warning: :warning: if you want to run this code block, you need to run `skyvern init --openai-api-key <your_openai_api_key>` command in your terminal to set up skyvern first.
 
@@ -67,7 +67,7 @@ from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
 from llama_index.core.tools import FunctionTool
-from skyvern_llamaindex.agent import SkyvernAgentToolSpec
+from skyvern_llamaindex.agent import SkyvernToolSpec
 
 async def sleep(seconds: int) -> str:
     await asyncio.sleep(seconds)
@@ -76,7 +76,7 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_tool = SkyvernAgentToolSpec()
+skyvern_tool = SkyvernToolSpec()
 
 sleep_tool = FunctionTool.from_defaults(
     async_fn=sleep,
@@ -84,7 +84,7 @@ sleep_tool = FunctionTool.from_defaults(
     name="sleep",
 )
 
-tools = skyvern_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
+tools = skyvern_tool.to_tool_list(["dispatch_task", "get_task"])
 tools.append(sleep_tool)
 
 agent = OpenAIAgent.from_tools(
@@ -94,7 +94,7 @@ agent = OpenAIAgent.from_tools(
     max_function_calls=10,
 )
 
-response = agent.chat("Queue a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
+response = agent.chat("Run a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
 print(response)
 
 ```
@@ -109,7 +109,7 @@ import asyncio
 from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
-from skyvern_llamaindex.client import SkyvernClientToolSpec
+from skyvern_llamaindex.client import SkyvernToolSpec
 
 
 async def sleep(seconds: int) -> str:
@@ -119,11 +119,11 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_client_tool = SkyvernClientToolSpec(
+skyvern_client_tool = SkyvernToolSpec(
     credential="<your_organization_api_key>",
 )
 
-tools = skyvern_client_tool.to_tool_list(["run_task_v2"])
+tools = skyvern_client_tool.to_tool_list(["run_task"])
 
 agent = OpenAIAgent.from_tools(
     tools=tools,
@@ -137,8 +137,8 @@ print(response)
 
 ```
 
-### Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
-> async task will return immediately and the task will be running in the background. You can use `GetSkyvernClientTaskV2Tool` tool to poll the task information until the task is finished.
+### Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
+> dispatch task will return immediately and the task will be running in the background. You can use `get_task` tool to poll the task information until the task is finished.
 
 no need to run `skyvern init` command in your terminal to set up skyvern before using this integration.
 
@@ -148,7 +148,7 @@ from dotenv import load_dotenv
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.llms.openai import OpenAI
 from llama_index.core.tools import FunctionTool
-from skyvern_llamaindex.client import SkyvernClientToolSpec
+from skyvern_llamaindex.client import SkyvernToolSpec
 
 
 async def sleep(seconds: int) -> str:
@@ -158,7 +158,7 @@ async def sleep(seconds: int) -> str:
 # load OpenAI API key from .env
 load_dotenv()
 
-skyvern_client_tool = SkyvernClientToolSpec(
+skyvern_client_tool = SkyvernToolSpec(
     credential="<your_organization_api_key>",
 )
 
@@ -168,7 +168,7 @@ sleep_tool = FunctionTool.from_defaults(
     name="sleep",
 )
 
-tools = skyvern_client_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
+tools = skyvern_client_tool.to_tool_list(["dispatch_task", "get_task"])
 tools.append(sleep_tool)
 
 agent = OpenAIAgent.from_tools(
@@ -178,7 +178,7 @@ agent = OpenAIAgent.from_tools(
     max_function_calls=10,
 )
 
-response = agent.chat("Queue a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
+response = agent.chat("Run a task with Skyvern. The task is about 'Navigate to the Hacker News homepage and get the top 3 posts.' Then, get this task information until it's completed. The task information re-get interval should be 60s.")
 print(response)
 
 ```

{skyvern_llamaindex-0.0.1 → skyvern_llamaindex-0.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "skyvern-llamaindex"
-version = "0.0.1"
+version = "0.0.2"
 description = "Skyvern integration for LlamaIndex"
 authors = ["lawyzheng <lawy@skyvern.com>"]
 packages = [{ include = "skyvern_llamaindex" }]

skyvern_llamaindex-0.0.2/skyvern_llamaindex/agent.py

@@ -0,0 +1,109 @@
+from typing import Any, Dict, List, Literal, Tuple
+
+from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
+from llama_index.core.tools.types import ToolMetadata
+from skyvern_llamaindex.schema import GetTaskInput, TaskV1Request, TaskV2Request
+
+from skyvern.agent import Agent
+from skyvern.forge.sdk.schemas.observers import ObserverTask
+from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
+
+
+class SkyvernToolSpec(BaseToolSpec):
+    spec_functions: List[SPEC_FUNCTION_TYPE] = [
+        "run_task",
+        "dispatch_task",
+        "get_task",
+    ]
+    spec_metadata: Dict[str, Dict[str, ToolMetadata]] = {
+        "TaskV1": {
+            "run_task": ToolMetadata(
+                name="run-skyvern-agent-task",
+                description="Use Skyvern agent to run a task. This function won't return until the task is finished.",
+                fn_schema=TaskV1Request,
+            ),
+            "dispatch_task": ToolMetadata(
+                name="dispatch-skyvern-agent-task",
+                description="Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.",
+                fn_schema=TaskV1Request,
+            ),
+            "get_task": ToolMetadata(
+                name="get-skyvern-agent-task",
+                description="Use Skyvern agent to get a task.",
+                fn_schema=GetTaskInput,
+            ),
+        },
+        "TaskV2": {
+            "run_task": ToolMetadata(
+                name="run-skyvern-agent-task",
+                description="Use Skyvern agent to run a task. This function won't return until the task is finished.",
+                fn_schema=TaskV2Request,
+            ),
+            "dispatch_task": ToolMetadata(
+                name="dispatch-skyvern-agent-task",
+                description="Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.",
+                fn_schema=TaskV2Request,
+            ),
+            "get_task": ToolMetadata(
+                name="get-skyvern-agent-task",
+                description="Use Skyvern agent to get a task.",
+                fn_schema=GetTaskInput,
+            ),
+        },
+    }
+
+    def __init__(self, *, engine: Literal["TaskV1", "TaskV2"] = "TaskV2") -> None:
+        self.agent = Agent()
+        self.engine = engine
+
+    def get_metadata_from_fn_name(
+        self, fn_name: str, spec_functions: List[str | Tuple[str, str]] | None = None
+    ) -> ToolMetadata | None:
+        try:
+            getattr(self, fn_name)
+        except AttributeError:
+            return None
+
+        return self.spec_metadata.get(self.engine, {}).get(fn_name)
+
+    async def run_task(self, **kwargs: Dict[str, Any]) -> TaskResponse | ObserverTask:
+        if self.engine == "TaskV1":
+            return await self.run_task_v1(**kwargs)
+        else:
+            return await self.run_task_v2(**kwargs)
+
+    async def dispatch_task(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse | ObserverTask:
+        if self.engine == "TaskV1":
+            return await self.dispatch_task_v1(**kwargs)
+        else:
+            return await self.dispatch_task_v2(**kwargs)
+
+    async def get_task(self, task_id: str) -> TaskResponse | ObserverTask | None:
+        if self.engine == "TaskV1":
+            return await self.get_task_v1(task_id)
+        else:
+            return await self.get_task_v2(task_id)
+
+    async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
+        task_request = TaskV1Request(**kwargs)
+        return await self.agent.run_task(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
+
+    async def dispatch_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
+        task_request = TaskV1Request(**kwargs)
+        return await self.agent.create_task(task_request=task_request)
+
+    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, **kwargs: Dict[str, Any]) -> ObserverTask:
+        task_request = TaskV2Request(**kwargs)
+        return await self.agent.run_observer_task_v_2(
+            task_request=task_request, timeout_seconds=task_request.timeout_seconds
+        )
+
+    async def dispatch_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
+        task_request = TaskV2Request(**kwargs)
+        return await self.agent.observer_task_v_2(task_request=task_request)
+
+    async def get_task_v2(self, task_id: str) -> ObserverTask | None:
+        return await self.agent.get_observer_task_v_2(task_id=task_id)

{skyvern_llamaindex-0.0.1 → skyvern_llamaindex-0.0.2}/skyvern_llamaindex/client.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Literal, Tuple
 
 from httpx import AsyncClient
 from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
@@ -9,56 +9,64 @@ from skyvern.client import AsyncSkyvern
 from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
 
 
-class SkyvernClientToolSpec(BaseToolSpec):
+class SkyvernToolSpec(BaseToolSpec):
     spec_functions: List[SPEC_FUNCTION_TYPE] = [
-        "run_task_v1",
-        "queue_task_v1",
-        "get_task_v1",
-        "run_task_v2",
-        "queue_task_v2",
-        "get_task_v2",
+        "run_task",
+        "dispatch_task",
+        "get_task",
     ]
 
-    spec_metadata: Dict[str, ToolMetadata] = {
-        "run_task_v1": ToolMetadata(
-            name="run-skyvern-client-task-v1",
-            description="Use Skyvern client to run a v1 task. It is usually used for the simple tasks. This function won't return until the task is finished.",
-            fn_schema=TaskV1Request,
-        ),
-        "queue_task_v1": ToolMetadata(
-            name="queue-skyvern-client-task-v1",
-            description="Use Skyvern client to queue a v1 task. It is usually used for the simple tasks. This function will return immediately and the task will be running in the background.",
-            fn_schema=TaskV1Request,
-        ),
-        "get_task_v1": ToolMetadata(
-            name="get-skyvern-client-task-v1",
-            description="Use Skyvern client to get a v1 task. v1 tasks are usually simple tasks.",
-            fn_schema=GetTaskInput,
-        ),
-        "run_task_v2": ToolMetadata(
-            name="run-skyvern-client-task-v2",
-            description="Use Skyvern client to run a v2 task. It is usually used for the complicated tasks. This function won't return until the task is finished.",
-            fn_schema=TaskV2Request,
-        ),
-        "queue_task_v2": ToolMetadata(
-            name="queue-skyvern-client-task-v2",
-            description="Use Skyvern client to queue a v2 task. It is usually used for the complicated tasks. This function will return immediately and the task will be running in the background.",
-            fn_schema=TaskV2Request,
-        ),
-        "get_task_v2": ToolMetadata(
-            name="get-skyvern-client-task-v2",
-            description="Use Skyvern client to get a v2 task. It is usually used for the complicated tasks.",
-            fn_schema=GetTaskInput,
-        ),
+    spec_metadata: Dict[str, Dict[str, ToolMetadata]] = {
+        "TaskV1": {
+            "run_task": ToolMetadata(
+                name="run-skyvern-client-task",
+                description="Use Skyvern client to run a task. This function won't return until the task is finished.",
+                fn_schema=TaskV1Request,
+            ),
+            "dispatch_task": ToolMetadata(
+                name="dispatch-skyvern-client-task",
+                description="Use Skyvern client to dispatch a task. This function will return immediately and the task will be running in the background.",
+                fn_schema=TaskV1Request,
+            ),
+            "get_task": ToolMetadata(
+                name="get-skyvern-client-task",
+                description="Use Skyvern client to get a task.",
+                fn_schema=GetTaskInput,
+            ),
+        },
+        "TaskV2": {
+            "run_task": ToolMetadata(
+                name="run-skyvern-client-task",
+                description="Use Skyvern client to run a task. This function won't return until the task is finished.",
+                fn_schema=TaskV2Request,
+            ),
+            "dispatch_task": ToolMetadata(
+                name="dispatch-skyvern-client-task",
+                description="Use Skyvern client to dispatch a task. This function will return immediately and the task will be running in the background.",
+                fn_schema=TaskV2Request,
+            ),
+            "get_task": ToolMetadata(
+                name="get-skyvern-client-task",
+                description="Use Skyvern client to get a task.",
+                fn_schema=GetTaskInput,
+            ),
+        },
     }
 
-    def __init__(self, credential: str, base_url: str = "https://api.skyvern.com"):
+    def __init__(
+        self,
+        credential: str,
+        *,
+        base_url: str = "https://api.skyvern.com",
+        engine: Literal["TaskV1", "TaskV2"] = "TaskV2",
+    ):
         httpx_client = AsyncClient(
             headers={
                 "Content-Type": "application/json",
                 "x-api-key": credential,
             },
         )
+        self.engine = engine
         self.client = AsyncSkyvern(base_url=base_url, httpx_client=httpx_client)
 
     def get_metadata_from_fn_name(
@@ -69,7 +77,25 @@ class SkyvernClientToolSpec(BaseToolSpec):
         except AttributeError:
             return None
 
-        return self.spec_metadata.get(fn_name)
+        return self.spec_metadata.get(self.engine, {}).get(fn_name)
+
+    async def run_task(self, **kwargs: Dict[str, Any]) -> TaskResponse | Dict[str, Any | None]:
+        if self.engine == "TaskV1":
+            return await self.run_task_v1(**kwargs)
+        else:
+            return await self.run_task_v2(**kwargs)
+
+    async def dispatch_task(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse | Dict[str, Any | None]:
+        if self.engine == "TaskV1":
+            return await self.dispatch_task_v1(**kwargs)
+        else:
+            return await self.dispatch_task_v2(**kwargs)
+
+    async def get_task(self, task_id: str) -> TaskResponse | Dict[str, Any | None]:
+        if self.engine == "TaskV1":
+            return await self.get_task_v1(task_id)
+        else:
+            return await self.get_task_v2(task_id)
 
     async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
         task_request = TaskV1Request(**kwargs)
@@ -92,7 +118,7 @@ class SkyvernClientToolSpec(BaseToolSpec):
             browser_session_id=task_request.browser_session_id,
         )
 
-    async def queue_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
+    async def dispatch_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
         task_request = TaskV1Request(**kwargs)
         return await self.client.agent.create_task(
             max_steps_override=task_request.max_steps,
@@ -129,7 +155,7 @@ class SkyvernClientToolSpec(BaseToolSpec):
             proxy_location=task_request.proxy_location,
         )
 
-    async def queue_task_v2(self, **kwargs: Dict[str, Any]) -> Dict[str, Any | None]:
+    async def dispatch_task_v2(self, **kwargs: Dict[str, Any]) -> Dict[str, Any | None]:
         task_request = TaskV2Request(**kwargs)
         return await self.client.agent.observer_task_v_2(
             max_iterations_override=task_request.max_iterations,

skyvern_llamaindex-0.0.1/skyvern_llamaindex/agent.py

@@ -1,95 +0,0 @@
-from typing import Any, Dict, List, Tuple
-
-from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
-from llama_index.core.tools.types import ToolMetadata
-from skyvern_llamaindex.schema import GetTaskInput, TaskV1Request, TaskV2Request
-
-from skyvern.agent import Agent
-from skyvern.forge.sdk.schemas.observers import ObserverTask
-from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
-
-
-class SkyvernAgentToolSpec(BaseToolSpec):
-    spec_functions: List[SPEC_FUNCTION_TYPE] = [
-        "run_task_v1",
-        "queue_task_v1",
-        "get_task_v1",
-        "run_task_v2",
-        "queue_task_v2",
-        "get_task_v2",
-    ]
-    spec_metadata: Dict[str, ToolMetadata] = {
-        "run_task_v1": ToolMetadata(
-            name="run-skyvern-agent-task-v1",
-            description="Use Skyvern agent to run a v1 task. It is usually used for the simple tasks. This function won't return until the task is finished.",
-            fn_schema=TaskV1Request,
-        ),
-        "queue_task_v1": ToolMetadata(
-            name="queue-skyvern-agent-task-v1",
-            description="Use Skyvern agent to queue a v1 task. It is usually used for the simple tasks. This function will return immediately and the task will be running in the background.",
-            fn_schema=TaskV1Request,
-        ),
-        "get_task_v1": ToolMetadata(
-            name="get-skyvern-agent-task-v1",
-            description="Use Skyvern agent to get a v1 task. v1 tasks are usually simple tasks.",
-            fn_schema=GetTaskInput,
-        ),
-        "run_task_v2": ToolMetadata(
-            name="run-skyvern-agent-task-v2",
-            description="Use Skyvern agent to run a v2 task. It is usually used for the complicated tasks. This function won't return until the task is finished.",
-            fn_schema=TaskV2Request,
-        ),
-        "queue_task_v2": ToolMetadata(
-            name="queue-skyvern-agent-task-v2",
-            description="Use Skyvern agent to queue a v2 task. It is usually used for the complicated tasks. This function will return immediately and the task will be running in the background.",
-            fn_schema=TaskV2Request,
-        ),
-        "get_task_v2": ToolMetadata(
-            name="get-skyvern-agent-task-v2",
-            description="Use Skyvern agent to get a v2 task. v2 tasks are usually complicated tasks.",
-            fn_schema=GetTaskInput,
-        ),
-    }
-
-    def __init__(self) -> None:
-        self.agent = Agent()
-
-    def get_metadata_from_fn_name(
-        self, fn_name: str, spec_functions: List[str | Tuple[str, str]] | None = None
-    ) -> ToolMetadata | None:
-        try:
-            getattr(self, fn_name)
-        except AttributeError:
-            return None
-
-        return self.spec_metadata.get(fn_name)
-
-    async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
-        """Use Skyvern agent to run a v1 task. It is usually used for the simple tasks. This function won't return until the task is finished."""
-        task_request = TaskV1Request(**kwargs)
-        return await self.agent.run_task(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
-
-    async def queue_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
-        """Use Skyvern agent to queue a v1 task. It is usually used for the simple tasks. This function will return immediately and the task will be running in the background."""
-        task_request = TaskV1Request(**kwargs)
-        return await self.agent.create_task(task_request=task_request)
-
-    async def get_task_v1(self, task_id: str) -> TaskResponse | None:
-        """Use Skyvern agent to get a v1 task. v1 tasks are usually simple tasks."""
-        return await self.agent.get_task(task_id=task_id)
-
-    async def run_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
-        """Use Skyvern agent to run a v2 task. It is usually used for the complicated tasks. This function won't return until the task is finished."""
-        task_request = TaskV2Request(**kwargs)
-        return await self.agent.run_observer_task_v_2(
-            task_request=task_request, timeout_seconds=task_request.timeout_seconds
-        )
-
-    async def queue_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
-        """Use Skyvern agent to queue a v2 task. It is usually used for the complicated tasks. This function will return immediately and the task will be running in the background."""
-        task_request = TaskV2Request(**kwargs)
-        return await self.agent.observer_task_v_2(task_request=task_request)
-
-    async def get_task_v2(self, task_id: str) -> ObserverTask | None:
-        """Use Skyvern agent to get a v2 task. v2 tasks are usually complicated tasks."""
-        return await self.agent.get_observer_task_v_2(task_id=task_id)