hyperpocket 0.0.1__py3-none-any.whl

hyperpocket/prompts.py

@@ -0,0 +1,15 @@
+def pocket_extended_tool_description(description: str):
+    return f'''
+This tool performs as stated in the <tool-description></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.
+'''

hyperpocket/repository/__init__.py

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

hyperpocket/repository/lock.py

@@ -0,0 +1,156 @@
+import abc
+import pathlib
+import shutil
+from typing import Optional
+
+import git
+from pydantic import BaseModel, Field
+
+from hyperpocket.config import settings, pocket_logger
+
+
+class Lock(BaseModel, abc.ABC):
+    tool_source: str = None
+
+    @abc.abstractmethod
+    def __str__(self):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def key(self) -> tuple[str, ...]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def sync(self, **kwargs):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def toolpkg_path(self) -> pathlib.Path:
+        raise NotImplementedError
+
+
+class LocalLock(Lock):
+    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()))
+
+    def __str__(self):
+        return f"local\t{self.tool_path}"
+
+    def key(self):
+        return self.tool_source, self.tool_path.rstrip('/')
+
+    def sync(self, **kwargs):
+        pocket_logger.info(f"Syncing path: {self.tool_path} ...")
+        pkg_path = self.toolpkg_path()
+        if pkg_path.exists():
+            shutil.rmtree(pkg_path)
+        shutil.copytree(self.tool_path, pkg_path)
+
+    def toolpkg_path(self) -> pathlib.Path:
+        pocket_pkgs = settings.toolpkg_path
+        return pocket_pkgs / 'local' / self.tool_path[1:]
+
+
+class GitLock(Lock):
+    tool_source: str = 'git'
+    repository_url: str
+    git_ref: str
+    ref_sha: Optional[str] = None
+
+    def __str__(self):
+        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
+
+    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://'):
+            cleansed_url = self.repository_url[7:]
+        elif self.repository_url.startswith('https://'):
+            cleansed_url = self.repository_url[8:]
+        elif self.repository_url.startswith('git@'):
+            cleansed_url = self.repository_url[4:]
+        return settings.toolpkg_path / cleansed_url / self.ref_sha
+
+    def sync(self, force_update: bool = False, **kwargs):
+        """
+        Synchronize the local git repository with the target remote branch.
+
+        1. Check if the SHA of the target ref in the remote repository matches the current local SHA.
+        2. If they do not match, fetch the target ref from the remote repository and do a hard reset
+           to align the local repository with the remote version.
+        """
+        try:
+            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}")
+
+            # check self.ref_sha should be updated
+            if self.ref_sha != new_sha:
+                if force_update or self.ref_sha is None:
+                    self.ref_sha = new_sha
+
+            # make pkg_version_path dir if not exists
+            pkg_version_path = self.toolpkg_path()
+            if not pkg_version_path.exists():
+                pkg_version_path.mkdir(parents=True)
+
+            # init git repo in local and set origin url
+            repo = git.Repo.init(pkg_version_path)
+            try:
+                remote = repo.remote('origin')
+                remote.set_url(self.repository_url)
+            except ValueError:
+                remote = repo.create_remote('origin', self.repository_url)
+
+            # check current local commit include new_sha
+            # if not included, fetch and do hard reset
+            exist_sha = None
+            try:
+                exist_sha = repo.head.commit.hexsha
+            except ValueError:
+                pass
+            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')
+        except Exception as 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):
+        """
+        get new sha in refs
+        First, check remote sha is matched to saved ref_sha
+        Second, check remote ref name is matched to saved ref name
+        Third, check local ref name is matched to saved ref name
+        And last, check tag ref name is matched to saved ref name
+        """
+        refs = git.cmd.Git().ls_remote(self.repository_url)
+
+        new_sha = None
+        for r in refs.split('\n'):
+            sha, ref = r.split('\t')
+            if sha == self.ref_sha:
+                new_sha = sha
+                break
+            elif ref == self.git_ref:
+                new_sha = sha
+                break
+            elif ref == f"refs/heads/{self.git_ref}":
+                new_sha = sha
+                break
+            elif ref == f"refs/tags/{self.git_ref}":
+                new_sha = sha
+                break
+        return new_sha

