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

hyperpocket/pocket_auth.py

@@ -3,13 +3,13 @@ import enum
 import uuid
 from typing import Optional, Type
 
-from hyperpocket.auth import AuthProvider, PREBUILT_AUTH_HANDLERS
+from hyperpocket.auth import PREBUILT_AUTH_HANDLERS, AuthProvider
 from hyperpocket.auth.context import AuthContext
-from hyperpocket.auth.handler import AuthHandlerInterface, AuthenticateRequest
+from hyperpocket.auth.handler import AuthenticateRequest, AuthHandlerInterface
 from hyperpocket.config import config, pocket_logger
 from hyperpocket.futures import FutureStore
 from hyperpocket.session import SESSION_STORAGE_LIST
-from hyperpocket.session.interface import SessionStorageInterface, BaseSessionValue
+from hyperpocket.session.interface import BaseSessionValue, SessionStorageInterface
 
 
 class AuthState(enum.Enum):
@@ -26,11 +26,12 @@ class PocketAuth(object):
     session_storage: SessionStorageInterface
 
     def __init__(
-            self,
-            handlers: Optional[list[Type[AuthHandlerInterface]]] = None,
-            session_storage: Optional[SessionStorageInterface] = None,
-            use_prebuilt_handlers: bool = None):
-        if config.auth.use_prebuilt_auth or use_prebuilt_handlers:
+        self,
+        handlers: Optional[list[Type[AuthHandlerInterface]]] = None,
+        session_storage: Optional[SessionStorageInterface] = None,
+        use_prebuilt_handlers: bool = None,
+    ):
+        if config().auth.use_prebuilt_auth or use_prebuilt_handlers:
             handlers = PREBUILT_AUTH_HANDLERS + (handlers or [])
         handler_impls = [C() for C in handlers] if handlers else []
         self.handlers = {handler.name: handler for handler in handler_impls}
@@ -38,21 +39,31 @@ class PocketAuth(object):
             self.session_storage = session_storage
         else:
             for session_type in SESSION_STORAGE_LIST:
-                if session_type.session_storage_type() == config.session.session_type:
-                    session_config = getattr(config.session, config.session.session_type.value)
-
-                    pocket_logger.info(f"init {session_type.session_storage_type()} session storage..")
+                if session_type.session_storage_type() == config().session.session_type:
+                    session_config = getattr(
+                        config().session, config().session.session_type.value
+                    )
+
+                    pocket_logger.info(
+                        f"init {session_type.session_storage_type()} session storage.."
+                    )
                     self.session_storage = session_type(session_config)
 
             if self.session_storage is None:
-                pocket_logger.error(f"not supported session type({config.session.session_type})")
-                raise RuntimeError(f"Not Supported Session Type({config.session.session_type})")
-
-    def make_request(self,
-                     auth_scopes: list[str] = None,
-                     auth_handler_name: Optional[str] = None,
-                     auth_provider: Optional[AuthProvider] = None,
-                     **kwargs) -> AuthenticateRequest:
+                pocket_logger.error(
+                    f"not supported session type({config().session.session_type})"
+                )
+                raise RuntimeError(
+                    f"Not Supported Session Type({config().session.session_type})"
+                )
+
+    def make_request(
+        self,
+        auth_scopes: list[str] = None,
+        auth_handler_name: Optional[str] = None,
+        auth_provider: Optional[AuthProvider] = None,
+        **kwargs,
+    ) -> AuthenticateRequest:
         """
         Make AuthenticationRequest based on authentication handler.
 
@@ -67,14 +78,16 @@ class PocketAuth(object):
         handler = self.find_handler_instance(auth_handler_name, auth_provider)
         return handler.make_request(auth_scopes, **kwargs)
 
-    def check(self,
-              auth_req: AuthenticateRequest,
-              auth_handler_name: Optional[str] = None,
-              auth_provider: Optional[AuthProvider] = None,
-              thread_id: str = "default",
-              profile: str = 'default',
-              *args, **kwargs
-              ) -> AuthState:
+    def check(
+        self,
+        auth_req: AuthenticateRequest,
+        auth_handler_name: Optional[str] = None,
+        auth_provider: Optional[AuthProvider] = None,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> AuthState:
         """
         Check current authentication state.
 
