hyperpocket 0.1.9__py3-none-any.whl → 0.2.0__py3-none-any.whl

hyperpocket/pocket_main.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import Any, Optional, List, Union
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, List, Optional, Union
 
 from hyperpocket.config import pocket_logger
 from hyperpocket.pocket_auth import PocketAuth
@@ -12,30 +13,38 @@ class Pocket(object):
     server: PocketServer
     core: PocketCore
 
-    def __init__(self,
-                 tools: list[ToolLike],
-                 auth: PocketAuth = None,
-                 lockfile_path: Optional[str] = None,
-                 force_update: bool = False,
-                 use_profile: bool = False):
+    def __init__(
+        self,
+        tools: list[ToolLike],
+        auth: PocketAuth = None,
+        lockfile_path: Optional[str] = None,
+        force_update: bool = False,
+        use_profile: bool = False,
+    ):
         self.use_profile = use_profile
-
         self.core = PocketCore(
             tools=tools,
             auth=auth,
             lockfile_path=lockfile_path,
-            force_update=force_update
+            force_update=force_update,
         )
-
         self.server = PocketServer()
         self.server.run(self.core)
-
-    def invoke(self,
-               tool_name: str,
-               body: Any,
-               thread_id: str = 'default',
-               profile: str = 'default',
-               *args, **kwargs) -> str:
+        try:
+            self.server.wait_initialized()
+        except Exception as e:
+            pocket_logger.error(f"Failed to initialize pocket server. error : {e}")
+            self._teardown_server()
+
+    def invoke(
+        self,
+        tool_name: str,
+        body: Any,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> str:
         """
         Invoke Tool synchronously
 
@@ -48,15 +57,19 @@ class Pocket(object):
         Returns:
             str: tool result
         """
-        result = asyncio.run(self.ainvoke(tool_name, body, thread_id, profile, *args, **kwargs))
-        return result
+        return asyncio.run(
+            self.ainvoke(tool_name, body, thread_id, profile, *args, **kwargs)
+        )
 
-    async def ainvoke(self,
-                      tool_name: str,
-                      body: Any,
-                      thread_id: str = 'default',
-                      profile: str = 'default',
-                      *args, **kwargs) -> str:
+    async def ainvoke(
+        self,
+        tool_name: str,
+        body: Any,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> str:
         """
         Invoke Tool asynchronously
 
@@ -79,12 +92,15 @@ class Pocket(object):
         )
         return result
 
-    def invoke_with_state(self,
-                          tool_name: str,
-                          body: Any,
-                          thread_id: str = 'default',
-                          profile: str = 'default',
-                          *args, **kwargs) -> tuple[str, bool]:
+    def invoke_with_state(
+        self,
+        tool_name: str,
+        body: Any,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> tuple[str, bool]:
         """
         Invoke Tool with state synchronously
         State indicates whether this tool is paused or not.
@@ -101,27 +117,33 @@ class Pocket(object):
         """
         try:
             loop = asyncio.new_event_loop()
-            result = loop.run_until_complete(
-                self.ainvoke_with_state(tool_name, body, thread_id, profile, *args, **kwargs))
-
-        except RuntimeError as e:
-            pocket_logger.warning("Can't execute sync def in event loop. use nest-asyncio")
+        except RuntimeError:
+            pocket_logger.warning(
+                "Can't execute sync def in event loop. use nest-asyncio"
+            )
 
             import nest_asyncio
+
             loop = asyncio.new_event_loop()
             nest_asyncio.apply(loop=loop)
 
-            result = loop.run_until_complete(
-                self.ainvoke_with_state(tool_name, body, thread_id, profile, *args, **kwargs))
+        result = loop.run_until_complete(
+            self.ainvoke_with_state(
+                tool_name, body, thread_id, profile, *args, **kwargs
+            )
+        )
 
         return result
 
-    async def ainvoke_with_state(self,
-                                 tool_name: str,
-                                 body: Any,
-                                 thread_id: str = 'default',
-                                 profile: str = 'default',
-                                 *args, **kwargs) -> tuple[str, bool]:
+    async def ainvoke_with_state(
+        self,
+        tool_name: str,
+        body: Any,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> tuple[str, bool]:
         """
         Invoke Tool with state synchronously
         State indicates whether this tool is paused or not.
@@ -140,10 +162,10 @@ class Pocket(object):
             PocketServerOperations.CALL,
             args,
             {
-                'tool_name': tool_name,
-                'body': body,
-                'thread_id': thread_id,
-                'profile': profile,
+                "tool_name": tool_name,
+                "body": body,
+                "thread_id": thread_id,
+                "profile": profile,
                 **kwargs,
             },
         )
@@ -152,10 +174,11 @@ class Pocket(object):
 
         return result, paused
 
-    async def initialize_tool_auth(self,
-                                   thread_id: str = 'default',
-                                   profile: str = 'default',
-                                   ) -> dict[str, str]:
+    async def initialize_tool_auth(
+        self,
+        thread_id: str = "default",
+        profile: str = "default",
+    ) -> dict[str, str]:
         """
         Initialize authentication for all tools.
 
@@ -177,19 +200,17 @@ class Pocket(object):
         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)
+                tool_name=tool_name_list, thread_id=thread_id, profile=profile
+            )
 
             if prepare is not None:
                 prepare_list[provider] = prepare
 
         return prepare_list
 
-    async def wait_tool_auth(self,
-                             thread_id: str = 'default',
-                             profile: str = 'default'
-                             ) -> bool:
+    async def wait_tool_auth(
+        self, thread_id: str = "default", profile: str = "default"
+    ) -> bool:
         """
         Wait until all tool authentications are completed.
 
@@ -214,10 +235,9 @@ class Pocket(object):
 
                 waiting_futures.append(
                     self.authenticate_in_subprocess(
-                        tool_name=tools[0].name,
-                        thread_id=thread_id,
-                        profile=profile
-                    ))
+                        tool_name=tools[0].name, thread_id=thread_id, profile=profile
+                    )
+                )
 
             await asyncio.gather(*waiting_futures)
 
