skyvern-llamaindex 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

skyvern_llamaindex/agent.py

@@ -1,23 +1,19 @@
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 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.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.tasks import CreateTaskResponse, TaskRequest, TaskResponse
-
-default_agent = Agent()
+from skyvern import Skyvern
+from skyvern.client.agent.types.agent_get_run_response import AgentGetRunResponse
+from skyvern.client.types.task_run_response import TaskRunResponse
+from skyvern.schemas.runs import RunEngine
 
 
 class SkyvernTool:
-    def __init__(self, agent: Optional[Agent] = None):
+    def __init__(self, agent: Optional[Skyvern] = None):
         if agent is None:
-            agent = default_agent
+            agent = Skyvern(base_url=None, api_key=None)
         self.agent = agent
 
     def run_task(self) -> FunctionTool:
@@ -43,23 +39,17 @@ class SkyvernTaskToolSpec(BaseToolSpec):
     def __init__(
         self,
         *,
-        agent: Optional[Agent] = None,
-        engine: Literal["TaskV1", "TaskV2"] = settings.engine,
+        agent: Skyvern | None = None,
+        engine: RunEngine = settings.engine,
         run_task_timeout_seconds: int = settings.run_task_timeout_seconds,
     ) -> None:
         if agent is None:
-            agent = Agent()
+            agent = Skyvern(base_url=None, api_key=None)
         self.agent = agent
         self.engine = engine
         self.run_task_timeout_seconds = run_task_timeout_seconds
 
-    # TODO: agent haven't exposed the task v1 generate function, we can migrate to use agent interface when it's available
-    async def _generate_v1_task_request(self, user_prompt: str) -> TaskGenerationBase:
-        llm_prompt = prompt_engine.load_prompt("generate-task", user_prompt=user_prompt)
-        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) -> TaskRunResponse:
         """
         Use Skyvern agent to run a task. This function won't return until the task is finished.
 
@@ -67,13 +57,15 @@ class SkyvernTaskToolSpec(BaseToolSpec):
             user_prompt[str]: The user's prompt describing the task.
             url (Optional[str]): The URL of the target website for the task.
         """
+        return await self.agent.run_task(
+            prompt=user_prompt,
+            url=url,
+            engine=self.engine,
+            timeout=self.run_task_timeout_seconds,
+            wait_for_completion=True,
+        )
 
-        if self.engine == "TaskV1":
-            return await self.run_task_v1(user_prompt=user_prompt, url=url)
-        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) -> TaskRunResponse:
         """
         Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background.
 
@@ -81,53 +73,19 @@ class SkyvernTaskToolSpec(BaseToolSpec):
             user_prompt[str]: The user's prompt describing the task.
             url (Optional[str]): The URL of the target website for the task.
         """
+        return await self.agent.run_task(
+            prompt=user_prompt,
+            url=url,
+            engine=self.engine,
+            timeout=self.run_task_timeout_seconds,
+            wait_for_completion=False,
+        )
 
-        if self.engine == "TaskV1":
-            return await self.dispatch_task_v1(user_prompt=user_prompt, url=url)
-        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) -> AgentGetRunResponse | None:
         """
         Use Skyvern agent to get a task.
 
         Args:
             task_id[str]: The id of the task.
         """
-
-        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, user_prompt: str, url: Optional[str] = None) -> TaskResponse:
-        task_generation = await self._generate_v1_task_request(user_prompt=user_prompt)
-        task_request = TaskRequest.model_validate(task_generation, from_attributes=True)
-        if url is not None:
-            task_request.url = url
-
-        return await self.agent.run_task(task_request=task_request, timeout_seconds=self.run_task_timeout_seconds)
-
-    async def dispatch_task_v1(self, user_prompt: str, url: Optional[str] = None) -> CreateTaskResponse:
-        task_generation = await self._generate_v1_task_request(user_prompt=user_prompt)
-        task_request = TaskRequest.model_validate(task_generation, from_attributes=True)
-        if url is not None:
-            task_request.url = url
-
-        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, user_prompt: str, url: Optional[str] = None) -> ObserverTask:
-        task_request = ObserverTaskRequest(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)
-        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)
+        return await self.agent.get_run(run_id=task_id)

skyvern_llamaindex/client.py

@@ -1,14 +1,14 @@
-from typing import Any, Dict, List, Literal, Optional
+from typing import List, Optional
 
-from httpx import AsyncClient
 from llama_index.core.tools import FunctionTool
 from llama_index.core.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
 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.tasks import CreateTaskResponse, TaskRequest, TaskResponse
+from skyvern import Skyvern
+from skyvern.client.agent.types.agent_get_run_response import AgentGetRunResponse
+from skyvern.client.types.task_run_response import TaskRunResponse
+from skyvern.schemas.runs import RunEngine
 
 
 class SkyvernTool(BaseModel):
@@ -52,20 +52,14 @@ class SkyvernTaskToolSpec(BaseToolSpec):
         *,
         api_key: str = settings.api_key,
         base_url: str = settings.base_url,