@@ -104,7 +117,9 @@ class PocketAuth(object):
         return auth_state
 
     @staticmethod
-    def get_session_state(session: Optional[BaseSessionValue], auth_req: Optional[AuthenticateRequest]) -> AuthState:
+    def get_session_state(
+        session: Optional[BaseSessionValue], auth_req: Optional[AuthenticateRequest]
+    ) -> AuthState:
         if not session:
             return AuthState.NO_SESSION
 
@@ -115,8 +130,9 @@ class PocketAuth(object):
 
             return AuthState.PENDING_RESOLVE
 
-        if auth_req is not None and not session.is_auth_applicable(auth_provider_name=session.auth_provider_name,
-                                                                   auth_req=auth_req):
+        if auth_req is not None and not session.is_auth_applicable(
+            auth_provider_name=session.auth_provider_name, auth_req=auth_req
+        ):
             return AuthState.DO_AUTH
 
         if session.is_near_expires():
@@ -124,13 +140,15 @@ class PocketAuth(object):
 
         return AuthState.SKIP_AUTH
 
-    def prepare(self,
-                auth_req: AuthenticateRequest,
-                auth_handler_name: Optional[str] = None,
-                auth_provider: Optional[AuthProvider] = None,
-                thread_id: str = "default",
-                profile: str = 'default',
-                **kwargs) -> Optional[str]:
+    def prepare(
+        self,
+        auth_req: AuthenticateRequest,
+        auth_handler_name: Optional[str] = None,
+        auth_provider: Optional[AuthProvider] = None,
+        thread_id: str = "default",
+        profile: str = "default",
+        **kwargs,
+    ) -> Optional[str]:
         """
         Prepare authentication.
 
@@ -158,8 +176,13 @@ class PocketAuth(object):
         )
 
         pocket_logger.debug(
-            f"[thread_id({thread_id}):profile({profile})] {auth_provider.name} provider current auth state in prepare : {auth_state}")
-        if auth_state in [AuthState.SKIP_AUTH, AuthState.DO_REFRESH, AuthState.RESOLVED]:
+            f"[thread_id({thread_id}):profile({profile})] {auth_provider.name} provider current auth state in prepare : {auth_state}"
+        )
+        if auth_state in [
+            AuthState.SKIP_AUTH,
+            AuthState.DO_REFRESH,
+            AuthState.RESOLVED,
+        ]:
             return None
 
         handler = self.find_handler_instance(auth_handler_name, auth_provider)
@@ -168,12 +191,13 @@ class PocketAuth(object):
         if session:
             scope = scope.union(session.auth_scopes)
 
-        modified_req = auth_req.model_copy(update={'auth_scopes': scope})
+        modified_req = auth_req.model_copy(update={"auth_scopes": scope})
 
         # session in pending
         if session and session.auth_resolve_uid:
             pocket_logger.debug(
-                f"[thread_id({thread_id}):profile({profile})] already exists pending session(auth_resolve_uid:{session.auth_resolve_uid}).")
+                f"[thread_id({thread_id}):profile({profile})] already exists pending session(auth_resolve_uid:{session.auth_resolve_uid})."
+            )
             future_uid = session.auth_resolve_uid
 
             # update session, in case of requesting new scopes before session pending resolved.
@@ -182,7 +206,7 @@ class PocketAuth(object):
                 future_uid=session.auth_resolve_uid,
                 profile=profile,
                 thread_id=thread_id,
-                scope=scope
+                scope=scope,
             )
         else:  # create new pending session
             future_uid = str(uuid.uuid4())
@@ -191,25 +215,34 @@ class PocketAuth(object):
                 future_uid=future_uid,
                 profile=profile,
                 thread_id=thread_id,
-                scope=scope)
+                scope=scope,
+            )
 
             pocket_logger.debug(
-                f"[thread_id({thread_id}):profile({profile})] create new pending session(auth_resolve_uid:{future_uid}).")
-            asyncio.create_task(self._check_session_pending_resolved(handler, thread_id, profile))
+                f"[thread_id({thread_id}):profile({profile})] create new pending session(auth_resolve_uid:{future_uid})."
+            )
+            asyncio.create_task(
+                self._check_session_pending_resolved(handler, thread_id, profile)
+            )
 
-        prepare_url = handler.prepare(modified_req, thread_id, profile, future_uid, **kwargs)
+        prepare_url = handler.prepare(
+            modified_req, thread_id, profile, future_uid, **kwargs
+        )
         pocket_logger.debug(
-            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s prepare_url : {prepare_url}.")
+            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s prepare_url : {prepare_url}."
+        )
 
         return prepare_url
 
-    async def authenticate_async(self,
-                                 auth_req: AuthenticateRequest,
-                                 auth_handler_name: Optional[str] = None,
-                                 auth_provider: Optional[AuthProvider] = None,
-                                 thread_id: str = "default",
-                                 profile: str = 'default',
-                                 **kwargs) -> AuthContext:
+    async def authenticate_async(
+        self,
+        auth_req: AuthenticateRequest,
+        auth_handler_name: Optional[str] = None,
+        auth_provider: Optional[AuthProvider] = None,
+        thread_id: str = "default",
+        profile: str = "default",
+        **kwargs,
+    ) -> AuthContext:
         """
         Performing authentication.
         It is performing authentication. and save the session in session storage.
