hyperpocket 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

hyperpocket/pocket_main.py

@@ -1,18 +1,34 @@
 import asyncio
-import uuid
-from typing import Any, List, Union
+import concurrent.futures
+from threading import Lock
+from typing import Any, List, Union, Callable, Optional
 
+from hyperpocket.builtin import get_builtin_tools
 from hyperpocket.config import pocket_logger
 from hyperpocket.pocket_auth import PocketAuth
-from hyperpocket.pocket_core import PocketCore
-from hyperpocket.server.server import PocketServer, PocketServerOperations
+from hyperpocket.server.server import PocketServer
+from hyperpocket.tool import Tool, from_func
+from hyperpocket.tool.dock import Dock
 from hyperpocket.tool_like import ToolLike
 
 
 class Pocket(object):
     server: PocketServer
-    core: PocketCore
-    _uid: str
+    auth: PocketAuth
+    tools: dict[str, Tool]
+
+    _cnt_pocket_count: int = 0
+    _pocket_count_lock = Lock()
+
+    @staticmethod
+    def _default_dock() -> Dock:
+        try:
+            from hyperdock_container.dock import ContainerDock
+            pocket_logger.info("hyperdock-container is loaded.")
+            return ContainerDock()
+        except ImportError as e:
+            pocket_logger.warning("Failed to import hyperdock_container.")
+            raise e
 
     def __init__(
         self,
@@ -21,30 +37,31 @@ class Pocket(object):
         use_profile: bool = False,
     ):
         try:
-            self._uid = str(uuid.uuid4())
+            if auth is None:
+                auth = PocketAuth()
+            self.auth = auth
             self.use_profile = use_profile
