From 3ce93c78021736b830cc9de9c18a432bc85b3427 Mon Sep 17 00:00:00 2001 From: Wei Ouyang Date: Fri, 31 Jan 2025 23:49:07 +0800 Subject: [PATCH] Add hypha-cli --- hypha/VERSION | 2 +- hypha/__main__.py | 26 ++++++++++++++++++++---- hypha/core/store.py | 46 +++++++++++++++++++++++++++++++++--------- hypha/interactive.py | 48 +++++++++++++++++++++++++++++++------------- hypha/server.py | 16 ++++++++++++++- setup.py | 7 ++++++- 6 files changed, 114 insertions(+), 31 deletions(-) diff --git a/hypha/VERSION b/hypha/VERSION index b919f4f2..7918342a 100644 --- a/hypha/VERSION +++ b/hypha/VERSION @@ -1,3 +1,3 @@ { - "version": "0.20.47.post5" + "version": "0.20.47.post6" } diff --git a/hypha/__main__.py b/hypha/__main__.py index c205a85d..eecdfc17 100644 --- a/hypha/__main__.py +++ b/hypha/__main__.py @@ -3,20 +3,38 @@ import asyncio from hypha.server import get_argparser, create_application from hypha.interactive import start_interactive_shell +import uvicorn +import os +import sys + + +def run_interactive_cli(): + """Run the interactive CLI.""" + + # Create the app instance + arg_parser = get_argparser(add_help=True) + opt = arg_parser.parse_args(sys.argv[1:]) + + # Force interactive server mode + if not opt.interactive_server: + opt.interactive = True + + app = create_application(opt) + + # Start the interactive shell + asyncio.run(start_interactive_shell(app, opt)) if __name__ == "__main__": + # Create the app instance arg_parser = get_argparser() opt = arg_parser.parse_args() app = create_application(opt) - if opt.interactive: + if opt.interactive or opt.interactive_server: asyncio.run(start_interactive_shell(app, opt)) else: - import uvicorn - uvicorn.run(app, host=opt.host, port=int(opt.port)) - else: # Create the app instance when imported by uvicorn arg_parser = get_argparser(add_help=False) diff --git a/hypha/core/store.py b/hypha/core/store.py index 583146e9..51d264b3 100644 --- a/hypha/core/store.py +++ b/hypha/core/store.py @@ -980,6 +980,7 @@ class WorkspaceManagerWrapper: def __init__(self, store): self._store = store self._interface = None + self._interface_cm = None def __call__(self, user_info=None, workspace=None): """Support both wm() and wm(context) syntax for context manager.""" @@ -989,21 +990,46 @@ def __call__(self, user_info=None, workspace=None): workspace or "public", ) + async def _ensure_interface(self): + """Ensure we have an active interface, creating one if needed.""" + if self._interface is None: + self._interface_cm = self._store.get_workspace_interface( + self._store._root_user, + "public", + client_id=None, + silent=True, + ) + self._interface = await self._interface_cm.__aenter__() + return self._interface + + async def _cleanup_interface(self): + """Clean up the interface if it exists.""" + if self._interface_cm is not None: + await self._interface_cm.__aexit__(None, None, None) + self._interface = None + self._interface_cm = None + def __getattr__(self, name): """Support direct method access by creating a temporary interface.""" async def wrapper(*args, **kwargs): - if self._interface is None: - self._interface = self._store.get_workspace_interface( - self._store._root_user, - "public", - client_id=None, - silent=True, - ) - interface = await self._interface.__aenter__() - method = getattr(interface, name) - return await method(*args, **kwargs) + try: + interface = await self._ensure_interface() + method = getattr(interface, name) + result = await method(*args, **kwargs) + return result + except Exception as e: + await self._cleanup_interface() + raise e return wrapper + async def __aenter__(self): + """Support using the wrapper itself as a context manager.""" + return await self._ensure_interface() + + async def __aexit__(self, exc_type, exc, tb): + """Clean up when used as a context manager.""" + await self._cleanup_interface() + return WorkspaceManagerWrapper(self) diff --git a/hypha/interactive.py b/hypha/interactive.py index 7fb685d6..80dca5bd 100644 --- a/hypha/interactive.py +++ b/hypha/interactive.py @@ -10,7 +10,7 @@ from fastapi import FastAPI import uvicorn -from hypha.server import get_argparser +from hypha.server import get_argparser, create_application def configure_ptpython(repl: PythonRepl) -> None: @@ -35,7 +35,7 @@ def configure_ptpython(repl: PythonRepl) -> None: You can run async code directly using top-level await. -Type 'exit' to quit. +Type 'exit()' or press Ctrl+D to quit. """ print(welcome_message) @@ -46,21 +46,37 @@ async def start_interactive_shell(app: FastAPI, args: Any) -> None: # Get the store from the app state store = app.state.store - # Start the server in the background - config = uvicorn.Config(app, host=args.host, port=int(args.port)) - server = uvicorn.Server(config) - # Run the server in a separate task - server_task = asyncio.create_task(server.serve()) - await store.get_event_bus().wait_for_local("startup") - print(f"\nServer started at http://{args.host}:{args.port}") + server = None + server_task = None + print("Initializing interactive shell...\n") + async def start_server(): + nonlocal server, server_task + config = uvicorn.Config(app, host=args.host, port=int(args.port)) + server = uvicorn.Server(config) + # Run the server in a separate task + server_task = asyncio.create_task(server.serve()) + await store.get_event_bus().wait_for_local("startup") + print(f"\nServer started at http://{args.host}:{args.port}") + + if args.interactive_server: + # Start the server in the background + await start_server() + # Prepare the local namespace local_ns = { "app": app, "store": store, - "server": server, } + if server: + local_ns["server"] = server + else: + args.startup_functions = args.startup_functions or [] + await store.init( + reset_redis=args.reset_redis, startup_functions=args.startup_functions + ) + local_ns["start_server"] = start_server # Start the interactive shell with stdout patching for better async support with patch_stdout(): @@ -73,9 +89,10 @@ async def start_interactive_shell(app: FastAPI, args: Any) -> None: ) # Cleanup when shell exits - print("\nShutting down server...") - server.should_exit = True - await server_task + if server: + print("\nShutting down server...") + server.should_exit = True + await server_task if store.is_ready(): print("Cleaning up store...") @@ -88,8 +105,11 @@ def main() -> None: arg_parser = get_argparser() args = arg_parser.parse_args() + # Create the FastAPI app instance + app = create_application(args) + try: - asyncio.run(start_interactive_shell(args)) + asyncio.run(start_interactive_shell(app, args)) except (KeyboardInterrupt, EOFError): print("\nExiting Hypha Interactive Shell...") except Exception as e: diff --git a/hypha/server.py b/hypha/server.py index 3ea4f220..468f2f60 100644 --- a/hypha/server.py +++ b/hypha/server.py @@ -509,10 +509,24 @@ def get_argparser(add_help=True): return parser +def add_interactive_arguments(parser): + """Add interactive-specific arguments to the parser.""" + parser.add_argument( + "--interactive-server", + action="store_true", + help="start an interactive server with the hypha store", + ) + return parser + + if __name__ == "__main__": import uvicorn arg_parser = get_argparser() + + # Only add interactive server argument when running as main script + add_interactive_arguments(arg_parser) + opt = arg_parser.parse_args() # Apply database migrations @@ -527,7 +541,7 @@ def get_argparser(add_help=True): command.upgrade(alembic_cfg, "head") app = create_application(opt) - if opt.interactive: + if opt.interactive or opt.interactive_server: from hypha.interactive import start_interactive_shell asyncio.run(start_interactive_shell(app, opt)) diff --git a/setup.py b/setup.py index bdd68e24..75f68e5d 100644 --- a/setup.py +++ b/setup.py @@ -92,5 +92,10 @@ ], }, zip_safe=False, - entry_points={"console_scripts": ["hypha = hypha.__main__:main"]}, + entry_points={ + "console_scripts": [ + "hypha = hypha.__main__:main", + "hypha-cli = hypha.__main__:run_interactive_cli", + ] + }, )