@@ -243,42 +276,61 @@ class PocketAuth(object):
         session = self.session_storage.get(handler.provider(), thread_id, profile)
         if session is None:
             pocket_logger.warning(
-                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'.")
+                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'."
+            )
             raise RuntimeError(
-                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'.")
+                f"[thread_id({thread_id}):profile({profile})] Session can't find. session should exist in 'authenticate'."
+            )
 
         pocket_logger.debug(
-            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s auth state : {auth_state}")
+            f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name})'s auth state : {auth_state}"
+        )
         try:
             if auth_state == AuthState.SKIP_AUTH:
                 context = session.auth_context
             elif auth_state == AuthState.DO_REFRESH:
                 try:
                     context = await asyncio.wait_for(
-                        handler.refresh(auth_req=auth_req, context=session.auth_context, **kwargs), timeout=300)
+                        handler.refresh(
+                            auth_req=auth_req, context=session.auth_context, **kwargs
+                        ),
+                        timeout=300,
+                    )
                 except Exception as e:
                     self.session_storage.delete(handler.provider(), thread_id, profile)
                     FutureStore.delete_future(session.auth_resolve_uid)
 
                     pocket_logger.warning(
-                        f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name}) failed to refresh the token.")
-                    raise RuntimeError("Failed to refresh the token. Please re-authenticate.") from e
-            elif auth_state == AuthState.PENDING_RESOLVE or auth_state == AuthState.RESOLVED:
+                        f"[thread_id({thread_id}):profile({profile})] auth_handler({auth_handler_name}) failed to refresh the token."
+                    )
+                    raise RuntimeError(
+                        "Failed to refresh the token. Please re-authenticate."
+                    ) from e
+            elif (
+                auth_state == AuthState.PENDING_RESOLVE
+                or auth_state == AuthState.RESOLVED
+            ):
                 future_uid = session.auth_resolve_uid
                 context = await asyncio.wait_for(
-                    handler.authenticate(auth_req=auth_req, future_uid=future_uid, **kwargs), timeout=300)
+                    handler.authenticate(
+                        auth_req=auth_req, future_uid=future_uid, **kwargs
+                    ),
+                    timeout=300,
+                )
             else:
                 # maybe auth_state is either AuthState.NO_SESSION or AuthState.DO_AUTH
                 pocket_logger.warning(
-                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}")
+                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}"
+                )
                 raise RuntimeError(
-                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}")
+                    f"[thread_id({thread_id}):profile({profile})] Invalid State. 'authenticate' cannot be reached while in state {auth_state}"
+                )
 
             session = await self._set_session_active(
                 context=context,
                 provider=handler.provider(),
                 profile=profile,
-                thread_id=thread_id
+                thread_id=thread_id,
             )
 
             return session.auth_context
@@ -288,33 +340,50 @@ class PocketAuth(object):
             FutureStore.delete_future(session.auth_resolve_uid)
             raise e
 
-    def get_auth_context(self, auth_provider: AuthProvider, thread_id: str = "default", profile: str = "default",
-                         **kwargs) -> Optional[AuthContext]:
+    def get_auth_context(
+        self,
+        auth_provider: AuthProvider,
+        thread_id: str = "default",
+        profile: str = "default",
+        **kwargs,
+    ) -> Optional[AuthContext]:
         session = self.session_storage.get(auth_provider, thread_id, profile, **kwargs)
         if session is None:
             return None
 
         return session.auth_context
 