-        engine: Literal["TaskV1", "TaskV2"] = settings.engine,
+        engine: RunEngine = settings.engine,
         run_task_timeout_seconds: int = settings.run_task_timeout_seconds,
     ):
-        httpx_client = AsyncClient(
-            headers={
-                "Content-Type": "application/json",
-                "x-api-key": api_key,
-            },
-        )
         self.engine = engine
         self.run_task_timeout_seconds = run_task_timeout_seconds
-        self.client = AsyncSkyvern(base_url=base_url, httpx_client=httpx_client)
+        self.client = Skyvern(base_url=base_url, api_key=api_key)
 
-    async def run_task(self, user_prompt: str, url: Optional[str] = None) -> TaskResponse | Dict[str, Any | None]:
+    async def run_task(self, user_prompt: str, url: Optional[str] = None) -> TaskRunResponse:
         """
         Use Skyvern client to run a task. This function won't return until the task is finished.
 
@@ -74,14 +68,15 @@ class SkyvernTaskToolSpec(BaseToolSpec):
             url (Optional[str]): The URL of the target website for the task.
         """
 
-        if self.engine == "TaskV1":
-            return await self.run_task_v1(user_prompt=user_prompt, url=url)
-        else:
-            return await self.run_task_v2(user_prompt=user_prompt, url=url)
+        return await self.client.run_task(
+            prompt=user_prompt,
+            url=url,
+            engine=self.engine,
+            timeout=self.run_task_timeout_seconds,
+            wait_for_completion=True,
+        )
 
-    async def dispatch_task(
-        self, user_prompt: str, url: Optional[str] = None
-    ) -> CreateTaskResponse | Dict[str, Any | None]:
+    async def dispatch_task(self, user_prompt: str, url: Optional[str] = None) -> TaskRunResponse:
         """
         Use Skyvern client to dispatch a task. This function will return immediately and the task will be running in the background.
 
@@ -90,12 +85,15 @@ class SkyvernTaskToolSpec(BaseToolSpec):
             url (Optional[str]): The URL of the target website for the task.
         """
 
-        if self.engine == "TaskV1":
-            return await self.dispatch_task_v1(user_prompt=user_prompt, url=url)
-        else:
-            return await self.dispatch_task_v2(user_prompt=user_prompt, url=url)
+        return await self.client.run_task(
+            prompt=user_prompt,
+            url=url,
+            engine=self.engine,
+            timeout=self.run_task_timeout_seconds,
+            wait_for_completion=False,
+        )
 
-    async def get_task(self, task_id: str) -> TaskResponse | Dict[str, Any | None]:
+    async def get_task(self, task_id: str) -> AgentGetRunResponse | None:
         """
         Use Skyvern client to get a task.
 
@@ -103,71 +101,4 @@ class SkyvernTaskToolSpec(BaseToolSpec):
             task_id[str]: The id of the task.
         """
 
-        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, user_prompt: str, url: Optional[str] = None) -> TaskResponse:
-        task_generation = await self.client.agent.generate_task(
-            prompt=user_prompt,
-        )
-        task_request = TaskRequest.model_validate(task_generation, from_attributes=True)
-        if url is not None:
-            task_request.url = url
-
-        return await self.client.agent.run_task(
-            timeout_seconds=self.run_task_timeout_seconds,
-            url=task_request.url,
-            title=task_request.title,
-            navigation_goal=task_request.navigation_goal,
-            data_extraction_goal=task_request.data_extraction_goal,
-            navigation_payload=task_request.navigation_goal,
-            error_code_mapping=task_request.error_code_mapping,
-            extracted_information_schema=task_request.extracted_information_schema,
-            complete_criterion=task_request.complete_criterion,
-            terminate_criterion=task_request.terminate_criterion,
-        )
-
-    async def dispatch_task_v1(self, user_prompt: str, url: Optional[str] = None) -> CreateTaskResponse:
-        task_generation = await self.client.agent.generate_task(
-            prompt=user_prompt,
-        )
-        task_request = TaskRequest.model_validate(task_generation, from_attributes=True)
-        if url is not None:
-            task_request.url = url
-
-        return await self.client.agent.create_task(
-            url=task_request.url,
-            title=task_request.title,
-            navigation_goal=task_request.navigation_goal,
-            data_extraction_goal=task_request.data_extraction_goal,
-            navigation_payload=task_request.navigation_goal,
-            error_code_mapping=task_request.error_code_mapping,
-            extracted_information_schema=task_request.extracted_information_schema,
-            complete_criterion=task_request.complete_criterion,
-            terminate_criterion=task_request.terminate_criterion,
-        )
-
-    async def get_task_v1(self, task_id: str) -> TaskResponse:
-        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)
-        return await self.client.agent.run_observer_task_v_2(
-            timeout_seconds=self.run_task_timeout_seconds,
-            user_prompt=task_request.user_prompt,
-            url=task_request.url,
-            browser_session_id=task_request.browser_session_id,
-        )
-
-    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)
-        return await self.client.agent.observer_task_v_2(
-            user_prompt=task_request.user_prompt,
-            url=task_request.url,
-            browser_session_id=task_request.browser_session_id,
-        )
-
-    async def get_task_v2(self, task_id: str) -> Dict[str, Any | None]:
-        return await self.client.agent.get_observer_task_v_2(task_id=task_id)
+        return await self.client.get_run(run_id=task_id)

