hyperpocket 0.0.2__py3-none-any.whl → 0.1.8__py3-none-any.whl

hyperpocket/pocket_main.py

@@ -1,86 +1,32 @@
 import asyncio
-import pathlib
-from typing import Union, Any, Optional, Callable
+from typing import Any, Optional, List, Union
 
 from hyperpocket.config import pocket_logger
 from hyperpocket.pocket_auth import PocketAuth
-from hyperpocket.repository import Lockfile
-from hyperpocket.repository.lock import LocalLock, GitLock
+from hyperpocket.pocket_core import PocketCore
 from hyperpocket.server.server import PocketServer, PocketServerOperations
-from hyperpocket.tool import Tool, ToolRequest
-from hyperpocket.tool.function.tool import FunctionTool
-from hyperpocket.tool.wasm import WasmTool
-from hyperpocket.tool.wasm.tool import WasmToolRequest
-
-ToolLike = Union[Tool, str, Callable, ToolRequest]
+from hyperpocket.tool_like import ToolLike
 
 
 class Pocket(object):
-    auth: PocketAuth
     server: PocketServer
-    tools: dict[str, Tool]
+    core: PocketCore
 
     def __init__(self,
                  tools: list[ToolLike],
                  auth: PocketAuth = None,
                  lockfile_path: Optional[str] = None,
                  force_update: bool = False):
-        if auth is None:
-            auth = PocketAuth()
-
-        if lockfile_path is None:
-            lockfile_path = "./pocket.lock"
-        lockfile_pathlib_path = pathlib.Path(lockfile_path)
-        lockfile = Lockfile(lockfile_pathlib_path)
-        tool_likes = []
-        for tool_like in tools:
-            if isinstance(tool_like, str):
-                if pathlib.Path(tool_like).exists():
-                    lock = LocalLock(tool_like)
-                    req = WasmToolRequest(lock, "")
-                else:
-                    lock = GitLock(repository_url=tool_like, git_ref='HEAD')
-                    req = WasmToolRequest(lock, "")
-                lockfile.add_lock(lock)
-                tool_likes.append(req)
-            elif isinstance(tool_like, WasmToolRequest):
-                lockfile.add_lock(tool_like.lock)
-                tool_likes.append(tool_like)
-            elif isinstance(tool_like, ToolRequest):
-                raise ValueError(f"unreachable. tool_like:{tool_like}")
-            elif isinstance(tool_like, WasmTool):
-                raise ValueError("WasmTool should pass ToolRequest instance instead.")
-            else:
-                tool_likes.append(tool_like)
-        lockfile.sync(force_update=force_update, referenced_only=True)
-
-        self.tools = dict()
-        for tool_like in tool_likes:
-            tool = self._load_tool(tool_like, lockfile)
-            if tool.name in self.tools:
-                pocket_logger.error(f"Duplicate tool name: {tool.name}.")
-                raise ValueError(f"Duplicate tool name: {tool.name}")
-            self.tools[tool.name] = tool
-
-        pocket_logger.info(f"All Tools Loaded successfully. total registered tools : {len(self.tools)}")
-
-        self.auth = auth
+
+        self.core = PocketCore(
+            tools=tools,
+            auth=auth,
+            lockfile_path=lockfile_path,
+            force_update=force_update
+        )
+
         self.server = PocketServer()
-        self.server.run(self)
-
-    def _load_tool(self, tool_like: ToolLike, lockfile: Lockfile) -> Tool:
-        pocket_logger.info(f"Loading Tool {tool_like}")
-        if isinstance(tool_like, Tool):
-            tool = tool_like
-        elif isinstance(tool_like, ToolRequest):
-            tool = Tool.from_tool_request(tool_like, lockfile=lockfile)
-        elif isinstance(tool_like, Callable):
-            tool = FunctionTool.from_func(tool_like)
-        else:
-            raise ValueError(f"Invalid tool type: {type(tool_like)}")
-
-        pocket_logger.info(f"Complete Loading Tool {tool.name}")
-        return tool
+        self.server.run(self.core)
 
     def invoke(self,
                tool_name: str,
@@ -204,131 +150,139 @@ class Pocket(object):
 
         return result, paused
 
-    # DO NOT EVER THINK ABOUT USING SERVER PROPERTY BELOW HERE
-    async def acall(self,
-                    tool_name: str,
-                    body: Any,
-                    thread_id: str = 'default',
-                    profile: str = 'default',
-                    *args, **kwargs) -> tuple[str, bool]:
+    async def initialize_tool_auth(self,
+                                   thread_id: str = 'default',
+                                   profile: str = 'default',
+                                   ) -> dict[str, str]:
         """
-        Invoke tool asynchronously, not that different from `Pocket.invoke`
-        But this method is called only in subprocess.
+        Initialize authentication for all tools.
 
-        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.
+        This method prepares all tools that require authentication by retrieving
+        their respective authentication URIs.
+
+        If no tool requires authentication, an empty list is returned.
 
         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
+            thread_id (str): The thread id. Defaults to 'default'.
+            profile (str): The profile to be used for authentication. Defaults to 'default'.
 
         Returns:
-            tuple[str, bool]: tool result and state.
+            List[str]: A list of authentication URIs for the tools that require authentication.
         """
-        tool = self._tool_instance(tool_name)
-        if tool.auth is not None:
-            callback_info = 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)
-        return result, False
-
-    def prepare_auth(self,
-                     tool_name: 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.
+        tool_by_provider = self.core.grouping_tool_by_auth_provider()
 
-        Args:
-            tool_name(str): tool name to invoke
-            thread_id(str): thread id
-            profile(str): profile name
+        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)
 