-    def list_session_state(self, thread_id: str, auth_provider: Optional[AuthProvider] = None):
-        session_list = self.session_storage.get_by_thread_id(thread_id=thread_id, auth_provider=auth_provider)
+    def list_session_state(
+        self, thread_id: str, auth_provider: Optional[AuthProvider] = None
+    ):
+        session_list = self.session_storage.get_by_thread_id(
+            thread_id=thread_id, auth_provider=auth_provider
+        )
         session_state_list = []
         for session in session_list:
             state = self.get_session_state(session=session, auth_req=None)
 
-            session_state_list.append({
-                "provider": session.auth_provider_name,
-                "scope": session.auth_scopes,
-                "state": state,
-            })
+            session_state_list.append(
+                {
+                    "provider": session.auth_provider_name,
+                    "scope": session.auth_scopes,
+                    "state": state,
+                }
+            )
 
         return session_state_list
 
-    def delete_session(self, auth_provider: AuthProvider, thread_id: str = "default", profile: str = "default") -> bool:
+    def delete_session(
+        self,
+        auth_provider: AuthProvider,
+        thread_id: str = "default",
+        profile: str = "default",
+    ) -> bool:
         return self.session_storage.delete(auth_provider, thread_id, profile)
 
-    def find_handler_instance(self, name: Optional[str] = None,
-                              auth_provider: Optional[AuthProvider] = None) -> AuthHandlerInterface:
+    def find_handler_instance(
+        self, name: Optional[str] = None, auth_provider: Optional[AuthProvider] = None
+    ) -> AuthHandlerInterface:
         if name:
             return self.handlers[name]
         if auth_provider:
@@ -323,19 +392,35 @@ class PocketAuth(object):
                     return handler
         raise ValueError("No handler found")
 
-    async def _check_session_pending_resolved(self, auth_handler: AuthHandlerInterface, thread_id: str = "default",
-                                              profile: str = "default", timeout_seconds=300, **kwargs):
+    async def _check_session_pending_resolved(
+        self,
+        auth_handler: AuthHandlerInterface,
+        thread_id: str = "default",
+        profile: str = "default",
+        timeout_seconds=300,
+        **kwargs,
+    ):
         await asyncio.sleep(timeout_seconds)
-        session = self.session_storage.get(auth_handler.provider(), thread_id, profile, **kwargs)
+        session = self.session_storage.get(
+            auth_handler.provider(), thread_id, profile, **kwargs
+        )
         if session.auth_resolve_uid is not None:
-            pocket_logger.info(f"session({session.auth_resolve_uid}) is not resolved yet and timeout. remove session")
+            pocket_logger.info(
+                f"session({session.auth_resolve_uid}) is not resolved yet and timeout. remove session"
+            )
             self.delete_session(auth_handler.provider(), thread_id, profile)
             FutureStore.delete_future(session.auth_resolve_uid)
 
         return
 
     def _upsert_pending_session(
-            self, auth_handler: AuthHandlerInterface, future_uid: str, profile: str, thread_id: str, scope: set[str]):
+        self,
+        auth_handler: AuthHandlerInterface,
+        future_uid: str,
+        profile: str,
+        thread_id: str,
+        scope: set[str],
+    ):
         return self.session_storage.set(
             auth_provider=auth_handler.provider(),
             thread_id=thread_id,
@@ -346,7 +431,9 @@ class PocketAuth(object):
             auth_resolve_uid=future_uid,
         )
 
-    async def _set_session_active(self, context: AuthContext, provider: AuthProvider, profile: str, thread_id: str):
+    async def _set_session_active(
+        self, context: AuthContext, provider: AuthProvider, profile: str, thread_id: str
+    ):
         session = self.session_storage.get(provider, thread_id, profile)
         if session is None:
             pocket_logger.error("the session to be active doesn't exist.")

hyperpocket/pocket_core.py

@@ -1,12 +1,12 @@
 import asyncio
 import pathlib
