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

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)