skyvern_llamaindex/settings.py

@@ -1,13 +1,13 @@
-from typing import Literal
-
 from dotenv import load_dotenv
 from pydantic_settings import BaseSettings
 
+from skyvern.schemas.runs import RunEngine
+
 
 class Settings(BaseSettings):
     api_key: str = ""
     base_url: str = "https://api.skyvern.com"
-    engine: Literal["TaskV1", "TaskV2"] = "TaskV2"
+    engine: RunEngine = RunEngine.skyvern_v2
     run_task_timeout_seconds: int = 60 * 60
 
     class Config:

{skyvern_llamaindex-0.0.4.dist-info → skyvern_llamaindex-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: skyvern-llamaindex
-Version: 0.0.4
+Version: 0.0.5
 Summary: Skyvern integration for LlamaIndex
 Author: lawyzheng
 Author-email: lawy@skyvern.com
@@ -8,12 +8,11 @@ Requires-Python: >=3.11,<3.12
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Requires-Dist: llama-index (>=0.12.19,<0.13.0)
-Requires-Dist: skyvern (>=0.1.56,<0.2.0)
+Requires-Dist: skyvern (>=0.1.84)
 Description-Content-Type: text/markdown
 
 <!-- START doctoc generated TOC please keep comment here to allow auto update -->
 <!-- DON'T EDIT THIS SECTION, INSTEAD RE-RUN doctoc TO UPDATE -->
-**Table of Contents**  *generated with [DocToc](https://github.com/thlorenz/doctoc)*
 
 - [Skyvern LlamaIndex](#skyvern-llamaindex)
   - [Installation](#installation)
@@ -45,7 +44,7 @@ pip install skyvern-llamaindex
 ### Run a task(sync) locally in your local environment
 > sync task won't return 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.
+:warning: :warning: if you want to run this code block, you need to run `skyvern init` command in your terminal to set up skyvern first.
 
 
 ```python
@@ -74,7 +73,7 @@ print(response)
 
 :warning: :warning: if you want to run the task in the background, you need to keep the agent running until the task is finished, otherwise the task will be killed when the agent finished the chat.
 
-: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.
+:warning: :warning: if you want to run this code block, you need to run `skyvern init` command in your terminal to set up skyvern first.
 
 ```python
 import asyncio
@@ -112,7 +111,7 @@ print(response)
 
 ### Get a task locally in your local environment
 
-: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.
+:warning: :warning: if you want to run this code block, you need to run `skyvern init` command in your terminal to set up skyvern first.
 
 ```python
 from dotenv import load_dotenv
@@ -228,7 +227,7 @@ To provide some examples of how to integrate Skyvern with other llama-index tool
 ### Dispatch a task(async) locally in your local environment and wait until the task is finished
 > 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.
+:warning: :warning: if you want to run this code block, you need to run `skyvern init` command in your terminal to set up skyvern first.
 
 ```python
 import asyncio

skyvern_llamaindex-0.0.5.dist-info/RECORD

@@ -0,0 +1,8 @@
+skyvern_llamaindex/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+skyvern_llamaindex/agent.py,sha256=0Qz1HOCcajL-uAzh6TT5pwS2AI8_IldBhzEIz3lm7SE,3213
+skyvern_llamaindex/client.py,sha256=0OCEUEmOP4UD6pR7Vt34lICFjU2NAS5ZI7ce0oIT4cA,3357
+skyvern_llamaindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+skyvern_llamaindex/settings.py,sha256=PG8SyXuwXvRmgWLEkF8i7bgn8pcvznin7opzTMFKEkM,396
+skyvern_llamaindex-0.0.5.dist-info/METADATA,sha256=qhOISJr-RsdWZMQRP2mKDIcJN3EoR9GbrLxhVmHj6Ak,10539
+skyvern_llamaindex-0.0.5.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+skyvern_llamaindex-0.0.5.dist-info/RECORD,,

skyvern_llamaindex-0.0.4.dist-info/RECORD

@@ -1,8 +0,0 @@
-skyvern_llamaindex/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-skyvern_llamaindex/agent.py,sha256=BrA9JxqyETTnJcVv_Hj18oFeDesOINpPpXpszh0F-pc,5622
-skyvern_llamaindex/client.py,sha256=JZZXghqLctZ_yV2XwNNOHhHcPYN7Q38jevOAuzUh4JU,6834
-skyvern_llamaindex/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-skyvern_llamaindex/settings.py,sha256=DsW7xaol9azsJ89b9qvJfoB2BWqSR3C5BkiL1aYx8Ws,386
-skyvern_llamaindex-0.0.4.dist-info/METADATA,sha256=xdaJbvcWjUyqbAslNuFAyDixt1JAKTqem67fWMDjYxI,10787
-skyvern_llamaindex-0.0.4.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-skyvern_llamaindex-0.0.4.dist-info/RECORD,,