-            self.server = PocketServer.get_instance_and_refcnt_up(self._uid)
-            self.server.wait_initialized()
-            self.core = PocketCore(
-                tools=tools,
-                auth=auth,
+            self.server = PocketServer.get_instance()
+
+            self._load_tools(tools)
+            pocket_logger.info(
+                f"All Registered Tools Loaded successfully. total registered tools : {len(self.tools)}"
+            )
+
+            # load builtin tool
+            builtin_tools = get_builtin_tools(self.auth)
+            for tool in builtin_tools:
+                self.tools[tool.name] = tool
+            pocket_logger.info(
+                f"All BuiltIn Tools Loaded successfully. total tools : {len(self.tools)}"
             )
+
+            with Pocket._pocket_count_lock:
+                Pocket._cnt_pocket_count += 1
         except Exception as e:
-            if hasattr(self, "server"):
-                self.server.refcnt_down(self._uid)
+            self.teardown()
             pocket_logger.error(f"Failed to initialize pocket server. error : {e}")
-            self._teardown_server()
             raise e
-        
-        try:
-            asyncio.get_running_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-        else:
-            import nest_asyncio
-            loop = asyncio.new_event_loop()
-            nest_asyncio.apply(loop=loop)
-        loop.run_until_complete(self.server.plug_core(self._uid, self.core))
 
     def invoke(
         self,
@@ -177,17 +194,12 @@ class Pocket(object):
             **kwargs,
         }
 
-        result, paused = await self.server.call_in_subprocess(
-            PocketServerOperations.CALL,
-            self._uid,
-            args,
-            {
-                "tool_name": tool_name,
-                "body": body,
-                "thread_id": thread_id,
-                "profile": profile,
-                **kwargs,
-            },
+        result, paused = await self.acall(
+            tool_name=tool_name,
+            body=body,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs
         )
         if not isinstance(result, str):
             result = str(result)
@@ -214,15 +226,16 @@ class Pocket(object):
         Returns:
             List[str]: A list of authentication URIs for the tools that require authentication.
         """
-        tool_by_provider = self.core.grouping_tool_by_auth_provider()
+        tool_by_provider = self.grouping_tool_by_auth_provider()
 
         prepare_list = {}
         for provider, tools in tool_by_provider.items():
             tool_name_list = [tool.name for tool in tools]
-            prepare = await self.prepare_in_subprocess(
-                tool_name=tool_name_list, thread_id=thread_id, profile=profile
+            prepare = await self.prepare_auth(
+                tool_name=tool_name_list,
+                thread_id=thread_id,
+                profile=profile,
             )
-
             if prepare is not None:
                 prepare_list[provider] = prepare
 
@@ -246,16 +259,17 @@ class Pocket(object):
             or `False` if the process was interrupted or failed.
         """
         try:
-            tool_by_provider = self.core.grouping_tool_by_auth_provider()
+            tool_by_provider = self.grouping_tool_by_auth_provider()
 
             waiting_futures = []
             for provider, tools in tool_by_provider.items():
                 if len(tools) == 0:
                     continue
-
                 waiting_futures.append(
-                    self.authenticate_in_subprocess(
-                        tool_name=tools[0].name, thread_id=thread_id, profile=profile
+                    self.authenticate(
+                        tool_name=tools[0].name,
+                        thread_id=thread_id,
+                        profile=profile,
                     )
                 )
 
@@ -267,51 +281,150 @@ class Pocket(object):
             pocket_logger.error("authentication time out.")
             raise e
 
-    async def prepare_in_subprocess(
+    async def acall(
         self,
-        tool_name: Union[str, List[str]],
+        tool_name: str,
+        body: Any,
         thread_id: str = "default",
         profile: str = "default",
         *args,
         **kwargs,
-    ):
-        prepare = await self.server.call_in_subprocess(
-            PocketServerOperations.PREPARE_AUTH,
-            self._uid,
-            args,
-            {
-                "tool_name": tool_name,
-                "thread_id": thread_id,
-                "profile": profile,
-                **kwargs,
-            },
+    ) -> tuple[str, bool]:
+        """
+        Invoke tool asynchronously, not that different from `Pocket.invoke`
+        But this method is called only in subprocess.
+
+        This function performs the following steps:
+            1. `prepare_auth` : preparing the authentication process for the tool if necessary.
+            2. `authenticate` : performing authentication that needs to invoke tool.
+            3. `tool_call` : Executing tool actually with authentication information.
+
+        Args:
+            tool_name(str): tool name to invoke
+            body(Any): tool arguments. should be json format
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            tuple[str, bool]: tool result and state.
+        """
+        pocket_logger.debug(f"{tool_name} tool call. body: {body}")
+        tool = self._tool_instance(tool_name)
+        if tool.auth is not None:
+            callback_info = await self.prepare_auth(tool_name, thread_id, profile, **kwargs)
+            if callback_info:
+                return callback_info, True
+        # 02. authenticate
+        credentials = await self.authenticate(tool_name, thread_id, profile, **kwargs)
+        # 03. call tool
+        result = await self.tool_call(tool_name, body=body, envs=credentials, **kwargs)
+        pocket_logger.debug(f"{tool_name} tool call result: {result}")
+        return result, False
+
+    async def prepare_auth(
+        self,
+        tool_name: Union[str, List[str]],
+        thread_id: str = "default",
+        profile: str = "default",
+        **kwargs,
+    ) -> Optional[str]:
+        """
+        Prepares the authentication process for the tool if necessary.
+        Returns callback URL and whether the tool requires authentication.
+
+        Args:
+            tool_name(Union[str,List[str]]): tool name to invoke
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            Optional[str]: callback URI if necessary
+        """
+
+        if isinstance(tool_name, str):
+            tool_name = [tool_name]
+
+        tools: List[Tool] = []
+        for name in tool_name:
+            tool = self._tool_instance(name)
+            if tool.auth is not None:
+                tools.append(tool)
+
+        if len(tools) == 0:
+            return None
+
+        auth_handler_name = tools[0].auth.auth_handler
+        auth_provider = tools[0].auth.auth_provider
+        auth_scopes = set()
+
+        for tool in tools:
+            if tool.auth.auth_handler != auth_handler_name:
+                pocket_logger.error(
+                    f"All Tools should have same auth handler. but it's different {tool.auth.auth_handler}, {auth_handler_name}"
+                )
+
+                return f"All Tools should have same auth handler. but it's different {tool.auth.auth_handler}, {auth_handler_name}"
+            if tool.auth.auth_provider != auth_provider:
+                pocket_logger.error(
+                    f"All Tools should have same auth provider. but it's different {tool.auth.auth_provider}, {auth_provider}"
+                )
+                return f"All Tools should have same auth provider. but it's different {tool.auth.auth_provider}, {auth_provider}"
+
+            if tool.auth.scopes is not None:
+                auth_scopes |= set(tool.auth.scopes)
+
+        auth_req = self.auth.make_request(
+            auth_handler_name=auth_handler_name,
+            auth_provider=auth_provider,
+            auth_scopes=list(auth_scopes),
         )
 
-        return prepare
+        return await self.auth.prepare(
+            auth_req=auth_req,
+            auth_handler_name=auth_handler_name,
+            auth_provider=auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs,
+        )
 
-    async def authenticate_in_subprocess(
+    async def authenticate(
         self,
         tool_name: str,
         thread_id: str = "default",
         profile: str = "default",
-        *args,
         **kwargs,
-    ):
-        credentials = await self.server.call_in_subprocess(
-            PocketServerOperations.AUTHENTICATE,
-            self._uid,
-            args,
-            {
-                "tool_name": tool_name,
-                "thread_id": thread_id,
-                "profile": profile,
-                **kwargs,
-            },
-        )
+    ) -> dict[str, str]:
+        """
+        Authenticates the handler included in the tool and returns credentials.
 
-        return credentials
+        Args:
+            tool_name(str): tool name to invoke
+            thread_id(str): thread id
+            profile(str): profile name
+
+        Returns:
+            dict[str, str]: credentials
+        """
+        tool = self._tool_instance(tool_name)
+        if tool.auth is None:
+            return {}
+        auth_req = self.auth.make_request(
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            auth_scopes=tool.auth.scopes,
+        )
+        auth_ctx = await self.auth.authenticate_async(
+            auth_req=auth_req,
+            auth_handler_name=tool.auth.auth_handler,
+            auth_provider=tool.auth.auth_provider,
+            thread_id=thread_id,
+            profile=profile,
+            **kwargs,
+        )
+        return auth_ctx.to_dict()
 
-    async def tool_call_in_subprocess(
+    async def tool_call(
         self,
         tool_name: str,
         body: Any,
@@ -320,38 +433,76 @@ class Pocket(object):
         *args,
         **kwargs,
     ):
-        result = await self.server.call_in_subprocess(
-            PocketServerOperations.TOOL_CALL,
-            self._uid,
-            args,
-            {
-                "tool_name": tool_name,
-                "body": body,
-                "thread_id": thread_id,
-                "profile": profile,
-                **kwargs,
-            },
-        )
+        try:
+            tool = self._tool_instance(tool_name)
+            result = await asyncio.wait_for(
+                tool.ainvoke(body=body, thread_id=thread_id, profile=profile, **kwargs), timeout=180)
+        except asyncio.TimeoutError:
+            pocket_logger.warning("Timeout tool call.")
+            return "timeout tool call"
+
+        # TODO(moon): extract
+        if tool.postprocessings is not None:
+            for postprocessing in tool.postprocessings:
+                try:
+                    result = postprocessing(result)
+                except Exception as e:
+                    exception_str = (
+                        f"Error in postprocessing `{postprocessing.__name__}`: {e}"
+                    )
+                    pocket_logger.error(exception_str)
+                    return exception_str
 
         return result
 
+    def grouping_tool_by_auth_provider(self) -> dict[str, List[Tool]]:
+        tool_by_provider = {}
+        for tool_name, tool in self.tools.items():
+            if tool.auth is None:
+                continue
+
+            auth_provider_name = tool.auth.auth_provider.name
+            if tool_by_provider.get(auth_provider_name):
+                tool_by_provider[auth_provider_name].append(tool)
+            else:
+                tool_by_provider[auth_provider_name] = [tool]
+        return tool_by_provider
+
+    def _load_tools(self, tools):
+        self.tools = dict()
+        dock = self._default_dock()
+
+        def _load(tool_like):
+            if isinstance(tool_like, str) or isinstance(tool_like, tuple):
+                return dock(tool_like)
+            elif isinstance(tool_like, Tool):
+                return tool_like
+            elif isinstance(tool_like, Callable):
+                return from_func(tool_like)
+            else:
+                raise ValueError(f"Invalid tool type: {type(tool_like)}")
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10, thread_name_prefix="tool-loader") as executor:
+            futures = [executor.submit(_load, tool_like) for tool_like in tools]
+            for future in concurrent.futures.as_completed(futures):
+                tool = future.result()
+                self.tools[tool.name] = tool
+
+    def _tool_instance(self, tool_name: str) -> Tool:
+        return self.tools[tool_name]
+
     def _teardown_server(self):
-        self.server.teardown()
+        self.teardown()
+
+    def teardown(self):
+        if hasattr(self, 'server'):
+            with Pocket._pocket_count_lock:
+                Pocket._cnt_pocket_count -= 1
+            if Pocket._cnt_pocket_count <= 0:
+                self.server.teardown()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.server.refcnt_down(self._uid)
-
-    def __del__(self):
-        self.server.refcnt_down(self._uid)
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        if "server" in state:
-            del state["server"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
+        self.teardown()

hyperpocket/server/auth/kraken.py

@@ -0,0 +1,58 @@
+from http import HTTPStatus
+
+from fastapi import APIRouter, Form
+from starlette.responses import HTMLResponse, RedirectResponse
+from hyperpocket.futures import FutureStore
+from hyperpocket.server.auth.token import add_query_params
+
+kraken_auth_router = APIRouter(prefix="/kraken")
+
+@kraken_auth_router.get("/keypair", response_class=HTMLResponse)
+async def keypair_form(redirect_uri: str, state: str = ""):
+    html = f"""    
+    <html>
+        <body>
+            <h2>Enter Token</h2>
+            <form action="submit" method="post">
+                <input type="hidden" name="redirect_uri" value="{redirect_uri}">
+                <input type="hidden" name="state" value="{state}">
+                
+                <label for="kraken_api_key">Kraken API Key:</label>
+                <input type="text" id="kraken_api_key" name="kraken_api_key" required>
+                
+                <label for="kraken_api_secret">Kraken API Secret:</label>
+                <input type="text" id="kraken_api_secret" name="kraken_api_secret" required>
+                
+                <button type="submit">submit</button>
+            </form>
+        </body>
+    </html>
+    """
+    return HTMLResponse(content=html)
+
+@kraken_auth_router.post("/submit", response_class=RedirectResponse)
+async def submit_keypair(
+    kraken_api_key: str = Form(...),
+    kraken_api_secret: str = Form(...),
+    redirect_uri: str = Form(...),
+    state: str = Form(...),
+):
+    new_callback_url = add_query_params(
+        redirect_uri, {
+            "kraken_api_key": kraken_api_key,
+            "kraken_api_secret": kraken_api_secret,
+            "state": state,
+        }
+    )
+    return RedirectResponse(url=new_callback_url, status_code=HTTPStatus.SEE_OTHER)
+
+@kraken_auth_router.get("/keypair/callback")
+async def kraken_keypair_callback(state: str, kraken_api_key: str, kraken_api_secret: str):
+    try:
+        FutureStore.resolve_future(state, {
+            "kraken_api_key": kraken_api_key,
+            "kraken_api_secret": kraken_api_secret,
+        })
+    except ValueError:
+        return HTMLResponse(content="failed")
+    return HTMLResponse(content="success")