hyperpocket 0.3.7__py3-none-any.whl → 0.4.0__py3-none-any.whl

hyperpocket/server/server.py

@@ -10,7 +10,6 @@ from uvicorn import Config, Server
 from hyperpocket.config import config, pocket_logger
 from hyperpocket.pocket_core import PocketCore
 from hyperpocket.server.auth import auth_router
-from hyperpocket.server.tool import tool_router
 
 
 class PocketServerOperations(enum.Enum):
@@ -18,10 +17,12 @@ class PocketServerOperations(enum.Enum):
     PREPARE_AUTH = "prepare_auth"
     AUTHENTICATE = "authenticate"
     TOOL_CALL = "tool_call"
+    PLUG_CORE = "plug_core"
 
 
 class PocketServer(object):
-    main_server: Server
+    fastapi_app: Optional[FastAPI]
+    main_server: Optional[Server]
     internal_server_port: int
     proxy_server: Optional[Server]
     proxy_port: int
@@ -29,15 +30,64 @@ class PocketServer(object):
     process: mp.Process
     future_store: dict[str, asyncio.Future]
     torn_down: bool = False
-
-    def __init__(
-        self,
-        internal_server_port: int = config().internal_server_port,
-        proxy_port: int = config().public_server_port,
-    ):
-        self.internal_server_port = internal_server_port
-        self.proxy_port = proxy_port
+    _uidset: set
+    _cores: dict[str, PocketCore]
+
+    def __init__(self):
+        self._initialized = False
+        self.internal_server_port = config().internal_server_port
+        self.proxy_port = config().public_server_port
+        self._uidset = set()
         self.future_store = dict()
+        self.fastapi_app = None
+        self.main_server = None
+        self._cores = dict()
+
+    # should be called in child process
+    def _plug_core(self, pocket_uid: str, pocket_core: PocketCore, *_a, **_kw):
+        # extend http routers from each docks
+        dock_routes = set([str(r) for r in self.fastapi_app.routes])
+
+        for dock in pocket_core.docks:
+            # check duplicated api route
+            dock_route = set([str(r) for r in dock.router.routes])
+            if dock_route in dock_routes:
+                continue
+
+            dock_routes.update(dock_route)
+            self.fastapi_app.include_router(dock.router)
+
+        # keep pocket core
+        self._cores[pocket_uid] = pocket_core
+
+    # should be called in parent process
+    async def plug_core(self, pocket_uid: str, pocket_core: PocketCore):
+        await self.call_in_subprocess(
+            PocketServerOperations.PLUG_CORE,
+            pocket_uid,
+            tuple(),
+            {
+                "pocket_uid": pocket_uid,
+                "pocket_core": pocket_core,
+            },
+        )
+
+    @classmethod
+    def get_instance_and_refcnt_up(cls, uid: str):
+        if cls.__dict__.get("_instance") is None:
+            cls._instance = cls()
+            cls._instance.run()
+        cls._instance.refcnt_up(uid)
+        return cls._instance
+
+    def refcnt_up(self, uid: str):
+        self._uidset.add(uid)
+
+    def refcnt_down(self, uid: str):
+        if uid in self._uidset:
+            self._uidset.remove(uid)
+            if len(self._uidset) == 0:
+                self._instance.teardown()
 
     def teardown(self):
         # @XXX(seokju) is it ok to call this method both in __del__ and __exit__?
@@ -63,20 +113,21 @@ class PocketServer(object):
         loop = asyncio.get_running_loop()
         _, conn = self.pipe
 
-        async def _acall(_conn, _op, _uid, a, kw):
+        async def _acall(_conn, _future_uid, _core_uid, _args, _kwargs):
             try:
-                result = await self.pocket_core.acall(*a, **kw)
+                core = self._cores[_core_uid]
+                result = await core.acall(*_args, **_kwargs)
                 error = None
             except Exception as e:
-                pocket_logger.error(f"failed to acall in pocket subprocess. error: {e}")
+                pocket_logger.error(f"failed in pocket subprocess. error: {e}")
                 result = None
                 error = e
