hyperpocket 0.4.5__py3-none-any.whl → 0.5.1__py3-none-any.whl

hyperpocket/builtin.py

@@ -12,7 +12,7 @@ def get_builtin_tools(pocket_auth: PocketAuth) -> List[Tool]:
     Builtin Tool can access to Pocket Core.
     """
 
-    def __get_current_thread_session_state(thread_id: str = "default") -> str:
+    async def __get_current_thread_session_state(thread_id: str = "default") -> str:
         """
         This tool retrieves the current session state list for the specified thread.
 
@@ -31,7 +31,7 @@ def get_builtin_tools(pocket_auth: PocketAuth) -> List[Tool]:
 
         This tool ensures transparency about the current session but must respect user-driven intent and should never be called automatically or without a specific user request.
         """
-        session_list = pocket_auth.list_session_state(thread_id)
+        session_list = await pocket_auth.list_session_state(thread_id)
         return str(session_list)
 
     def __delete_session(
@@ -58,7 +58,7 @@ def get_builtin_tools(pocket_auth: PocketAuth) -> List[Tool]:
         return str(is_deleted)
 
     builtin_tools = [
-        from_func(func=__get_current_thread_session_state),
+        from_func(func=__get_current_thread_session_state, afunc=__get_current_thread_session_state),
         from_func(func=__delete_session),
     ]
 

hyperpocket/cli/__main__.py

@@ -2,7 +2,6 @@ import click
 
 from hyperpocket.cli.eject import eject
 from hyperpocket.cli.pull import pull
-from hyperpocket.cli.sync import sync
 from hyperpocket.cli.eject import eject
 from hyperpocket.cli.auth_token import create_token_auth_template
 from hyperpocket.cli.auth_oauth2 import create_oauth2_auth_template
@@ -28,7 +27,6 @@ devtool.add_command(build_tool)
 devtool.add_command(export_tool)
 
 cli.add_command(pull)
-cli.add_command(sync)
 cli.add_command(eject)
 
 cli()

hyperpocket/config/logger.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import pathlib
 from logging.handlers import RotatingFileHandler
 from pathlib import Path
 
@@ -27,9 +28,7 @@ class ColorFormatter(logging.Formatter):
 
 
 def get_logger():
-    # init log file
-    package_path = Path(os.path.dirname(hyperpocket.__file__))
-    log_dir = package_path / ".log"
+    log_dir = pathlib.Path(os.getcwd()) / ".log"
     os.makedirs(log_dir, exist_ok=True)
     log_file = log_dir / "pocket.log"
     if not log_file.exists():

hyperpocket/futures/futurestore.py

@@ -25,7 +25,6 @@ class FutureStore(object):
                 f"the future already exists. the existing future is returned. uid: {uid}"
             )
             return future
-
         loop = asyncio.get_running_loop()
         future = loop.create_future()
         future_data = FutureData(future=future, data=data)
@@ -42,7 +41,13 @@ class FutureStore(object):
         if not future_data:
             raise ValueError(f"Future not found for uid={uid}")
         if not future_data.future.done():
-            future_data.future.set_result(value)
+            # if the future loop is running, it should be executed in same event loop
+            loop = future_data.future.get_loop()
+            if loop.is_running():
+                loop.call_soon_threadsafe(future_data.future.set_result, value)
+            # if the future loop is not running, it can be executed from anywhere.
+            else:
+                future_data.future.set_result(value)
 
     def delete_future(self, uid: str):
         self.futures.pop(uid, None)

hyperpocket/pocket_auth.py

@@ -78,7 +78,7 @@ class PocketAuth(object):
         handler = self.find_handler_instance(auth_handler_name, auth_provider)
         return handler.make_request(auth_scopes, **kwargs)
 
-    def check(
+    async def check(
         self,
         auth_req: AuthenticateRequest,
         auth_handler_name: Optional[str] = None,
@@ -112,12 +112,12 @@ class PocketAuth(object):
         """
         handler = self.find_handler_instance(auth_handler_name, auth_provider)
         session = self.session_storage.get(handler.provider(), thread_id, profile)
-        auth_state = self.get_session_state(session=session, auth_req=auth_req)
+        auth_state = await self.get_session_state(session=session, auth_req=auth_req)
 
         return auth_state
 
     @staticmethod
-    def get_session_state(
+    async def get_session_state(
         session: Optional[BaseSessionValue], auth_req: Optional[AuthenticateRequest]
     ) -> AuthState:
         if not session:
@@ -125,6 +125,8 @@ class PocketAuth(object):
 
         if session.auth_resolve_uid:
             future_data = FutureStore.get_future(session.auth_resolve_uid)
+            # it yields before checking future's state, because the future is being resolved on another thread's event loop.
+            await asyncio.sleep(0)
             if future_data is not None and future_data.future.done():
                 return AuthState.RESOLVED
 
@@ -140,7 +142,7 @@ class PocketAuth(object):
 
         return AuthState.SKIP_AUTH
 
-    def prepare(
+    async def prepare(
         self,
         auth_req: AuthenticateRequest,
         auth_handler_name: Optional[str] = None,
@@ -166,7 +168,7 @@ class PocketAuth(object):
         Returns:
             Optional[str]: authentication URL
         """
-        auth_state = self.check(
+        auth_state = await self.check(
             auth_req=auth_req,
             auth_handler_name=auth_handler_name,
             auth_provider=auth_provider,
@@ -263,7 +265,7 @@ class PocketAuth(object):
         Returns:
             AuthContext: authentication context
         """
-        auth_state = self.check(
+        auth_state = await self.check(
             auth_req=auth_req,
             auth_handler_name=auth_handler_name,
             auth_provider=auth_provider,
@@ -353,7 +355,7 @@ class PocketAuth(object):
 
         return session.auth_context
 
-    def list_session_state(
+    async def list_session_state(
         self, thread_id: str, auth_provider: Optional[AuthProvider] = None
     ):
         session_list = self.session_storage.get_by_thread_id(
@@ -361,7 +363,7 @@ class PocketAuth(object):
         )
         session_state_list = []
         for session in session_list:
-            state = self.get_session_state(session=session, auth_req=None)
+            state = await self.get_session_state(session=session, auth_req=None)
 
             session_state_list.append(
                 {

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,29 +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:
             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,
@@ -176,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)
@@ -213,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
 
@@ -245,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,
                     )
                 )
 
@@ -266,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,
@@ -319,42 +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'):
-            self.server.refcnt_down(self._uid)
+            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.teardown()
-
-    def __del__(self):
-        self.teardown()
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        if "server" in state:
-            del state["server"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)