hyperpocket/repository/lockfile.py

@@ -0,0 +1,56 @@
+import pathlib
+from concurrent.futures.thread import ThreadPoolExecutor
+
+from hyperpocket.repository.lock import Lock, LocalLock, GitLock
+
+
+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:
+                for line in f:
+                    split = line.strip().split('\t')
+                    source = split[0]
+                    if source == 'local':
+                        lock = LocalLock(tool_path=split[1])
+                    elif source == 'git':
+                        lock = GitLock(
+                            repository_url=split[1],
+                            git_ref=split[2],
+                            ref_sha=split[3],
+                        )
+                    else:
+                        raise ValueError(f"Unknown tool source: {source}")
+                    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 get_lock(self, key: tuple[str, ...]):
+        return self.locks[key]
+
+    def sync(self, force_update: bool, referenced_only: bool = False):
+        if referenced_only:
+            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)
+        self.write()
+    
+    def write(self):
+        with open(self.path, 'w') as f:
+            for lock in self.locks.values():
+                f.write(str(lock) + '\n')
+

hyperpocket/repository/repository.py

@@ -0,0 +1,18 @@
+import pathlib
+
+from hyperpocket.repository.lock import LocalLock, GitLock
+from hyperpocket.repository.lockfile import Lockfile
+
+
+def pull(lockfile: Lockfile, urllike: str, git_ref: str):
+    path = pathlib.Path(urllike)
+    if path.exists():
+        lockfile.add_lock(LocalLock(tool_path=str(path)))
+    else:
+        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()

hyperpocket/server/__init__.py

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

hyperpocket/server/auth/__init__.py

@@ -0,0 +1,15 @@
+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")
+
+routers = get_objects_from_subpackage("hyperpocket.server.auth", APIRouter)
+for r in routers:
+    auth_router.include_router(r)

hyperpocket/server/auth/calendly.py

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

hyperpocket/server/auth/github.py

@@ -0,0 +1,25 @@
+from fastapi import APIRouter, Request
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+github_auth_router = APIRouter(prefix="/github")
+
+
+@github_auth_router.get("/oauth2/callback")
+async def github_oauth2_callback(request: Request, state: str, code: str):
+    try:
+        FutureStore.resolve_future(state, code)
+    except ValueError:
+        return HTMLResponse(content="failed")
+
+    return HTMLResponse(content="success")
+
+@github_auth_router.get("/token/callback")
+async def github_token_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/google.py

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

hyperpocket/server/auth/linear.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+linear_auth_router = APIRouter(
+    prefix="/linear"
+)
+
+
+@linear_auth_router.get("/token/callback")
+async def slack_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/slack.py

@@ -0,0 +1,28 @@
+from fastapi import APIRouter
+from starlette.responses import HTMLResponse
+
+from hyperpocket.futures import FutureStore
+
+slack_auth_router = APIRouter(
+    prefix="/slack"
+)
+
+
+@slack_auth_router.get("/oauth2/callback")
+async def slack_oauth2_callback(state: str, code: str):
+    try:
+        FutureStore.resolve_future(state, code)
+    except ValueError:
+        return HTMLResponse(content="failed")
+
+    return HTMLResponse(content="success")
+
+
+@slack_auth_router.get("/token/callback")
+async def slack_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/token.py

@@ -0,0 +1,51 @@
+from http import HTTPStatus
+from urllib.parse import urlencode, urlunparse, urlparse, parse_qs
+
+from fastapi import APIRouter, Form
+from starlette.responses import HTMLResponse, RedirectResponse
+
+token_router = APIRouter()
+
+
+@token_router.get("/token", response_class=HTMLResponse)
+async def token_form(redirect_uri: str, state: str = ""):
+    html = f"""    
+    <html>
+        <body>
+            <h2>Enter Token</h2>
+                     <form action="submit" method="post">
+                <input type="hidden" name="redirect_uri" value="{redirect_uri}">
+                <input type="hidden" name="state" value="{state}">
+                
+                <label for="user_token">Token:</label>
+                <input type="text" id="user_token" name="user_token" required>
+                
+                <button type="submit">submit</button>
+            </form>
+        </body>
+    </html>
+    """
+    return HTMLResponse(content=html)
+
+
+@token_router.post("/submit", response_class=RedirectResponse)
+async def submit_token(user_token: str = Form(...), redirect_uri: str = Form(...), state: str = Form(...)):
+    new_callback_url = add_query_params(redirect_uri, {"token": user_token, "state": state})
+    return RedirectResponse(url=new_callback_url, status_code=HTTPStatus.SEE_OTHER)
+
+
+def add_query_params(url: str, params: dict):
+    url_parts = urlparse(url)
+    query_params = parse_qs(url_parts.query)
+    query_params.update(params)
+    new_query = urlencode(query_params, doseq=True)
+
+    new_url = urlunparse((
+        url_parts.scheme,
+        url_parts.netloc,
+        url_parts.path,
+        url_parts.params,
+        new_query,
+        url_parts.fragment
+    ))
+    return new_url

