vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

update llamaindex integration (#2390)

Browse Source
This commit is contained in:
LawyZheng
2025-05-20 00:16:25 +08:00
committed by GitHub
parent 24a73b7af0
commit 2b5bb17506
7 changed files with 937 additions and 444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
integrations/langchain/skyvern_langchain/client.py
View File

@@ -6,6 +6,7 @@ from skyvern_langchain.schema import CreateTaskInput, GetTaskInput
from skyvern_langchain.settings import settings
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
@@ -58,5 +59,5 @@ class GetTask(SkyvernTaskBaseTool):
description: str = """Use Skyvern client to get a task."""
args_schema: Type[BaseModel] = GetTaskInput
async def _arun(self, task_id: str) -> TaskRunResponse:
async def _arun(self, task_id: str) -> AgentGetRunResponse | None:
return await self.get_client().get_run(run_id=task_id)

8
integrations/llama_index/README.md
View File

@@ -31,7 +31,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
@@ -60,7 +60,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
@@ -98,7 +98,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
@@ -214,7 +214,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

1145
integrations/llama_index/poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

4
integrations/llama_index/pyproject.toml
View File

@@ -1,6 +1,6 @@
[tool.poetry]
name = "skyvern-llamaindex"
version = "0.0.4"
version = "0.0.5"
description = "Skyvern integration for LlamaIndex"
authors = ["lawyzheng <lawy@skyvern.com>"]
packages = [{ include = "skyvern_llamaindex" }]
@@ -8,7 +8,7 @@ readme = "README.md"
[tool.poetry.dependencies]
python = "^3.11,<3.12"
skyvern = "^0.1.56"
skyvern = ">=0.1.84"
llama-index = "^0.12.19"

96
integrations/llama_index/skyvern_llamaindex/agent.py
View File

@@ -1,21 +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 SkyvernAgent
from skyvern.forge import app
from skyvern.forge.prompts import prompt_engine
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
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[SkyvernAgent] = None):
def __init__(self, agent: Optional[Skyvern] = None):
if agent is None:
agent = SkyvernAgent()
agent = Skyvern(base_url=None, api_key=None)
self.agent = agent
def run_task(self) -> FunctionTool:
@@ -41,23 +39,17 @@ class SkyvernTaskToolSpec(BaseToolSpec):
def __init__(
self,
*,
agent: SkyvernAgent | None = 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 = SkyvernAgent()
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 | TaskV2:
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.
@@ -65,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 | TaskV2:
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.
@@ -79,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 | TaskV2 | 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_v1(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_v1(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) -> 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) -> 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) -> TaskV2 | None:
return await self.agent.get_observer_task_v_2(task_id=task_id)
return await self.agent.get_run(run_id=task_id)

119
integrations/llama_index/skyvern_llamaindex/client.py
View File

@@ -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.task_v2 import TaskV2Request
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_v1(
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 = 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,
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 = 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,
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)

6
integrations/llama_index/skyvern_llamaindex/settings.py
View File

@@ -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: