update llamaindex integration (#2390)
This commit is contained in:
LawyZheng
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user