+            _conn.send((_future_uid, result, error))
 
-            _conn.send((_op, _uid, result, error))
-
-        async def _prepare(_conn, _op, _uid, a, kw):
+        async def _prepare(_conn, _future_uid, _core_uid, a, kw):
             try:
-                result = self.pocket_core.prepare_auth(*a, **kw)
+                core = self._cores[_core_uid]
+                result = core.prepare_auth(*a, **kw)
                 error = None
             except Exception as e:
                 pocket_logger.error(
@@ -85,11 +136,12 @@ class PocketServer(object):
                 result = None
                 error = e
 
-            _conn.send((_op, _uid, result, error))
+            _conn.send((_future_uid, result, error))
 
-        async def _authenticate(_conn, _op, _uid, a, kw):
+        async def _authenticate(_conn, _future_uid, _core_uid, a, kw):
             try:
-                result = await self.pocket_core.authenticate(*a, **kw)
+                core = self._cores[_core_uid]
+                result = await core.authenticate(*a, **kw)
                 error = None
             except Exception as e:
                 pocket_logger.error(
@@ -98,11 +150,12 @@ class PocketServer(object):
                 result = None
                 error = e
 
-            _conn.send((_op, _uid, result, error))
+            _conn.send((_future_uid, result, error))
 
-        async def _tool_call(_conn, _op, _uid, a, kw):
+        async def _tool_call(_conn, _future_uid, _core_uid, a, kw):
             try:
-                result = await self.pocket_core.tool_call(*a, **kw)
+                core = self._cores[_core_uid]
+                result = await core.tool_call(*a, **kw)
                 error = None
             except Exception as e:
                 pocket_logger.error(
@@ -111,38 +164,50 @@ class PocketServer(object):
                 result = None
                 error = e
 
-            _conn.send((_op, _uid, result, error))
+            _conn.send((_future_uid, result, error))
+
+        async def _plug_core(_conn, _future_uid, _core_uid, a, kw):
+            try:
+                self._plug_core(*a, **kw)
+                _conn.send((_future_uid, None, None))
+            except Exception as e:
+                pocket_logger.error(
+                    f"failed to plug_core in pocket subprocess. error: {e}"
+                )
+                _conn.send((_future_uid, None, e))
 
         while True:
             if conn.poll():
-                op, uid, args, kwargs = conn.recv()
-                if op == PocketServerOperations.CALL.value:
-                    loop.create_task(_acall(conn, op, uid, args, kwargs))
-                elif op == PocketServerOperations.PREPARE_AUTH.value:
-                    loop.create_task(_prepare(conn, op, uid, args, kwargs))
-                elif op == PocketServerOperations.AUTHENTICATE.value:
-                    loop.create_task(_authenticate(conn, op, uid, args, kwargs))
-                elif op == PocketServerOperations.TOOL_CALL.value:
-                    loop.create_task(_tool_call(conn, op, uid, args, kwargs))
+                op, future_uid, core_uid, args, kwargs = conn.recv()
+                if op == PocketServerOperations.CALL:
+                    loop.create_task(_acall(conn, future_uid, core_uid, args, kwargs))
+                elif op == PocketServerOperations.PREPARE_AUTH:
+                    loop.create_task(_prepare(conn, future_uid, core_uid, args, kwargs))
+                elif op == PocketServerOperations.AUTHENTICATE:
+                    loop.create_task(_authenticate(conn, future_uid, core_uid, args, kwargs))
+                elif op == PocketServerOperations.TOOL_CALL:
+                    loop.create_task(_tool_call(conn, future_uid, core_uid, args, kwargs))
+                elif op == PocketServerOperations.PLUG_CORE:
+                    loop.create_task(_plug_core(conn, future_uid, core_uid, args, kwargs))
                 else:
-                    raise AttributeError(f"Can't find operations. op:{op}")
+                    raise ValueError(f"Unknown operation: {op}")
             else:
                 await asyncio.sleep(0)
 
-    def send_in_parent(self, op: PocketServerOperations, args: tuple, kwargs: dict):
+    def send_in_parent(self, op: PocketServerOperations, pocket_uid: str, args: tuple, kwargs: dict):
         conn, _ = self.pipe
-        uid = str(uuid.uuid4())
-        message = (op.value, uid, args, kwargs)
+        future_uid = str(uuid.uuid4())
+        message = (op, future_uid, pocket_uid, args, kwargs)
         future = asyncio.Future()
-        self.future_store[uid] = future
+        self.future_store[future_uid] = future
         conn.send(message)
-        return uid
+        return future_uid
 
     async def poll_in_parent(self):
         conn, _ = self.pipe
         while True:
             if conn.poll():
-                op, uid, result, error = conn.recv()
+                uid, result, error = conn.recv()
                 future = self.future_store[uid]
                 if error:
                     future.set_exception(error)
@@ -153,19 +218,21 @@ class PocketServer(object):
                 await asyncio.sleep(0)
 
     async def call_in_subprocess(
-        self, op: PocketServerOperations, args: tuple, kwargs: dict
+        self, op: PocketServerOperations, pocket_uid: str, args: tuple, kwargs: dict
     ):
-        uid = self.send_in_parent(op, args, kwargs)
+        uid = self.send_in_parent(op, pocket_uid, args, kwargs)
         loop = asyncio.get_running_loop()
         loop.create_task(self.poll_in_parent())
         return await self.future_store[uid]
 
-    def run(self, pocket_core: PocketCore):
+    def run(self):
         self._set_mp_start_method()
 
         error_queue = mp.Queue()
         self.pipe = mp.Pipe()
-        self.process = mp.Process(target=self._run, args=(pocket_core,))
+        self.process = mp.Process(
+            target=self._run
+        )
         self.process.start()  # process start
 
         if not error_queue.empty():
@@ -174,27 +241,25 @@ class PocketServer(object):
 
     def _report_initialized(self, error: Optional[Exception] = None):
         _, conn = self.pipe
-        conn.send(
-            (
-                "server-initialization",
-                error,
-            )
-        )
+        conn.send(('server-initialization', error,))
 
     def wait_initialized(self):
+        if self._initialized:
+            return
         conn, _ = self.pipe
         while True:
             if conn.poll():
                 _, error = conn.recv()
                 if error:
                     raise error
-                return
+                break
+        self._initialized = True
 
-    def _run(self, pocket_core):
+    def _run(self):
         try:
             # init process
-            self.pocket_core = pocket_core
-            self.main_server = self._create_main_server()
+            self.fastapi_app = self._create_fastapi_app()
+            self.main_server = self._create_main_server(self.fastapi_app)
             self.proxy_server = self._create_https_proxy_server()
             self._report_initialized()
 
@@ -204,20 +269,21 @@ class PocketServer(object):
         except Exception as error:
             self._report_initialized(error)
 
-    def _create_main_server(self) -> Server:
+    def _create_fastapi_app(self) -> FastAPI:
         app = FastAPI()
+        app.include_router(auth_router)
+        app.add_api_route("/health", lambda: {"status": "ok"}, methods=["GET"])
+        return app
+
+    def _create_main_server(self, app: FastAPI) -> Server:
         _config = Config(
             app,
             host="0.0.0.0",
             port=self.internal_server_port,
             log_level=config().log_level,
         )
-        app.include_router(tool_router)
-        app.include_router(auth_router)
-        app.add_api_route("/health", lambda: {"status": "ok"}, methods=["GET"])
-
-        app = Server(_config)
-        return app
+        server = Server(_config)
+        return server
 
     def _create_https_proxy_server(self) -> Optional[Server]:
         if not config().enable_local_callback_proxy:

hyperpocket/session/in_memory.py

@@ -19,10 +19,20 @@ InMemorySessionValue = BaseSessionValue
 class InMemorySessionStorage(
     SessionStorageInterface[InMemorySessionKey, InMemorySessionValue]
 ):
-    # TODO(moon) : Force it to always take SessionConfig as an input
+    _instance = None
+    _is_initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(InMemorySessionStorage, cls).__new__(cls)
+            cls._instance._is_initialized = False
+        return cls._instance
+
     def __init__(self, session_config: SessionConfigInMemory):
-        super().__init__()
-        self.storage: Dict[InMemorySessionKey, InMemorySessionValue] = {}
+        if not self._is_initialized:
+            super().__init__()
+            self.storage: Dict[InMemorySessionKey, InMemorySessionValue] = {}
+            self._is_initialized = True
 
     @classmethod
     def session_storage_type(cls) -> SessionType:

hyperpocket/tool/__init__.py

@@ -1,13 +1,10 @@
 from hyperpocket.tool.function import from_dock, from_func, function_tool
 from hyperpocket.tool.tool import Tool, ToolAuth, ToolRequest
-from hyperpocket.tool.wasm.tool import from_git, from_local
 
 __all__ = [
     "Tool",
     "ToolAuth",
     "ToolRequest",
-    "from_local",
-    "from_git",
     "from_dock",
     "from_func",
     "function_tool",

hyperpocket/tool/dock/__init__.py

@@ -0,0 +1,3 @@
+from hyperpocket.tool.dock.dock import Dock
+
+__all__ = ["Dock"]

hyperpocket/tool/dock/dock.py

@@ -0,0 +1,34 @@
+import abc
+from typing import Any
+
+from fastapi import APIRouter
+
+from hyperpocket.tool import ToolRequest
+from hyperpocket.tool.function import FunctionTool
+
+
+class Dock(abc.ABC):
+    _tool_requests: list[ToolRequest]
+    _dock_http_router: APIRouter
+    _dock_vars: dict[str, str]
+
+    def __init__(self, dock_vars: dict[str, str] = None):
+        self._dock_http_router = APIRouter()
+        self._tool_requests = []
+        self._dock_vars = dock_vars if dock_vars is not None else {}
+
+    @property
+    def router(self):
+        return self._dock_http_router
+
+    @abc.abstractmethod
+    def plug(self, req_like: Any, **kwargs):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def tools(self) -> list[FunctionTool]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def teardown(self):
+        raise NotImplementedError

hyperpocket/tool/function/__init__.py

@@ -4,4 +4,4 @@ from hyperpocket.tool.function.tool import FunctionTool
 from_func = FunctionTool.from_func
 from_dock = FunctionTool.from_dock
 
-__all__ = ["from_func", "from_dock", "function_tool"]
+__all__ = ["from_func", "from_dock", "function_tool", "FunctionTool"]

hyperpocket/tool/function/tool.py

@@ -17,19 +17,36 @@ class FunctionTool(Tool):
 
     func: Optional[Callable[..., str]]
     afunc: Optional[Callable[..., Coroutine[Any, Any, str]]]
+    keep_structured_arguments: bool = False
 
     def invoke(self, **kwargs) -> str:
         binding_args = self._get_binding_args(kwargs)
         if self.func is None:
             if self.afunc is None:
                 raise ValueError("Both func and afunc are None")
+            try:
+                loop = asyncio.get_running_loop()
+                return str(loop.run_until_complete(self.afunc(**binding_args)))
+            except RuntimeError:
+                pass
+            except Exception as e:
+                import traceback
+                traceback.print_exc()
+                traceback.print_stack()
+                return "There was an error while executing the tool: " + str(e)
             try:
                 return str(asyncio.run(self.afunc(**binding_args)))
             except Exception as e:
+                import traceback
+                traceback.print_exc()
+                traceback.print_stack()
                 return "There was an error while executing the tool: " + str(e)
         try:
             return str(self.func(**binding_args))
         except Exception as e:
+            import traceback
+            traceback.print_exc()
+            traceback.print_stack()
             return "There was an error while executing the tool: " + str(e)
 
     async def ainvoke(self, **kwargs) -> str:
@@ -39,9 +56,16 @@ class FunctionTool(Tool):
             binding_args = self._get_binding_args(kwargs)
             return str(await self.afunc(**binding_args))
         except Exception as e:
+            import traceback
+            traceback.print_exc()
+            traceback.print_stack()
             return "There was an error while executing the tool: " + str(e)
 
     def _get_binding_args(self, kwargs):
+        if self.keep_structured_arguments:
+            if kwargs.get("envs") is not None:
+                kwargs["envs"] |= self.tool_vars
+            return kwargs
         _kwargs = copy.deepcopy(kwargs)
 
         # make body args to model
@@ -54,28 +78,23 @@ class FunctionTool(Tool):
 
         # binding args
         binding_args = {}
-        if self.func.__dict__.get("__model__") is not None:
-            # when a function signature is not inferrable from the function itself
-            binding_args = args.copy()
-            binding_args |= _kwargs.get("envs", {}) | self.tool_vars
-        else:
-            sig = inspect.signature(self.func)
-            for param_name, param in sig.parameters.items():
-                if param_name not in args:
-                    continue
+        sig = inspect.signature(self.func)
+        for param_name, param in sig.parameters.items():
+            if param_name not in args:
+                continue
 
-                if param.kind == param.VAR_KEYWORD:
-                    # var keyword args should be passed by plain dict
-                    binding_args |= args[param_name]
-                    binding_args |= _kwargs.get("envs", {}) | self.tool_vars
+            if param.kind == param.VAR_KEYWORD:
+                # var keyword args should be passed by plain dict
+                binding_args |= args[param_name]
+                binding_args |= _kwargs.get("envs", {}) | self.tool_vars
 
-                    if "envs" in _kwargs:
-                        _kwargs.pop("envs")
+                if "envs" in _kwargs:
+                    _kwargs.pop("envs")
 
-                    binding_args |= _kwargs  # add other kwargs
-                    continue
+                binding_args |= _kwargs  # add other kwargs
+                continue
 
-                binding_args[param_name] = args[param_name]
+            binding_args[param_name] = args[param_name]
 
         return binding_args
 
@@ -89,11 +108,15 @@ class FunctionTool(Tool):
 
     @classmethod
     def from_func(
-        cls,
-        func: Callable | "FunctionTool",
-        afunc: Callable[..., Coroutine[Any, Any, str]] | "FunctionTool" = None,
-        auth: Optional[ToolAuth] = None,
-        tool_vars: dict[str, str] = None,
+            cls,
+            func: Callable | "FunctionTool",
+            afunc: Callable[..., Coroutine[Any, Any, str]] | "FunctionTool" = None,
+            name: str = None,
+            description: str = None,
+            json_schema: dict[str, Any] = None,
+            auth: Optional[ToolAuth] = None,
+            tool_vars: dict[str, str] = None,
+            keep_structured_arguments: bool = False,
     ) -> "FunctionTool":
         if tool_vars is None:
             tool_vars = dict()
@@ -109,28 +132,35 @@ class FunctionTool(Tool):
         elif not callable(func) and not callable(afunc):
             raise ValueError("FunctionTool can only be created from a callable")
 
-        model = function_to_model(func)
-        argument_json_schema = flatten_json_schema(model.model_json_schema())
+        name = name or (func and func.__name__) or (afunc and afunc.__name__)
+        description = description or (func and func.__doc__) or (afunc and func.__doc__)
+        schema = json_schema or \
+                 (func and function_to_model(func).model_json_schema()) or \
+                 (afunc and function_to_model(afunc).model_json_schema())
+        argument_json_schema = flatten_json_schema(schema)
 
         return cls(
             func=func,
             afunc=afunc,
-            name=func.__name__,
-            description=func.__doc__ if func.__doc__ is not None else "",
+            name=name,
+            description=description,
             argument_json_schema=argument_json_schema,
             auth=auth,
             default_tool_vars=tool_vars,
+            keep_structured_arguments=keep_structured_arguments,
         )
 
     @classmethod
     def from_dock(
-        cls,
-        dock: list[Callable[..., str]],
-        tool_vars: Optional[dict[str, str]] = None,
+            cls,
+            dock: list[Callable[..., str]],
+            tool_vars: Optional[dict[str, str]] = None,
     ) -> list["FunctionTool"]:
         if tool_vars is None:
             tool_vars = dict()
         tools = []
+
+        # @moon: will be refactored.
         for func in dock:
             if (_model := func.__dict__.get("__model__")) is not None:
                 model = _model
@@ -155,7 +185,7 @@ class FunctionTool(Tool):
                         argument_json_schema=argument_json_schema,
                         auth=auth,
                         default_tool_vars=(
-                            tool_vars | func.__dict__.get("__vars__", {})
+                                tool_vars | func.__dict__.get("__vars__", {})
                         ),
                     )
                 )
@@ -169,7 +199,7 @@ class FunctionTool(Tool):
                         argument_json_schema=argument_json_schema,
                         auth=auth,
                         default_tool_vars=(
-                            tool_vars | func.__dict__.get("__vars__", {})
+                                tool_vars | func.__dict__.get("__vars__", {})
                         ),
                     )
                 )