-        Returns:
-            Optional[str]: callback URI if necessary
-        """
-        tool = self._tool_instance(tool_name)
-        if tool.auth is None:
-            return None
-
-        auth_req = self.auth.make_request(
-            auth_handler_name=tool.auth.auth_handler,
-            auth_provider=tool.auth.auth_provider,
-            auth_scopes=tool.auth.scopes)
-
-        return self.auth.prepare(
-            auth_req=auth_req,
-            auth_handler_name=tool.auth.auth_handler,
-            auth_provider=tool.auth.auth_provider,
-            thread_id=thread_id,
-            profile=profile,
-            **kwargs
-        )
+            if prepare is not None:
+                prepare_list[provider] = prepare
 
-    async def authenticate(
-            self,
-            tool_name: str,
-            thread_id: str = 'default',
-            profile: str = 'default',
-            **kwargs) -> dict[str, str]:
+        return prepare_list
+
+    async def wait_tool_auth(self,
+                             thread_id: str = 'default',
+                             profile: str = 'default'
+                             ) -> bool:
         """
-        Authenticates the handler included in the tool and returns credentials.
+        Wait until all tool authentications are completed.
+
+        This method waits until all tools associated with the given
+        `thread_id` and `profile` have completed their authentication process.
 
         Args:
-            tool_name(str): tool name to invoke
-            thread_id(str): thread id
-            profile(str): profile name
+            thread_id (str): The thread id. Defaults to 'default'.
+            profile (str): The profile to be used for authentication. Defaults to 'default'.
 
         Returns:
-            dict[str, str]: credentials
+            bool: Returns `True` if all tool authentications are successfully completed,
+            or `False` if the process was interrupted or failed.
         """
-        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,
+        try:
+            tool_by_provider = self.core.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
+                    ))
+
+            await asyncio.gather(*waiting_futures)
+
+            return True
+
+        except asyncio.TimeoutError as e:
+            pocket_logger.error("authentication time out.")
+            raise e
+
+    async def prepare_in_subprocess(self,
+                                    tool_name: Union[str, List[str]],
+                                    thread_id: str = 'default',
+                                    profile: str = 'default',
+                                    *args, **kwargs):
+        prepare = await self.server.call_in_subprocess(
+            PocketServerOperations.PREPARE_AUTH,
+            args,
+            {
+                'tool_name': tool_name,
+                'thread_id': thread_id,
+                'profile': profile,
+                **kwargs,
+            },
         )
-        return auth_ctx.to_dict()
 
-    async def tool_call(self, tool_name: str, **kwargs) -> str:
-        """
-        Executing tool actually
+        return prepare
 
-        Args:
-            tool_name(str): tool name to invoke
-            kwargs(dict): keyword arguments. authentication information is passed through this.
+    async def authenticate_in_subprocess(self,
+                                         tool_name: str,
+                                         thread_id: str = 'default',
+                                         profile: str = 'default',
+                                         *args, **kwargs):
+        credentials = await self.server.call_in_subprocess(
+            PocketServerOperations.AUTHENTICATE,
+            args,
+            {
+                'tool_name': tool_name,
+                'thread_id': thread_id,
+                'profile': profile,
+                **kwargs,
+            },
+        )
 