@@ -227,56 +247,65 @@ class Pocket(object):
             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):
+    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,
+                "tool_name": tool_name,
+                "thread_id": thread_id,
+                "profile": profile,
                 **kwargs,
             },
         )
 
         return prepare
 
-    async def authenticate_in_subprocess(self,
-                                         tool_name: str,
-                                         thread_id: str = 'default',
-                                         profile: str = 'default',
-                                         *args, **kwargs):
+    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,
+                "tool_name": tool_name,
+                "thread_id": thread_id,
+                "profile": profile,
                 **kwargs,
             },
         )
 
         return credentials
 
-    async def tool_call_in_subprocess(self,
-                                      tool_name: str,
-                                      body: Any,
-                                      thread_id: str = 'default',
-                                      profile: str = 'default',
-                                      *args, **kwargs):
+    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,
+                "tool_name": tool_name,
+                "body": body,
+                "thread_id": thread_id,
+                "profile": profile,
                 **kwargs,
             },
         )
@@ -290,17 +319,17 @@ class Pocket(object):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.__dict__.get('server'):
+        if self.__dict__.get("server"):
             self.server.teardown()
 
     def __del__(self):
-        if self.__dict__.get('server'):
+        if self.__dict__.get("server"):
             self.server.teardown()
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        if 'server' in state:
-            del state['server']
+        if "server" in state:
+            del state["server"]
         return state
 
     def __setstate__(self, state):

hyperpocket/prompts.py

@@ -1,5 +1,5 @@
 def pocket_extended_tool_description(description: str):
