update langchain integration (#1815)

2025-03-03 10:38:00 +08:00
parent 889c8f1963
commit b3b548ed38
6 changed files with 369 additions and 211 deletions
--- a/integrations/langchain/skyvern_langchain/agent.py
+++ b/integrations/langchain/skyvern_langchain/agent.py
@@ -1,48 +1,97 @@
-from typing import Any, Dict
+from typing import Any, Literal, Type

-from langchain.tools import tool
-from skyvern_langchain.schema import GetTaskInput, TaskV1Request, TaskV2Request
+from langchain.tools import BaseTool
+from litellm import BaseModel
+from pydantic import Field
+from skyvern_langchain.schema import CreateTaskInput, GetTaskInput
+from skyvern_langchain.settings import settings

 from skyvern.agent import Agent
-from skyvern.forge.sdk.schemas.observers import ObserverTask
-from skyvern.forge.sdk.schemas.tasks import CreateTaskResponse, TaskResponse
+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
+
+agent = Agent()


-@tool("run-skyvern-agent-task-v1", args_schema=TaskV1Request)
-async def run_task_v1(**kwargs: Dict[str, Any]) -> TaskResponse:
-    """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."""
-    task_request = TaskV1Request(**kwargs)
-    return await Agent().run_task(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
+class SkyvernTaskBaseTool(BaseTool):
+    engine: Literal["TaskV1", "TaskV2"] = Field(default=settings.engine)
+    timeout_seconds: int = Field(default=settings.run_task_timeout)
+    agent: Agent = agent
+
+    def _run(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("skyvern task tool does not support sync")
+
+    # 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)


-@tool("queue-skyvern-agent-task-v1", args_schema=TaskV1Request)
-async def queue_task_v1(**kwargs: Dict[str, Any]) -> CreateTaskResponse:
-    """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."""
-    task_request = TaskV1Request(**kwargs)
-    return await Agent().create_task(task_request=task_request)
+class RunTask(SkyvernTaskBaseTool):
+    name: str = "run-skyvern-agent-task"
+    description: str = """Use Skyvern agent to run a task. This function won't return until the task is finished."""
+    args_schema: Type[BaseModel] = CreateTaskInput
+
+    async def _arun(self, user_prompt: str, url: str | None = None) -> TaskResponse | ObserverTask:
+        if self.engine == "TaskV1":
+            return await self._arun_task_v1(user_prompt=user_prompt, url=url)
+        else:
+            return await self._arun_task_v2(user_prompt=user_prompt, url=url)
+
+    async def _arun_task_v1(self, user_prompt: str, url: str | None = 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.timeout_seconds)
+
+    async def _arun_task_v2(self, user_prompt: str, url: str | None = 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.timeout_seconds)


-@tool("get-skyvern-agent-task-v1", args_schema=GetTaskInput)
-async def get_task_v1(task_id: str) -> TaskResponse | None:
-    """Use Skyvern agent to get a v1 task. v1 tasks are usually simple tasks."""
-    return await Agent().get_task(task_id=task_id)
+class DispatchTask(SkyvernTaskBaseTool):
+    name: str = "dispatch-skyvern-agent-task"
+    description: str = """Use Skyvern agent to dispatch a task. This function will return immediately and the task will be running in the background."""
+    args_schema: Type[BaseModel] = CreateTaskInput
+
+    async def _arun(self, user_prompt: str, url: str | None = None) -> CreateTaskResponse | ObserverTask:
+        if self.engine == "TaskV1":
+            return await self._arun_task_v1(user_prompt=user_prompt, url=url)
+        else:
+            return await self._arun_task_v2(user_prompt=user_prompt, url=url)
+
+    async def _arun_task_v1(self, user_prompt: str, url: str | None = 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 _arun_task_v2(self, user_prompt: str, url: str | None = None) -> ObserverTask:
+        task_request = ObserverTaskRequest(user_prompt=user_prompt, url=url)
+        return await self.agent.observer_task_v_2(task_request=task_request)


-@tool("run-skyvern-agent-task-v2", args_schema=TaskV2Request)
-async def run_task_v2(**kwargs: Dict[str, Any]) -> ObserverTask:
-    """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."""
-    task_request = TaskV2Request(**kwargs)
-    return await Agent().run_observer_task_v_2(task_request=task_request, timeout_seconds=task_request.timeout_seconds)
+class GetTask(SkyvernTaskBaseTool):
+    name: str = "get-skyvern-agent-task"
+    description: str = """Use Skyvern agent to get a task."""
+    args_schema: Type[BaseModel] = GetTaskInput

+    async def _arun(self, task_id: str) -> TaskResponse | ObserverTask | None:
+        if self.engine == "TaskV1":
+            return await self._arun_task_v1(task_id=task_id)
+        else:
+            return await self._arun_task_v2(task_id=task_id)

-@tool("queue-skyvern-agent-task-v2", args_schema=TaskV2Request)
-async def queue_task_v2(**kwargs: Dict[str, Any]) -> ObserverTask:
-    """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."""
-    task_request = TaskV2Request(**kwargs)
-    return await Agent().observer_task_v_2(task_request=task_request)
+    async def _arun_task_v1(self, task_id: str) -> TaskResponse | None:
+        return await self.agent.get_task(task_id=task_id)

-
-@tool("get-skyvern-agent-task-v2", args_schema=GetTaskInput)
-async def get_task_v2(task_id: str) -> ObserverTask | None:
-    """Use Skyvern agent to get a v2 task. v2 tasks are usually complicated tasks."""
-    return await Agent().get_observer_task_v_2(task_id=task_id)
+    async def _arun_task_v2(self, task_id: str) -> ObserverTask | None:
+        return await self.agent.get_observer_task_v_2(task_id=task_id)