hyperpocket/server/proxy.py

@@ -0,0 +1,63 @@
+import httpx
+from fastapi import FastAPI, Request
+from starlette.responses import HTMLResponse
+
+from hyperpocket.config import config, pocket_logger
+
+
+async def proxy(request: Request, path: str):
+    async with httpx.AsyncClient() as client:
+        resp = await client.request(
+            method=request.method,
+            url=f"{config.internal_base_url}/{path}",
+            headers=request.headers,
+            content=await request.body(),
+            params=request.query_params,
+            timeout=300,
+        )
+        return HTMLResponse(content=resp.text, headers=resp.headers, status_code=resp.status_code)
+
+
+def add_callback_proxy(app: FastAPI):
+    app.add_api_route(f"/{config.callback_url_rewrite_prefix}/{{path:path}}", proxy,
+                      methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
+
+
+https_proxy_app = None
+if config.enable_local_callback_proxy:
+    https_proxy_app = FastAPI()
+    add_callback_proxy(https_proxy_app)
+
+
+def _generate_ssl_certificates(ssl_keypath, ssl_certpath):
+    import subprocess
+
+    pocket_logger.info("generate default ssl file")
+
+    subj = (
+        "/C=US"
+        "/ST=California"
+        "/L=San Jose"
+        "/O=local"
+        "/OU=local"
+        "/CN=localhost"
+        "/emailAddress=local@example.com"
+    )
+    command = [
+        "openssl", "req", "-x509",
+        "-newkey", "rsa:4096",
+        "-keyout", ssl_keypath,
+        "-out", ssl_certpath,
+        "-days", "1",
+        "-nodes",
+        '-subj', subj,
+        "-sha256",
+    ]
+
+    try:
+        # 명령 실행
+        subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        pocket_logger.info("SSL Certificates generated: callback_server.key, callback_server.crt")
+    except subprocess.CalledProcessError as e:
+        pocket_logger.warning(f"An error occurred while generating certificates: {e}")
+        raise e

hyperpocket/server/server.py

@@ -0,0 +1,178 @@
+import asyncio
+import enum
+import uuid
+from typing import Optional
+
+import multiprocess as mp
+from fastapi import FastAPI
+from uvicorn import Config, Server
+
+from hyperpocket.config import config, pocket_logger
+from hyperpocket.server.auth import auth_router
+from hyperpocket.server.tool import tool_router
+
+
+class PocketServerOperations(enum.Enum):
+    CALL = "call"
+    PREPARE_AUTH = "prepare_auth"
+    AUTHENTICATE = "authenticate"
+    TOOL_CALL = "tool_call"
+
+
+class PocketServer(object):
+    main_server: Server
+    internal_server_port: int
+    proxy_server: Optional[Server]
+    proxy_port: int
+    pipe: mp.Pipe
+    process: mp.Process
+    future_store: dict[str, asyncio.Future]
+
+    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
+        self.future_store = dict()
+
+    def teardown(self):
+        if self.process and self.process.is_alive():
+            self.process.terminate()
+            self.process.join()
+
+    async def _run_async(self):
+        try:
+            await asyncio.gather(
+                self.main_server.serve(),
+                self.proxy_server.serve() if self.proxy_server is not None else asyncio.sleep(0),
+                self.poll_in_child(),
+            )
+        except Exception as e:
+            pocket_logger.warning(f"failed to start pocket server. error : {e}")
+
+    async def poll_in_child(self):
+        loop = asyncio.get_running_loop()
+        _, conn = self.pipe
+
+        async def _acall(_conn, _op, _uid, a, kw):
+            result = await self.child_pocket.acall(*a, **kw)
+            _conn.send((_op, _uid, result))
+
+        async def _prepare(_conn, _op, _uid, a, kw):
+            result = self.child_pocket.prepare_auth(*a, **kw)
+            _conn.send((_op, _uid, result))
+
+        async def _authenticate(_conn, _op, _uid, a, kw):
+            result = await self.child_pocket.authenticate(*a, **kw)
+            _conn.send((_op, _uid, result))
+
+        async def _tool_call(_conn, _op, _uid, a, kw):
+            result = await self.child_pocket.tool_call(*a, **kw)
+            _conn.send((_op, _uid, result))
+
+        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))
+                else:
+                    raise AttributeError(f"Can't find operations. op:{op}")
+            else:
+                await asyncio.sleep(0)
+
+    def send_in_parent(self,
+                       op: PocketServerOperations,
+                       args: tuple,
+                       kwargs: dict):
+        conn, _ = self.pipe
+        uid = str(uuid.uuid4())
+        message = (op.value, uid, args, kwargs)
+        future = asyncio.Future()
+        self.future_store[uid] = future
+        conn.send(message)
+        return uid
+
+    async def poll_in_parent(self):
+        conn, _ = self.pipe
+        while True:
+            if conn.poll():
+                op, uid, result = conn.recv()
+                future = self.future_store[uid]
+                future.set_result(result)
+                break
+            else:
+                await asyncio.sleep(0)
+
+    async def call_in_subprocess(self,
+                                 op: PocketServerOperations,
+                                 args: tuple,
+                                 kwargs: dict):
+        uid = self.send_in_parent(op, args, kwargs)
+        loop = asyncio.get_running_loop()
+        loop.create_task(self.poll_in_parent())
+        return await self.future_store[uid]
+
+    def run(self, child_pocket):
+        self._set_mp_start_method()
+
+        self.pipe = mp.Pipe()
+        self.process = mp.Process(target=self._run, args=(child_pocket,), daemon=True)
+        self.process.start()
+
+    def _run(self, child_pocket):
+        self.child_pocket = child_pocket
+        self.main_server = self._create_main_server()
+        self.proxy_server = self._create_https_proxy_server()
+        loop = asyncio.new_event_loop()
+        loop.run_until_complete(self._run_async())
+        loop.close()
+
+    def _create_main_server(self) -> Server:
+        app = FastAPI()
+        _config = Config(app, host="0.0.0.0", port=self.internal_server_port)
+        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
+
+    def _create_https_proxy_server(self) -> Optional[Server]:
+        if not config.enable_local_callback_proxy:
+            return None
+        from hyperpocket.server.proxy import _generate_ssl_certificates
+        from hyperpocket.server.proxy import https_proxy_app
+
+        from hyperpocket.config.settings import pocket_root
+        ssl_keypath = pocket_root / "callback_server.key"
+        ssl_certpath = pocket_root / "callback_server.crt"
+
+        if not ssl_keypath.exists() or not ssl_certpath.exists():
+            _generate_ssl_certificates(ssl_keypath, ssl_certpath)
+
+        _config = Config(https_proxy_app, host="0.0.0.0", port=self.proxy_port, ssl_keyfile=ssl_keypath,
+                         ssl_certfile=ssl_certpath)
+        proxy_server = Server(_config)
+        return proxy_server
+
+    def _set_mp_start_method(self):
+        import platform
+
+        os_name = platform.system()
+        if os_name == "Windows":
+            mp.set_start_method("spawn", force=True)
+            pocket_logger.debug("Process start method set to 'spawn' for Windows.")
+        elif os_name == "Darwin":  # macOS
+            mp.set_start_method("spawn", force=True)
+            pocket_logger.debug("Process start method set to 'spawn' for macOS.")
+        elif os_name == "Linux":
+            mp.set_start_method("fork", force=True)
+            pocket_logger.debug("Process start method set to 'fork' for Linux.")
+        else:
+            pocket_logger.debug(f"Unrecognized OS: {os_name}. Default start method will be used.")