-    return f'''
+    return f"""
 This tool functions as described in the <tool-description> XML tag:
 <tool-description>
 {description}
@@ -10,4 +10,4 @@ Arguments required:
 - '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, eject
+from hyperpocket.repository.repository import eject, pull, sync
 
 __all__ = ["Lock", "Lockfile", "pull", "sync", "eject"]

hyperpocket/repository/lock.py

@@ -1,12 +1,13 @@
 import abc
 import pathlib
 import shutil
-from typing import Optional, Tuple
+from typing import Optional, Tuple, ClassVar
 
 import git
 from pydantic import BaseModel, Field
+from pydantic.fields import ModelPrivateAttr
 
-from hyperpocket.config import settings, pocket_logger
+from hyperpocket.config import pocket_logger, settings
 
 
 class Lock(BaseModel, abc.ABC):
@@ -34,17 +35,19 @@ class Lock(BaseModel, abc.ABC):
 
 
 class LocalLock(Lock):
-    tool_source: str = Field(default='local')
+    tool_source: str = Field(default="local")
     tool_path: str
 
     def __init__(self, tool_path: str):
-        super().__init__(tool_source="local", tool_path=str(pathlib.Path(tool_path).resolve()))
+        super().__init__(
+            tool_source="local", tool_path=str(pathlib.Path(tool_path).resolve())
+        )
 
     def __str__(self):
         return f"local\t{self.tool_path}"
 
     def key(self):
-        return self.tool_source, self.tool_path.rstrip('/')
+        return self.tool_source, self.tool_path.rstrip("/")
 
     def sync(self, **kwargs):
         pocket_logger.info(f"Syncing path: {self.tool_path} ...")
@@ -55,11 +58,12 @@ class LocalLock(Lock):
 
     def toolpkg_path(self) -> pathlib.Path:
         pocket_pkgs = settings.toolpkg_path
-        return pocket_pkgs / 'local' / self.tool_path[1:]
+        return pocket_pkgs / "local" / self.tool_path[1:]
 
 
 class GitLock(Lock):
-    tool_source: str = 'git'
+    _remote_cache: ClassVar[dict[str, dict[str, str]]]
+    tool_source: str = "git"
     repository_url: str
     git_ref: str
     ref_sha: Optional[str] = None
@@ -68,17 +72,17 @@ class GitLock(Lock):
         return f"git\t{self.repository_url}\t{self.git_ref}\t{self.ref_sha}"
 
     def key(self):
-        return self.tool_source, self.repository_url.rstrip('/'), self.git_ref
+        return self.tool_source, self.repository_url.rstrip("/"), self.git_ref
 
     def toolpkg_path(self) -> pathlib.Path:
         if not self.ref_sha:
             raise ValueError("ref_sha is not set")
         cleansed_url = self.repository_url
-        if self.repository_url.startswith('http://'):
+        if self.repository_url.startswith("http://"):
             cleansed_url = self.repository_url[7:]
-        elif self.repository_url.startswith('https://'):
+        elif self.repository_url.startswith("https://"):
             cleansed_url = self.repository_url[8:]
-        elif self.repository_url.startswith('git@'):
+        elif self.repository_url.startswith("git@"):
             cleansed_url = self.repository_url[4:]
         return settings.toolpkg_path / cleansed_url / self.ref_sha
 
@@ -91,12 +95,16 @@ class GitLock(Lock):
            to align the local repository with the remote version.
         """
         try:
-            pocket_logger.info(f"Syncing git: {self.repository_url} @ ref: {self.git_ref} ...")
+            pocket_logger.info(
+                f"Syncing git: {self.repository_url} @ ref: {self.git_ref} ..."
+            )
 
             # get new sha from refs
             new_sha = self._get_new_sha_if_exists_in_remote()
             if new_sha is None:
-                raise ValueError(f"Could not find ref {self.git_ref} in {self.repository_url}")
+                raise ValueError(
+                    f"Could not find ref {self.git_ref} in {self.repository_url}"
+                )
 
             # check self.ref_sha should be updated
             if self.ref_sha != new_sha:
@@ -111,10 +119,10 @@ class GitLock(Lock):
             # init git repo in local and set origin url
             repo = git.Repo.init(pkg_version_path)
             try:
-                remote = repo.remote('origin')
+                remote = repo.remote("origin")
                 remote.set_url(self.repository_url)
             except ValueError:
-                remote = repo.create_remote('origin', self.repository_url)
+                remote = repo.create_remote("origin", self.repository_url)
 
             # check current local commit include new_sha
             # if not included, fetch and do hard reset
@@ -126,10 +134,12 @@ class GitLock(Lock):
             if exist_sha is None or exist_sha != self.ref_sha:
                 remote.fetch(depth=1, refspec=self.ref_sha)
                 repo.git.checkout(new_sha)
-                repo.git.reset('--hard', new_sha)
-                repo.git.clean('-fd')
+                repo.git.reset("--hard", new_sha)
+                repo.git.clean("-fd")
         except Exception as e:
-            pocket_logger.error(f"failed to sync git: {self.repository_url} @ ref: {self.git_ref}. reason : {e}")
+            pocket_logger.error(
+                f"failed to sync git: {self.repository_url} @ ref: {self.git_ref}. reason : {e}"
+            )
             raise e
 
     def _get_new_sha_if_exists_in_remote(self):
@@ -143,8 +153,8 @@ class GitLock(Lock):
         refs = git.cmd.Git().ls_remote(self.repository_url)
 
         new_sha = None
-        for r in refs.split('\n'):
-            sha, ref = r.split('\t')
+        for r in refs.split("\n"):
+            sha, ref = r.split("\t")
             if sha == self.ref_sha:
                 new_sha = sha
                 break
@@ -161,16 +171,19 @@ class GitLock(Lock):
 
     @classmethod
     def get_git_branches(cls, repo_url):
-        ls_lists = git.cmd.Git().ls_remote(repo_url)
-
-        branches = {}
-        for line in ls_lists.split("\n"):
-            sha, ref = line.split("\t")
-            if ref.startswith("refs/heads/"):
-                branch_name = ref.replace("refs/heads/", "")
-                branches[branch_name] = sha
-
-        return branches
+        if not hasattr(cls, "_remote_cache"):
+            cls._remote_cache = {}
+        if cls._remote_cache.get(repo_url) is None:
+            ls_lists = git.cmd.Git().ls_remote(repo_url)
+
+            branches = {}
+            for line in ls_lists.split("\n"):
+                sha, ref = line.split("\t")
+                if ref.startswith("refs/heads/"):
+                    branch_name = ref.replace("refs/heads/", "")
+                    branches[branch_name] = sha
+            cls._remote_cache[repo_url] = branches
+        return cls._remote_cache[repo_url]
 
     @classmethod
     def parse_repo_url(cls, repo_url: str) -> Tuple[str, str, str]:
@@ -195,7 +208,7 @@ class GitLock(Lock):
         # Parse base repo URL and remaining path
         tree_index = repo_path_list.index("tree")
         base_repo = f"https://github.com/{'/'.join(repo_path_list[:tree_index])}"
-        sub_path = repo_path_list[tree_index + 1:]
+        sub_path = repo_path_list[tree_index + 1 :]
 
         # Fetch branch information
         branches = cls.get_git_branches(base_repo)
@@ -204,14 +217,15 @@ class GitLock(Lock):
         for idx in range(1, len(sub_path) + 1):
             branch_name = "/".join(sub_path[:idx])
             if branch_name in branches:
-                directory_path = "/".join(sub_path[idx:]) if idx < len(sub_path) else None
+                directory_path = (
+                    "/".join(sub_path[idx:]) if idx < len(sub_path) else None
+                )
                 return base_repo, branch_name, directory_path
 
         # If no valid branch is found, raise an error
         raise ValueError("Branch not found in repository")
 
     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} ..."

hyperpocket/repository/lockfile.py

@@ -1,7 +1,7 @@
 import pathlib
 from concurrent.futures.thread import ThreadPoolExecutor
 
-from hyperpocket.repository.lock import Lock, LocalLock, GitLock
+from hyperpocket.repository.lock import GitLock, LocalLock, Lock
 
 
 class Lockfile:
@@ -51,7 +51,7 @@ class Lockfile:
         else:
             locks = list(self.locks.values())
         with ThreadPoolExecutor(
-                max_workers=min(len(locks) + 1, 100), thread_name_prefix="repository_loader"
+            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()

hyperpocket/repository/repository.py

@@ -1,6 +1,6 @@
 import pathlib
 
-from hyperpocket.repository.lock import LocalLock, GitLock
+from hyperpocket.repository.lock import GitLock, LocalLock
 from hyperpocket.repository.lockfile import Lockfile
 
 

hyperpocket/server/__init__.py

@@ -1,3 +1,3 @@
 from hyperpocket.server.proxy import add_callback_proxy
 
-__all__ = ["add_callback_proxy"]
+__all__ = ["add_callback_proxy"]

hyperpocket/server/auth/github.py

@@ -15,6 +15,7 @@ async def github_oauth2_callback(request: Request, state: str, code: str):
 
     return HTMLResponse(content="success")
 
+
 @github_auth_router.get("/token/callback")
 async def github_token_callback(request: Request, state: str, token: str):
     try:
@@ -22,4 +23,4 @@ async def github_token_callback(request: Request, state: str, token: str):
     except ValueError:
         return HTMLResponse(content="failed")
 
-    return HTMLResponse(content="success")
+    return HTMLResponse(content="success")

hyperpocket/server/auth/linear.py

@@ -3,9 +3,7 @@ from starlette.responses import HTMLResponse
 
 from hyperpocket.futures import FutureStore
 
-linear_auth_router = APIRouter(
-    prefix="/linear"
-)
+linear_auth_router = APIRouter(prefix="/linear")
 
 
 @linear_auth_router.get("/token/callback")

hyperpocket/server/auth/notion.py

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

hyperpocket/server/auth/slack.py

@@ -3,9 +3,7 @@ from starlette.responses import HTMLResponse
 
 from hyperpocket.futures import FutureStore
 
-slack_auth_router = APIRouter(
-    prefix="/slack"
-)
+slack_auth_router = APIRouter(prefix="/slack")
 
 
 @slack_auth_router.get("/oauth2/callback")