From 342d6c7ab6e35037bb33bf5997ba624e481fa51d Mon Sep 17 00:00:00 2001 From: Karthik Raja A <75213811+kitrak-rev@users.noreply.github.com> Date: Tue, 31 Oct 2023 07:04:56 +0530 Subject: [PATCH] Multi on client toolkit (#12392) Replace this entire comment with: -Add MultiOn close function and update key value and add async functionality - solved the key value TabId not found.. (updated to use latest key value) @hwchase17 --- .../agents/agent_toolkits/multion/toolkit.py | 3 +- .../langchain/tools/multion/__init__.py | 4 +- .../langchain/tools/multion/close_session.py | 67 +++++++++++++++++++ .../langchain/tools/multion/create_session.py | 22 +++++- .../langchain/tools/multion/update_session.py | 41 ++++++++---- 5 files changed, 120 insertions(+), 17 deletions(-) create mode 100644 libs/langchain/langchain/tools/multion/close_session.py diff --git a/libs/langchain/langchain/agents/agent_toolkits/multion/toolkit.py b/libs/langchain/langchain/agents/agent_toolkits/multion/toolkit.py index 5bab2972d5d..224da334efa 100644 --- a/libs/langchain/langchain/agents/agent_toolkits/multion/toolkit.py +++ b/libs/langchain/langchain/agents/agent_toolkits/multion/toolkit.py @@ -5,6 +5,7 @@ from typing import List from langchain.agents.agent_toolkits.base import BaseToolkit from langchain.tools import BaseTool +from langchain.tools.multion.close_session import MultionCloseSession from langchain.tools.multion.create_session import MultionCreateSession from langchain.tools.multion.update_session import MultionUpdateSession @@ -29,4 +30,4 @@ class MultionToolkit(BaseToolkit): def get_tools(self) -> List[BaseTool]: """Get the tools in the toolkit.""" - return [MultionCreateSession(), MultionUpdateSession()] + return [MultionCreateSession(), MultionUpdateSession(), MultionCloseSession()] diff --git a/libs/langchain/langchain/tools/multion/__init__.py b/libs/langchain/langchain/tools/multion/__init__.py index bf8068935ca..d6844572f8f 100644 --- a/libs/langchain/langchain/tools/multion/__init__.py +++ b/libs/langchain/langchain/tools/multion/__init__.py @@ -1,6 +1,6 @@ """MutliOn Client API tools.""" - +from langchain.tools.multion.close_session import MultionCloseSession from langchain.tools.multion.create_session import MultionCreateSession from langchain.tools.multion.update_session import MultionUpdateSession -__all__ = ["MultionCreateSession", "MultionUpdateSession"] +__all__ = ["MultionCreateSession", "MultionUpdateSession", "MultionCloseSession"] diff --git a/libs/langchain/langchain/tools/multion/close_session.py b/libs/langchain/langchain/tools/multion/close_session.py new file mode 100644 index 00000000000..c78557e56cc --- /dev/null +++ b/libs/langchain/langchain/tools/multion/close_session.py @@ -0,0 +1,67 @@ +import asyncio +from typing import TYPE_CHECKING, Optional, Type + +from langchain.callbacks.manager import ( + AsyncCallbackManagerForToolRun, + CallbackManagerForToolRun, +) +from langchain.pydantic_v1 import BaseModel, Field +from langchain.tools.base import BaseTool + +if TYPE_CHECKING: + # This is for linting and IDE typehints + import multion +else: + try: + # We do this so pydantic can resolve the types when instantiating + import multion + except ImportError: + pass + + +class CloseSessionSchema(BaseModel): + """Input for UpdateSessionTool.""" + + sessionId: str = Field( + ..., + description="""The sessionId, received from one of the createSessions + or updateSessions run before""", + ) + + +class MultionCloseSession(BaseTool): + """Tool that closes an existing Multion Browser Window with provided fields. + + Attributes: + name: The name of the tool. Default: "close_multion_session" + description: The description of the tool. + args_schema: The schema for the tool's arguments. Default: UpdateSessionSchema + """ + + name: str = "close_multion_session" + description: str = """Use this tool to close \ +an existing corresponding Multion Browser Window with provided fields. \ +Note: SessionId must be received from previous Browser window creation.""" + args_schema: Type[CloseSessionSchema] = CloseSessionSchema + sessionId: str = "" + + def _run( + self, + sessionId: str, + run_manager: Optional[CallbackManagerForToolRun] = None, + ) -> None: + try: + try: + multion.close_session(sessionId) + except Exception as e: + print(f"{e}, retrying...") + except Exception as e: + raise Exception(f"An error occurred: {e}") + + async def _arun( + self, + sessionId: str, + run_manager: Optional[AsyncCallbackManagerForToolRun] = None, + ) -> None: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._run, sessionId) diff --git a/libs/langchain/langchain/tools/multion/create_session.py b/libs/langchain/langchain/tools/multion/create_session.py index 58a1951b592..008d5be0ae3 100644 --- a/libs/langchain/langchain/tools/multion/create_session.py +++ b/libs/langchain/langchain/tools/multion/create_session.py @@ -1,6 +1,10 @@ +import asyncio from typing import TYPE_CHECKING, Optional, Type -from langchain.callbacks.manager import CallbackManagerForToolRun +from langchain.callbacks.manager import ( + AsyncCallbackManagerForToolRun, + CallbackManagerForToolRun, +) from langchain.pydantic_v1 import BaseModel, Field from langchain.tools.base import BaseTool @@ -57,6 +61,20 @@ class MultionCreateSession(BaseTool): ) -> dict: try: response = multion.new_session({"input": query, "url": url}) - return {"tabId": response["tabId"], "Response": response["message"]} + return { + "sessionId": response["session_id"], + "Response": response["message"], + } except Exception as e: raise Exception(f"An error occurred: {e}") + + async def _arun( + self, + query: str, + url: Optional[str] = "https://www.google.com/", + run_manager: Optional[AsyncCallbackManagerForToolRun] = None, + ) -> dict: + loop = asyncio.get_running_loop() + result = await loop.run_in_executor(None, self._run, query, url) + + return result diff --git a/libs/langchain/langchain/tools/multion/update_session.py b/libs/langchain/langchain/tools/multion/update_session.py index 12e148d19a0..5710de95316 100644 --- a/libs/langchain/langchain/tools/multion/update_session.py +++ b/libs/langchain/langchain/tools/multion/update_session.py @@ -1,6 +1,10 @@ +import asyncio from typing import TYPE_CHECKING, Optional, Type -from langchain.callbacks.manager import CallbackManagerForToolRun +from langchain.callbacks.manager import ( + AsyncCallbackManagerForToolRun, + CallbackManagerForToolRun, +) from langchain.pydantic_v1 import BaseModel, Field from langchain.tools.base import BaseTool @@ -18,8 +22,10 @@ else: class UpdateSessionSchema(BaseModel): """Input for UpdateSessionTool.""" - tabId: str = Field( - ..., description="The tabID, received from one of the createSessions run before" + sessionId: str = Field( + ..., + description="""The sessionID, + received from one of the createSessions run before""", ) query: str = Field( ..., @@ -44,28 +50,39 @@ class MultionUpdateSession(BaseTool): name: str = "update_multion_session" description: str = """Use this tool to update \ an existing corresponding Multion Browser Window with provided fields. \ -Note: TabId must be received from previous Browser window creation.""" +Note: sessionId must be received from previous Browser window creation.""" args_schema: Type[UpdateSessionSchema] = UpdateSessionSchema - tabId: str = "" + sessionId: str = "" def _run( self, - tabId: str, + sessionId: str, query: str, url: Optional[str] = "https://www.google.com/", run_manager: Optional[CallbackManagerForToolRun] = None, ) -> dict: try: try: - response = multion.update_session(tabId, {"input": query, "url": url}) - content = {"tabId": tabId, "Response": response["message"]} - self.tabId = tabId + response = multion.update_session( + sessionId, {"input": query, "url": url} + ) + content = {"sessionId": sessionId, "Response": response["message"]} + self.sessionId = sessionId return content except Exception as e: print(f"{e}, retrying...") return {"error": f"{e}", "Response": "retrying..."} - # response = multion.new_session({"input": query, "url": url}) - # self.tabID = response["tabId"] - # return {"tabId": response["tabId"], "Response": response["message"]} except Exception as e: raise Exception(f"An error occurred: {e}") + + async def _arun( + self, + sessionId: str, + query: str, + url: Optional[str] = "https://www.google.com/", + run_manager: Optional[AsyncCallbackManagerForToolRun] = None, + ) -> dict: + loop = asyncio.get_running_loop() + result = await loop.run_in_executor(None, self._run, sessionId, query, url) + + return result