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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: skyvern-llamaindex
3
- Version: 0.0.1
3
+ Version: 0.0.2
4
4
  Summary: Skyvern integration for LlamaIndex
5
5
  Author: lawyzheng
6
6
  Author-email: lawy@skyvern.com
@@ -19,9 +19,9 @@ Description-Content-Type: text/markdown
19
19
  - [Installation](#installation)
20
20
  - [Usage](#usage)
21
21
  - [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)
22
- - [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)
22
+ - [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)
23
23
  - [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)
24
- - [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)
24
+ - [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)
25
25
 
26
26
  <!-- END doctoc generated TOC please keep comment here to allow auto update -->
27
27
 
@@ -48,14 +48,14 @@ import asyncio
48
48
  from dotenv import load_dotenv
49
49
  from llama_index.agent.openai import OpenAIAgent
50
50
  from llama_index.llms.openai import OpenAI
51
- from skyvern_llamaindex.agent import SkyvernAgentToolSpec
51
+ from skyvern_llamaindex.agent import SkyvernToolSpec
52
52
 
53
53
  # load OpenAI API key from .env
54
54
  load_dotenv()
55
55
 
56
- skyvern_tool = SkyvernAgentToolSpec()
56
+ skyvern_tool = SkyvernToolSpec()
57
57
 
58
- tools = skyvern_tool.to_tool_list(["run_task_v2"])
58
+ tools = skyvern_tool.to_tool_list(["run_task"])
59
59
 
60
60
  agent = OpenAIAgent.from_tools(
61
61
  tools=tools,
@@ -69,8 +69,8 @@ response = agent.chat("Run the task with skyvern. The task is about 'Navigate to
69
69
  print(response)
70
70
  ```
71
71
 
72
- ### Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
73
- > 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.
72
+ ### Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
73
+ > 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.
74
74
 
75
75
  :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.
76
76
 
@@ -80,7 +80,7 @@ from dotenv import load_dotenv
80
80
  from llama_index.agent.openai import OpenAIAgent
81
81
  from llama_index.llms.openai import OpenAI
82
82
  from llama_index.core.tools import FunctionTool
83
- from skyvern_llamaindex.agent import SkyvernAgentToolSpec
83
+ from skyvern_llamaindex.agent import SkyvernToolSpec
84
84
 
85
85
  async def sleep(seconds: int) -> str:
86
86
  await asyncio.sleep(seconds)
@@ -89,7 +89,7 @@ async def sleep(seconds: int) -> str:
89
89
  # load OpenAI API key from .env
90
90
  load_dotenv()
91
91
 
92
- skyvern_tool = SkyvernAgentToolSpec()
92
+ skyvern_tool = SkyvernToolSpec()
93
93
 
94
94
  sleep_tool = FunctionTool.from_defaults(
95
95
  async_fn=sleep,
@@ -97,7 +97,7 @@ sleep_tool = FunctionTool.from_defaults(
97
97
  name="sleep",
98
98
  )
99
99
 
100
- tools = skyvern_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
100
+ tools = skyvern_tool.to_tool_list(["dispatch_task", "get_task"])
101
101
  tools.append(sleep_tool)
102
102
 
103
103
  agent = OpenAIAgent.from_tools(
@@ -107,7 +107,7 @@ agent = OpenAIAgent.from_tools(
107
107
  max_function_calls=10,
108
108
  )
109
109
 
110
- 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.")
110
+ 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.")
111
111
  print(response)
112
112
 
113
113
  ```
@@ -122,7 +122,7 @@ import asyncio
122
122
  from dotenv import load_dotenv
123
123
  from llama_index.agent.openai import OpenAIAgent
124
124
  from llama_index.llms.openai import OpenAI
125
- from skyvern_llamaindex.client import SkyvernClientToolSpec
125
+ from skyvern_llamaindex.client import SkyvernToolSpec
126
126
 
127
127
 
128
128
  async def sleep(seconds: int) -> str:
@@ -132,11 +132,11 @@ async def sleep(seconds: int) -> str:
132
132
  # load OpenAI API key from .env
133
133
  load_dotenv()
134
134
 
135
- skyvern_client_tool = SkyvernClientToolSpec(
135
+ skyvern_client_tool = SkyvernToolSpec(
136
136
  credential="<your_organization_api_key>",
137
137
  )
138
138
 
139
- tools = skyvern_client_tool.to_tool_list(["run_task_v2"])
139
+ tools = skyvern_client_tool.to_tool_list(["run_task"])
140
140
 
141
141
  agent = OpenAIAgent.from_tools(
142
142
  tools=tools,
@@ -150,8 +150,8 @@ print(response)
150
150
 
151
151
  ```
152
152
 
153
- ### Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
154
- > 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.
153
+ ### Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
154
+ > 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.
155
155
 
156
156
  no need to run `skyvern init` command in your terminal to set up skyvern before using this integration.
157
157
 
@@ -161,7 +161,7 @@ from dotenv import load_dotenv
161
161
  from llama_index.agent.openai import OpenAIAgent
162
162
  from llama_index.llms.openai import OpenAI
163
163
  from llama_index.core.tools import FunctionTool
164
- from skyvern_llamaindex.client import SkyvernClientToolSpec
164
+ from skyvern_llamaindex.client import SkyvernToolSpec
165
165
 
166
166
 
167
167
  async def sleep(seconds: int) -> str:
@@ -171,7 +171,7 @@ async def sleep(seconds: int) -> str:
171
171
  # load OpenAI API key from .env
172
172
  load_dotenv()
173
173
 
174
- skyvern_client_tool = SkyvernClientToolSpec(
174
+ skyvern_client_tool = SkyvernToolSpec(
175
175
  credential="<your_organization_api_key>",
176
176
  )
177
177
 
@@ -181,7 +181,7 @@ sleep_tool = FunctionTool.from_defaults(
181
181
  name="sleep",
182
182
  )
183
183
 
184
- tools = skyvern_client_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
184
+ tools = skyvern_client_tool.to_tool_list(["dispatch_task", "get_task"])
185
185
  tools.append(sleep_tool)
186
186
 
187
187
  agent = OpenAIAgent.from_tools(
@@ -191,7 +191,7 @@ agent = OpenAIAgent.from_tools(
191
191
  max_function_calls=10,
192
192
  )
193
193
 
194
- 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.")
194
+ 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.")
195
195
  print(response)
196
196
 
197
197
  ```
@@ -6,9 +6,9 @@
6
6
  - [Installation](#installation)
7
7
  - [Usage](#usage)
8
8
  - [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)
9
- - [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)
9
+ - [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)
10
10
  - [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)
11
- - [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)
11
+ - [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)
12
12
 
13
13
  <!-- END doctoc generated TOC please keep comment here to allow auto update -->
14
14
 
@@ -35,14 +35,14 @@ import asyncio
35
35
  from dotenv import load_dotenv
36
36
  from llama_index.agent.openai import OpenAIAgent
37
37
  from llama_index.llms.openai import OpenAI
38
- from skyvern_llamaindex.agent import SkyvernAgentToolSpec
38
+ from skyvern_llamaindex.agent import SkyvernToolSpec
39
39
 
40
40
  # load OpenAI API key from .env
41
41
  load_dotenv()
42
42
 
43
- skyvern_tool = SkyvernAgentToolSpec()
43
+ skyvern_tool = SkyvernToolSpec()
44
44
 
45
- tools = skyvern_tool.to_tool_list(["run_task_v2"])
45
+ tools = skyvern_tool.to_tool_list(["run_task"])
46
46
 
47
47
  agent = OpenAIAgent.from_tools(
48
48
  tools=tools,
@@ -56,8 +56,8 @@ response = agent.chat("Run the task with skyvern. The task is about 'Navigate to
56
56
  print(response)
57
57
  ```
58
58
 
59
- ### Run a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
60
- > 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.
59
+ ### Dispatch a task(async) with skyvern agent (calling skyvern agent function directly in the tool)
60
+ > 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.
61
61
 
62
62
  :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.
63
63
 
@@ -67,7 +67,7 @@ from dotenv import load_dotenv
67
67
  from llama_index.agent.openai import OpenAIAgent
68
68
  from llama_index.llms.openai import OpenAI
69
69
  from llama_index.core.tools import FunctionTool
70
- from skyvern_llamaindex.agent import SkyvernAgentToolSpec
70
+ from skyvern_llamaindex.agent import SkyvernToolSpec
71
71
 
72
72
  async def sleep(seconds: int) -> str:
73
73
  await asyncio.sleep(seconds)
@@ -76,7 +76,7 @@ async def sleep(seconds: int) -> str:
76
76
  # load OpenAI API key from .env
77
77
  load_dotenv()
78
78
 
79
- skyvern_tool = SkyvernAgentToolSpec()
79
+ skyvern_tool = SkyvernToolSpec()
80
80
 
81
81
  sleep_tool = FunctionTool.from_defaults(
82
82
  async_fn=sleep,
@@ -84,7 +84,7 @@ sleep_tool = FunctionTool.from_defaults(
84
84
  name="sleep",
85
85
  )
86
86
 
87
- tools = skyvern_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
87
+ tools = skyvern_tool.to_tool_list(["dispatch_task", "get_task"])
88
88
  tools.append(sleep_tool)
89
89
 
90
90
  agent = OpenAIAgent.from_tools(
@@ -94,7 +94,7 @@ agent = OpenAIAgent.from_tools(
94
94
  max_function_calls=10,
95
95
  )
96
96
 
97
- 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.")
97
+ 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.")
98
98
  print(response)
99
99
 
100
100
  ```
@@ -109,7 +109,7 @@ import asyncio
109
109
  from dotenv import load_dotenv
110
110
  from llama_index.agent.openai import OpenAIAgent
111
111
  from llama_index.llms.openai import OpenAI
112
- from skyvern_llamaindex.client import SkyvernClientToolSpec
112
+ from skyvern_llamaindex.client import SkyvernToolSpec
113
113
 
114
114
 
115
115
  async def sleep(seconds: int) -> str:
@@ -119,11 +119,11 @@ async def sleep(seconds: int) -> str:
119
119
  # load OpenAI API key from .env
120
120
  load_dotenv()
121
121
 
122
- skyvern_client_tool = SkyvernClientToolSpec(
122
+ skyvern_client_tool = SkyvernToolSpec(
123
123
  credential="<your_organization_api_key>",
124
124
  )
125
125
 
126
- tools = skyvern_client_tool.to_tool_list(["run_task_v2"])
126
+ tools = skyvern_client_tool.to_tool_list(["run_task"])
127
127
 
128
128
  agent = OpenAIAgent.from_tools(
129
129
  tools=tools,
@@ -137,8 +137,8 @@ print(response)
137
137
 
138
138
  ```
139
139
 
140
- ### Run a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
141
- > 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.
140
+ ### Dispatch a task(async) with skyvern client (calling skyvern OpenAPI in the tool)
141
+ > 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.
142
142
 
143
143
  no need to run `skyvern init` command in your terminal to set up skyvern before using this integration.
144
144
 
@@ -148,7 +148,7 @@ from dotenv import load_dotenv
148
148
  from llama_index.agent.openai import OpenAIAgent
149
149
  from llama_index.llms.openai import OpenAI
150
150
  from llama_index.core.tools import FunctionTool
151
- from skyvern_llamaindex.client import SkyvernClientToolSpec
151
+ from skyvern_llamaindex.client import SkyvernToolSpec
152
152
 
153
153
 
154
154
  async def sleep(seconds: int) -> str:
@@ -158,7 +158,7 @@ async def sleep(seconds: int) -> str:
158
158
  # load OpenAI API key from .env
159
159
  load_dotenv()
160
160
 
161
- skyvern_client_tool = SkyvernClientToolSpec(
161
+ skyvern_client_tool = SkyvernToolSpec(
162
162
  credential="<your_organization_api_key>",
163
163
  )
164
164
 
@@ -168,7 +168,7 @@ sleep_tool = FunctionTool.from_defaults(
168
168
  name="sleep",
169
169
  )
170
170
 
171
- tools = skyvern_client_tool.to_tool_list(["queue_task_v2", "get_task_v2"])
171
+ tools = skyvern_client_tool.to_tool_list(["dispatch_task", "get_task"])
172
172
  tools.append(sleep_tool)
173
173
 
174
174
  agent = OpenAIAgent.from_tools(
@@ -178,7 +178,7 @@ agent = OpenAIAgent.from_tools(
178
178
  max_function_calls=10,
179
179
  )
180
180
 
181
- 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.")
181
+ 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.")
182
182
  print(response)
183
183
 
184
184
  ```
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "skyvern-llamaindex"
3
- version = "0.0.1"
3
+ version = "0.0.2"
4
4
  description = "Skyvern integration for LlamaIndex"
5
5
  authors = ["lawyzheng <lawy@skyvern.com>"]
6
6
  packages = [{ include = "skyvern_llamaindex" }]
@@ -0,0 +1,109 @@
1
+ from typing import Any, Dict, List, Literal, Tuple
2
+
3
+ from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
4
+ from llama_index.core.tools.types import ToolMetadata
5
+ from skyvern_llamaindex.schema import GetTaskInput, TaskV1Request, TaskV2Request
6
+
7
+ from skyvern.agent import Agent
8
+ from skyvern.forge.sdk.schemas.observers import ObserverTask
9
+ from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
10
+
11
+
12
+ class SkyvernToolSpec(BaseToolSpec):
13
+ spec_functions: List[SPEC_FUNCTION_TYPE] = [
14
+ "run_task",
15
+ "dispatch_task",
16
+ "get_task",
17
+ ]
18
+ spec_metadata: Dict[str, Dict[str, ToolMetadata]] = {
19
+ "TaskV1": {
20
+ "run_task": ToolMetadata(
21
+ name="run-skyvern-agent-task",
22
+ description="Use Skyvern agent to run a task. This function won't return until the task is finished.",
23
+ fn_schema=TaskV1Request,
24
+ ),
25
+ "dispatch_task": ToolMetadata(
26
+ name="dispatch-skyvern-agent-task",
27
+ description="Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.",
28
+ fn_schema=TaskV1Request,
29
+ ),
30
+ "get_task": ToolMetadata(
31
+ name="get-skyvern-agent-task",
32
+ description="Use Skyvern agent to get a task.",
33
+ fn_schema=GetTaskInput,
34
+ ),
35
+ },
36
+ "TaskV2": {
37
+ "run_task": ToolMetadata(
38
+ name="run-skyvern-agent-task",
39
+ description="Use Skyvern agent to run a task. This function won't return until the task is finished.",
40
+ fn_schema=TaskV2Request,
41
+ ),
42
+ "dispatch_task": ToolMetadata(
43
+ name="dispatch-skyvern-agent-task",
44
+ description="Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.",
45
+ fn_schema=TaskV2Request,
46
+ ),
47
+ "get_task": ToolMetadata(
48
+ name="get-skyvern-agent-task",
49
+ description="Use Skyvern agent to get a task.",
50
+ fn_schema=GetTaskInput,
51
+ ),
52
+ },
53
+ }
54
+
55
+ def __init__(self, *, engine: Literal["TaskV1", "TaskV2"] = "TaskV2") -> None:
56
+ self.agent = Agent()
57
+ self.engine = engine
58
+
59
+ def get_metadata_from_fn_name(
60
+ self, fn_name: str, spec_functions: List[str | Tuple[str, str]] | None = None
61
+ ) -> ToolMetadata | None:
62
+ try:
63
+ getattr(self, fn_name)
64
+ except AttributeError:
65
+ return None
66
+
67
+ return self.spec_metadata.get(self.engine, {}).get(fn_name)
68
+
69
+ async def run_task(self, **kwargs: Dict[str, Any]) -> TaskResponse | ObserverTask:
70
+ if self.engine == "TaskV1":
71
+ return await self.run_task_v1(**kwargs)
72
+ else:
73
+ return await self.run_task_v2(**kwargs)
74
+
75
+ async def dispatch_task(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse | ObserverTask:
76
+ if self.engine == "TaskV1":
77
+ return await self.dispatch_task_v1(**kwargs)
78
+ else:
79
+ return await self.dispatch_task_v2(**kwargs)
80
+
81
+ async def get_task(self, task_id: str) -> TaskResponse | ObserverTask | None:
82
+ if self.engine == "TaskV1":
83
+ return await self.get_task_v1(task_id)
84
+ else:
85
+ return await self.get_task_v2(task_id)
86
+
87
+ async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
88
+ task_request = TaskV1Request(**kwargs)
89
+ return await self.agent.run_task(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
90
+
91
+ async def dispatch_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
92
+ task_request = TaskV1Request(**kwargs)
93
+ return await self.agent.create_task(task_request=task_request)
94
+
95
+ async def get_task_v1(self, task_id: str) -> TaskResponse | None:
96
+ return await self.agent.get_task(task_id=task_id)
97
+
98
+ async def run_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
99
+ task_request = TaskV2Request(**kwargs)
100
+ return await self.agent.run_observer_task_v_2(
101
+ task_request=task_request, timeout_seconds=task_request.timeout_seconds
102
+ )
103
+
104
+ async def dispatch_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
105
+ task_request = TaskV2Request(**kwargs)
106
+ return await self.agent.observer_task_v_2(task_request=task_request)
107
+
108
+ async def get_task_v2(self, task_id: str) -> ObserverTask | None:
109
+ return await self.agent.get_observer_task_v_2(task_id=task_id)
@@ -1,4 +1,4 @@
1
- from typing import Any, Dict, List, Tuple
1
+ from typing import Any, Dict, List, Literal, Tuple
2
2
 
3
3
  from httpx import AsyncClient
4
4
  from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
@@ -9,56 +9,64 @@ from skyvern.client import AsyncSkyvern
9
9
  from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
10
10
 
11
11
 
12
- class SkyvernClientToolSpec(BaseToolSpec):
12
+ class SkyvernToolSpec(BaseToolSpec):
13
13
  spec_functions: List[SPEC_FUNCTION_TYPE] = [
14
- "run_task_v1",
15
- "queue_task_v1",
16
- "get_task_v1",
17
- "run_task_v2",
18
- "queue_task_v2",
19
- "get_task_v2",
14
+ "run_task",
15
+ "dispatch_task",
16
+ "get_task",
20
17
  ]
21
18
 
22
- spec_metadata: Dict[str, ToolMetadata] = {
23
- "run_task_v1": ToolMetadata(
24
- name="run-skyvern-client-task-v1",
25
- 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.",
26
- fn_schema=TaskV1Request,
27
- ),
28
- "queue_task_v1": ToolMetadata(
29
- name="queue-skyvern-client-task-v1",
30
- 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.",
31
- fn_schema=TaskV1Request,
32
- ),
33
- "get_task_v1": ToolMetadata(
34
- name="get-skyvern-client-task-v1",
35
- description="Use Skyvern client to get a v1 task. v1 tasks are usually simple tasks.",
36
- fn_schema=GetTaskInput,
37
- ),
38
- "run_task_v2": ToolMetadata(
39
- name="run-skyvern-client-task-v2",
40
- 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.",
41
- fn_schema=TaskV2Request,
42
- ),
43
- "queue_task_v2": ToolMetadata(
44
- name="queue-skyvern-client-task-v2",
45
- 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.",
46
- fn_schema=TaskV2Request,
47
- ),
48
- "get_task_v2": ToolMetadata(
49
- name="get-skyvern-client-task-v2",
50
- description="Use Skyvern client to get a v2 task. It is usually used for the complicated tasks.",
51
- fn_schema=GetTaskInput,
52
- ),
19
+ spec_metadata: Dict[str, Dict[str, ToolMetadata]] = {
20
+ "TaskV1": {
21
+ "run_task": ToolMetadata(
22
+ name="run-skyvern-client-task",
23
+ description="Use Skyvern client to run a task. This function won't return until the task is finished.",
24
+ fn_schema=TaskV1Request,
25
+ ),
26
+ "dispatch_task": ToolMetadata(
27
+ name="dispatch-skyvern-client-task",
28
+ description="Use Skyvern client to dispatch a task. This function will return immediately and the task will be running in the background.",
29
+ fn_schema=TaskV1Request,
30
+ ),
31
+ "get_task": ToolMetadata(
32
+ name="get-skyvern-client-task",
33
+ description="Use Skyvern client to get a task.",
34
+ fn_schema=GetTaskInput,
35
+ ),
36
+ },
37
+ "TaskV2": {
38
+ "run_task": ToolMetadata(
39
+ name="run-skyvern-client-task",
40
+ description="Use Skyvern client to run a task. This function won't return until the task is finished.",
41
+ fn_schema=TaskV2Request,
42
+ ),
43
+ "dispatch_task": ToolMetadata(
44
+ name="dispatch-skyvern-client-task",
45
+ description="Use Skyvern client to dispatch a task. This function will return immediately and the task will be running in the background.",
46
+ fn_schema=TaskV2Request,
47
+ ),
48
+ "get_task": ToolMetadata(
49
+ name="get-skyvern-client-task",
50
+ description="Use Skyvern client to get a task.",
51
+ fn_schema=GetTaskInput,
52
+ ),
53
+ },
53
54
  }
54
55
 
55
- def __init__(self, credential: str, base_url: str = "https://api.skyvern.com"):
56
+ def __init__(
57
+ self,
58
+ credential: str,
59
+ *,
60
+ base_url: str = "https://api.skyvern.com",
61
+ engine: Literal["TaskV1", "TaskV2"] = "TaskV2",
62
+ ):
56
63
  httpx_client = AsyncClient(
57
64
  headers={
58
65
  "Content-Type": "application/json",
59
66
  "x-api-key": credential,
60
67
  },
61
68
  )
69
+ self.engine = engine
62
70
  self.client = AsyncSkyvern(base_url=base_url, httpx_client=httpx_client)
63
71
 
64
72
  def get_metadata_from_fn_name(
@@ -69,7 +77,25 @@ class SkyvernClientToolSpec(BaseToolSpec):
69
77
  except AttributeError:
70
78
  return None
71
79
 
72
- return self.spec_metadata.get(fn_name)
80
+ return self.spec_metadata.get(self.engine, {}).get(fn_name)
81
+
82
+ async def run_task(self, **kwargs: Dict[str, Any]) -> TaskResponse | Dict[str, Any | None]:
83
+ if self.engine == "TaskV1":
84
+ return await self.run_task_v1(**kwargs)
85
+ else:
86
+ return await self.run_task_v2(**kwargs)
87
+
88
+ async def dispatch_task(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse | Dict[str, Any | None]:
89
+ if self.engine == "TaskV1":
90
+ return await self.dispatch_task_v1(**kwargs)
91
+ else:
92
+ return await self.dispatch_task_v2(**kwargs)
93
+
94
+ async def get_task(self, task_id: str) -> TaskResponse | Dict[str, Any | None]:
95
+ if self.engine == "TaskV1":
96
+ return await self.get_task_v1(task_id)
97
+ else:
98
+ return await self.get_task_v2(task_id)
73
99
 
74
100
  async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
75
101
  task_request = TaskV1Request(**kwargs)
@@ -92,7 +118,7 @@ class SkyvernClientToolSpec(BaseToolSpec):
92
118
  browser_session_id=task_request.browser_session_id,
93
119
  )
94
120
 
95
- async def queue_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
121
+ async def dispatch_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
96
122
  task_request = TaskV1Request(**kwargs)
97
123
  return await self.client.agent.create_task(
98
124
  max_steps_override=task_request.max_steps,
@@ -129,7 +155,7 @@ class SkyvernClientToolSpec(BaseToolSpec):
129
155
  proxy_location=task_request.proxy_location,
130
156
  )
131
157
 
132
- async def queue_task_v2(self, **kwargs: Dict[str, Any]) -> Dict[str, Any | None]:
158
+ async def dispatch_task_v2(self, **kwargs: Dict[str, Any]) -> Dict[str, Any | None]:
133
159
  task_request = TaskV2Request(**kwargs)
134
160
  return await self.client.agent.observer_task_v_2(
135
161
  max_iterations_override=task_request.max_iterations,
File without changes
@@ -1,95 +0,0 @@
1
- from typing import Any, Dict, List, Tuple
2
-
3
- from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
4
- from llama_index.core.tools.types import ToolMetadata
5
- from skyvern_llamaindex.schema import GetTaskInput, TaskV1Request, TaskV2Request
6
-
7
- from skyvern.agent import Agent
8
- from skyvern.forge.sdk.schemas.observers import ObserverTask
9
- from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
10
-
11
-
12
- class SkyvernAgentToolSpec(BaseToolSpec):
13
- spec_functions: List[SPEC_FUNCTION_TYPE] = [
14
- "run_task_v1",
15
- "queue_task_v1",
16
- "get_task_v1",
17
- "run_task_v2",
18
- "queue_task_v2",
19
- "get_task_v2",
20
- ]
21
- spec_metadata: Dict[str, ToolMetadata] = {
22
- "run_task_v1": ToolMetadata(
23
- name="run-skyvern-agent-task-v1",
24
- 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.",
25
- fn_schema=TaskV1Request,
26
- ),
27
- "queue_task_v1": ToolMetadata(
28
- name="queue-skyvern-agent-task-v1",
29
- 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.",
30
- fn_schema=TaskV1Request,
31
- ),
32
- "get_task_v1": ToolMetadata(
33
- name="get-skyvern-agent-task-v1",
34
- description="Use Skyvern agent to get a v1 task. v1 tasks are usually simple tasks.",
35
- fn_schema=GetTaskInput,
36
- ),
37
- "run_task_v2": ToolMetadata(
38
- name="run-skyvern-agent-task-v2",
39
- 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.",
40
- fn_schema=TaskV2Request,
41
- ),
42
- "queue_task_v2": ToolMetadata(
43
- name="queue-skyvern-agent-task-v2",
44
- 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.",
45
- fn_schema=TaskV2Request,
46
- ),
47
- "get_task_v2": ToolMetadata(
48
- name="get-skyvern-agent-task-v2",
49
- description="Use Skyvern agent to get a v2 task. v2 tasks are usually complicated tasks.",
50
- fn_schema=GetTaskInput,
51
- ),
52
- }
53
-
54
- def __init__(self) -> None:
55
- self.agent = Agent()
56
-
57
- def get_metadata_from_fn_name(
58
- self, fn_name: str, spec_functions: List[str | Tuple[str, str]] | None = None
59
- ) -> ToolMetadata | None:
60
- try:
61
- getattr(self, fn_name)
62
- except AttributeError:
63
- return None
64
-
65
- return self.spec_metadata.get(fn_name)
66
-
67
- async def run_task_v1(self, **kwargs: Dict[str, Any]) -> TaskResponse:
68
- """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."""
69
- task_request = TaskV1Request(**kwargs)
70
- return await self.agent.run_task(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
71
-
72
- async def queue_task_v1(self, **kwargs: Dict[str, Any]) -> CreateTaskResponse:
73
- """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."""
74
- task_request = TaskV1Request(**kwargs)
75
- return await self.agent.create_task(task_request=task_request)
76
-
77
- async def get_task_v1(self, task_id: str) -> TaskResponse | None:
78
- """Use Skyvern agent to get a v1 task. v1 tasks are usually simple tasks."""
79
- return await self.agent.get_task(task_id=task_id)
80
-
81
- async def run_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
82
- """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."""
83
- task_request = TaskV2Request(**kwargs)
84
- return await self.agent.run_observer_task_v_2(
85
- task_request=task_request, timeout_seconds=task_request.timeout_seconds
86
- )
87
-
88
- async def queue_task_v2(self, **kwargs: Dict[str, Any]) -> ObserverTask:
89
- """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."""
90
- task_request = TaskV2Request(**kwargs)
91
- return await self.agent.observer_task_v_2(task_request=task_request)
92
-
93
- async def get_task_v2(self, task_id: str) -> ObserverTask | None:
94
- """Use Skyvern agent to get a v2 task. v2 tasks are usually complicated tasks."""
95
- return await self.agent.get_observer_task_v_2(task_id=task_id)