hyperpocket/tool/tool.py

@@ -116,21 +116,13 @@ class Tool(BaseModel, abc.ABC):
             return self.description
 
     def override_tool_variables(self, override_vars: dict[str, str]) -> "Tool":
-        self.overridden_tool_vars = override_vars
+        self.overridden_tool_vars |= override_vars
         return self
 
     @property
     def tool_vars(self) -> dict[str, str]:
         return self.default_tool_vars | self.overridden_tool_vars
 
-    @classmethod
-    def from_tool_request(cls, tool_req: ToolRequest, **kwargs) -> "Tool":
-        from hyperpocket.tool.wasm.tool import WasmTool, WasmToolRequest
-
-        if isinstance(tool_req, WasmToolRequest):
-            return WasmTool.from_tool_request(tool_req, **kwargs)
-        raise ValueError("Unknown tool request type")
-
     @classmethod
     def _get_schema_model(
         cls, name: str, json_schema: Optional[dict], use_profile: bool

hyperpocket/tool_like.py

@@ -1,5 +1,6 @@
 from typing import Callable, Union
 
 from hyperpocket.tool import Tool, ToolRequest
+from hyperpocket.tool.dock import Dock
 
-ToolLike = Union[Tool, str, Callable, ToolRequest]
+ToolLike = Union[Tool, str, tuple, Callable, ToolRequest, Dock]

hyperpocket/util/generate_slug.py

@@ -0,0 +1,4 @@
+import string, random
+
+def generate_slug(length: int = 6) -> str:
+    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

hyperpocket/util/json_schema_to_model.py

@@ -60,13 +60,17 @@ def json_schema_to_model(
             fields["additional_properties"] = (dict[str, additional_model], {})
 
     # Create the model
-    model = create_model(model_name, **fields)
+    model = create_model(f"{model_name}", **fields)
 
     # Add custom Config class to handle extra properties
     class Config:
         extra = config_extra
 
     model.Config = Config
+    
+    # workaround for pickling dynamic classes
+    model.__module__ = "__main__"
+    model.__qualname__ = model.__name__.split('.')[-1]
 
     return model
 

{hyperpocket-0.3.7.dist-info → hyperpocket-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyperpocket
-Version: 0.3.7
+Version: 0.4.0
 Summary: Building AI agent with hyperpocket tool in a flash
 Project-URL: Homepage, https://vessl-ai.github.io/hyperpocket
 Project-URL: Repository, https://github.com/vessl-ai/hyperpocket
@@ -13,6 +13,7 @@ Requires-Dist: gitpython>=3.1.43
 Requires-Dist: httpx==0.27
 Requires-Dist: jinja2>=3.1.4
 Requires-Dist: multiprocess>=0.70.17
+Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: playwright>=1.49.0
 Requires-Dist: pydantic>=2.10.2
 Requires-Dist: pygithub>=2.5.0
@@ -21,6 +22,8 @@ Requires-Dist: redis>=5.2.1
 Requires-Dist: requests>=2.32.3
 Requires-Dist: toml>=0.10.2
 Requires-Dist: uvicorn>=0.32.1
+Provides-Extra: standard
+Requires-Dist: hyperdock-container; extra == 'standard'
 Description-Content-Type: text/markdown
 
 <p align="center">