-        Returns:
-            str: tool result
-        """
-        tool = self._tool_instance(tool_name)
-        try:
-            return await asyncio.wait_for(tool.ainvoke(**kwargs), timeout=180)
-        except asyncio.TimeoutError:
-            pocket_logger.warning("Timeout tool call.")
-            return "timeout tool call"
+        return credentials
+
+    async def tool_call_in_subprocess(self,
+                                      tool_name: str,
+                                      body: Any,
+                                      thread_id: str = 'default',
+                                      profile: str = 'default',
+                                      *args, **kwargs):
+        result = await self.server.call_in_subprocess(
+            PocketServerOperations.TOOL_CALL,
+            args,
+            {
+                'tool_name': tool_name,
+                'body': body,
+                'thread_id': thread_id,
+                'profile': profile,
+                **kwargs,
+            },
+        )
+
+        return result
 
-    def _tool_instance(self, tool_name: str) -> Tool:
-        return self.tools[tool_name]
+    def _teardown_server(self):
+        self.server.teardown()
 
     def __enter__(self):
         return self

hyperpocket/prompts.py

@@ -1,15 +1,13 @@
 def pocket_extended_tool_description(description: str):
     return f'''
-This tool performs as stated in the <tool-description></tool-description> XML tag.
+This tool functions as described in the <tool-description> XML tag:
 <tool-description>
 {description}
 </tool-description>
 
-This tool requires arguments as follows:
-- 'thread_id': The ID of the chat thread where the tool is invoked. Omitted when unknown.
-- 'profile': The profile of the user invoking the tool. Inferred from user's messages.
-Users can request tools to be invoked in specific personas, which is called a profile.
-If the user's profile name can be inferred from the query, pass it as a string in the 'profile' JSON property.
-Omitted when unknown.
-- 'body': The argument of the tool. The argument is passed as JSON property 'body' in the JSON schema.
+Arguments required:
+- 'thread_id': Chat thread ID (omit if unknown).
+- 'profile': The profile of the user invoking the tool. (infer from messages; omit if unknown).
+Users can request tools to be invoked in specific personas.
+- 'body': Tool argument passed as a JSON 'body'.
 '''

hyperpocket/repository/__init__.py

@@ -1,5 +1,5 @@
 from hyperpocket.repository.lock import Lock
 from hyperpocket.repository.lockfile import Lockfile
-from hyperpocket.repository.repository import pull, sync
+from hyperpocket.repository.repository import pull, sync, eject
 
-__all__ = ["Lock", "Lockfile", "pull", "sync"]
+__all__ = ["Lock", "Lockfile", "pull", "sync", "eject"]

hyperpocket/repository/lock.py

@@ -28,6 +28,10 @@ class Lock(BaseModel, abc.ABC):
     def toolpkg_path(self) -> pathlib.Path:
         raise NotImplementedError
 
+    def eject_to_path(self, dest_path: pathlib.Path, src_sub_path: str = None):
+        ## local locks are already tracked by git
+        raise NotImplementedError
+
 
 class LocalLock(Lock):
     tool_source: str = Field(default='local')
@@ -154,3 +158,18 @@ class GitLock(Lock):
                 new_sha = sha
                 break
         return new_sha
+
+    def eject_to_path(self, dest_path: pathlib.Path, src_sub_path: str = None):
+
+        # clone the git repository to the target path
+        pocket_logger.info(
+            f"Ejecting git: {self.repository_url} @ ref: {self.git_ref} source in path: {src_sub_path} to {dest_path} ..."
+        )
+        if dest_path.exists():
+            shutil.rmtree(dest_path)
+
+        if src_sub_path:
+            src_path = self.toolpkg_path() / src_sub_path
+        else:
+            src_path = self.toolpkg_path()
+        shutil.copytree(src_path, dest_path)

hyperpocket/repository/lockfile.py

@@ -8,19 +8,19 @@ class Lockfile:
     path: pathlib.Path = None
     locks: dict[tuple, Lock] = None
     referenced_locks: set[tuple] = None
-    
+
     def __init__(self, path: pathlib.Path):
         self.path = path
         self.locks = {}
         self.referenced_locks = set()
         if self.path.exists():
-            with open(self.path, 'r') as f:
+            with open(self.path, "r") as f:
                 for line in f:
-                    split = line.strip().split('\t')
+                    split = line.strip().split("\t")
                     source = split[0]
-                    if source == 'local':
+                    if source == "local":
                         lock = LocalLock(tool_path=split[1])
-                    elif source == 'git':
+                    elif source == "git":
                         lock = GitLock(
                             repository_url=split[1],
                             git_ref=split[2],
@@ -31,12 +31,17 @@ class Lockfile:
                     self.locks[lock.key()] = lock
         else:
             self.path.touch()
-    
+
     def add_lock(self, lock: Lock):
         if lock.key() not in self.locks:
             self.locks[lock.key()] = lock
         self.referenced_locks.add(lock.key())
-    
+
+    def remove_lock(self, key: tuple[str, ...]):
+        self.locks.pop(key)
+        if key in self.referenced_locks:
+            self.referenced_locks.remove(key)
+
     def get_lock(self, key: tuple[str, ...]):
         return self.locks[key]
 
@@ -45,12 +50,13 @@ class Lockfile:
             locks = [self.get_lock(key) for key in self.referenced_locks]
         else:
             locks = list(self.locks.values())
-        with ThreadPoolExecutor(max_workers=min(len(locks), 100), thread_name_prefix="repository_loader") as executor:
-            executor.map(lambda l: l.sync(force_update=force_update), locks)
+        with ThreadPoolExecutor(
+                max_workers=min(len(locks) + 1, 100), thread_name_prefix="repository_loader"
+        ) as executor:
+            executor.map(lambda lock: lock.sync(force_update=force_update), locks)
         self.write()
-    
+
     def write(self):
-        with open(self.path, 'w') as f:
+        with open(self.path, "w") as f:
             for lock in self.locks.values():
-                f.write(str(lock) + '\n')
-
+                f.write(str(lock) + "\n")

hyperpocket/repository/repository.py

@@ -12,7 +12,32 @@ def pull(lockfile: Lockfile, urllike: str, git_ref: str):
         lockfile.add_lock(GitLock(repository_url=urllike, git_ref=git_ref))
     lockfile.sync(force_update=False)
     lockfile.write()
-    
+
+
 def sync(lockfile: Lockfile, force_update: bool):
     lockfile.sync(force_update=force_update)
     lockfile.write()
+
+
+def eject(url: str, ref: str, remote_path: str, lockfile: Lockfile):
+    path = pathlib.Path(url)
+    if path.exists():
+        raise ValueError("Local tools are already ejected")
+    else:
+        lock = lockfile.get_lock(GitLock(repository_url=url, git_ref=ref).key())
+
+    # first copy source to ejected directory
+    working_dir = pathlib.Path.cwd()
+    ejected_dir = working_dir / "ejected_tools"
+    if not ejected_dir.exists():
+        ejected_dir.mkdir()
+
+    tool_name = pathlib.Path(remote_path).name
+    eject_path = ejected_dir / tool_name
+    lock.eject_to_path(eject_path, remote_path)
+
+    # then update lockfile
+    lockfile.remove_lock(lock.key())
+    lockfile.add_lock(LocalLock(tool_path=str(eject_path)))
+    lockfile.sync(force_update=True)
+    lockfile.write()

hyperpocket/server/auth/__init__.py

@@ -1,11 +1,5 @@
 from fastapi import APIRouter
 
-from hyperpocket.server.auth.github import github_auth_router
-from hyperpocket.server.auth.google import google_auth_router
-from hyperpocket.server.auth.linear import linear_auth_router
-from hyperpocket.server.auth.slack import slack_auth_router
-from hyperpocket.server.auth.token import token_router
-from hyperpocket.server.auth.calendly import calendly_auth_router
 from hyperpocket.util.get_objects_from_subpackage import get_objects_from_subpackage
 
 auth_router = APIRouter(prefix="/auth")

hyperpocket/server/auth/gumloop.py

@@ -0,0 +1,16 @@
+from fastapi import APIRouter, Request
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+gumloop_auth_router = APIRouter(prefix="/gumloop")
+
+
+@gumloop_auth_router.get("/token/callback")
+async def gumloop_oauth2_callback(request: Request, state: str, token: str):
+    try:
+        FutureStore.resolve_future(state, token)
+    except ValueError:
+        return HTMLResponse(content="failed")
+
+    return HTMLResponse(content="success")

hyperpocket/server/auth/notion.py

@@ -0,0 +1,19 @@
+
+from fastapi import APIRouter
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+notion_auth_router = APIRouter(
+    prefix="/notion"
+)
+
+
+@notion_auth_router.get("/token/callback")
+async def notion_token_callback(state: str, token: str):
+    try:
+        FutureStore.resolve_future(state, token)
+    except ValueError:
+        return HTMLResponse(content="failed")
+
+    return HTMLResponse(content="success")         

hyperpocket/server/auth/reddit.py

@@ -0,0 +1,16 @@
+from fastapi import APIRouter
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+reddit_auth_router = APIRouter(prefix="/reddit")
+
+
+@reddit_auth_router.get("/oauth2/callback")
+async def reddit_oauth2_callback(state: str, code: str):
+    try:
+        FutureStore.resolve_future(state, code)
+    except ValueError:
+        return HTMLResponse(content="failed")
+
+    return HTMLResponse(content="success")