-from typing import Any, Optional, Callable, List, Union
+from typing import Any, Callable, List, Optional, Union
 
 from hyperpocket.builtin import get_builtin_tools
 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.repository.lock import GitLock, LocalLock
 from hyperpocket.tool import Tool, ToolRequest
 from hyperpocket.tool.function import from_func
 from hyperpocket.tool.wasm import WasmTool
@@ -18,11 +18,13 @@ class PocketCore:
     auth: PocketAuth
     tools: dict[str, Tool]
 
-    def __init__(self,
-                 tools: list[ToolLike],
-                 auth: PocketAuth = None,
-                 lockfile_path: Optional[str] = None,
-                 force_update: bool = False):
+    def __init__(
+        self,
+        tools: list[ToolLike],
+        auth: PocketAuth = None,
+        lockfile_path: Optional[str] = None,
+        force_update: bool = False,
+    ):
         if auth is None:
             auth = PocketAuth()
         self.auth = auth
@@ -37,8 +39,10 @@ class PocketCore:
                 if pathlib.Path(tool_like).exists():
                     lock = LocalLock(tool_like)
                     req = WasmToolRequest(lock, "")
-                else:
-                    base_repo_url, git_ref, rel_path = GitLock.parse_repo_url(repo_url=tool_like)
+                elif tool_like.startswith("https://github.com"):
+                    base_repo_url, git_ref, rel_path = GitLock.parse_repo_url(
+                        repo_url=tool_like
+                    )
                     lock = GitLock(repository_url=base_repo_url, git_ref=git_ref)
                     req = WasmToolRequest(lock=lock, rel_path=rel_path, tool_vars={})
 
@@ -63,19 +67,26 @@ class PocketCore:
                 raise ValueError(f"Duplicate tool name: {tool.name}")
             self.tools[tool.name] = tool
 
-        pocket_logger.info(f"All Registered Tools Loaded successfully. total registered tools : {len(self.tools)}")
+        pocket_logger.info(
+            f"All Registered Tools Loaded successfully. total registered tools : {len(self.tools)}"
+        )
 
         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)}")
-
-    async def acall(self,
-                    tool_name: str,
-                    body: Any,
-                    thread_id: str = 'default',
-                    profile: str = 'default',
-                    *args, **kwargs) -> tuple[str, bool]:
+        pocket_logger.info(
+            f"All BuiltIn Tools Loaded successfully. total tools : {len(self.tools)}"
+        )
+
+    async def acall(
+        self,
+        tool_name: str,
+        body: Any,
+        thread_id: str = "default",
+        profile: str = "default",
+        *args,
+        **kwargs,
+    ) -> tuple[str, bool]:
         """
         Invoke tool asynchronously, not that different from `Pocket.invoke`
         But this method is called only in subprocess.
@@ -105,11 +116,13 @@ class PocketCore:
         result = await self.tool_call(tool_name, body=body, envs=credentials, **kwargs)
         return result, False
 
-    def prepare_auth(self,
-                     tool_name: Union[str, List[str]],
-                     thread_id: str = 'default',
-                     profile: str = 'default',
-                     **kwargs) -> Optional[str]:
+    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.
@@ -142,12 +155,14 @@ class PocketCore:
         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}")
+                    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}")
+                    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:
@@ -156,7 +171,8 @@ class PocketCore:
         auth_req = self.auth.make_request(
             auth_handler_name=auth_handler_name,
             auth_provider=auth_provider,
-            auth_scopes=list(auth_scopes))
+            auth_scopes=list(auth_scopes),
+        )
 
         return self.auth.prepare(
             auth_req=auth_req,
@@ -164,15 +180,16 @@ class PocketCore:
             auth_provider=auth_provider,
             thread_id=thread_id,
             profile=profile,
-            **kwargs
+            **kwargs,
         )
 
     async def authenticate(
-            self,
-            tool_name: str,
-            thread_id: str = 'default',
-            profile: str = 'default',
-            **kwargs) -> dict[str, str]:
+        self,
+        tool_name: str,
+        thread_id: str = "default",
+        profile: str = "default",
+        **kwargs,
+    ) -> dict[str, str]:
         """
         Authenticates the handler included in the tool and returns credentials.
 
@@ -190,7 +207,8 @@ class PocketCore:
         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_scopes=tool.auth.scopes,
+        )
         auth_ctx = await self.auth.authenticate_async(
             auth_req=auth_req,
             auth_handler_name=tool.auth.auth_handler,