autogen/python/packages/autogen-studio/autogenstudio/cli.py

125 lines
3.5 KiB
Python

import os
import tempfile
import warnings
from typing import Optional
import typer
import uvicorn
from typing_extensions import Annotated
from .version import VERSION
app = typer.Typer()
# Ignore deprecation warnings from websockets
warnings.filterwarnings("ignore", message="websockets.legacy is deprecated*")
warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated*")
@app.command()
def ui(
host: str = "127.0.0.1",
port: int = 8081,
workers: int = 1,
reload: Annotated[bool, typer.Option("--reload")] = False,
docs: bool = True,
appdir: str = None,
database_uri: Optional[str] = None,
upgrade_database: bool = False,
):
"""
Run the AutoGen Studio UI.
Args:
host (str, optional): Host to run the UI on. Defaults to 127.0.0.1 (localhost).
port (int, optional): Port to run the UI on. Defaults to 8081.
workers (int, optional): Number of workers to run the UI with. Defaults to 1.
reload (bool, optional): Whether to reload the UI on code changes. Defaults to False.
docs (bool, optional): Whether to generate API docs. Defaults to False.
appdir (str, optional): Path to the AutoGen Studio app directory. Defaults to None.
database-uri (str, optional): Database URI to connect to. Defaults to None.
"""
# Create temporary env file to share configuration with uvicorn workers
temp_env = tempfile.NamedTemporaryFile(mode="w", suffix=".env", delete=True)
# Write configuration
env_vars = {
"AUTOGENSTUDIO_HOST": host,
"AUTOGENSTUDIO_PORT": port,
"AUTOGENSTUDIO_API_DOCS": str(docs),
}
if appdir:
env_vars["AUTOGENSTUDIO_APPDIR"] = appdir
if database_uri:
env_vars["AUTOGENSTUDIO_DATABASE_URI"] = database_uri
if upgrade_database:
env_vars["AUTOGENSTUDIO_UPGRADE_DATABASE"] = "1"
for key, value in env_vars.items():
temp_env.write(f"{key}={value}\n")
temp_env.flush()
uvicorn.run(
"autogenstudio.web.app:app",
host=host,
port=port,
workers=workers,
reload=reload,
reload_excludes=["**/alembic/*", "**/alembic.ini", "**/versions/*"] if reload else None,
env_file=temp_env.name,
)
@app.command()
def serve(
team: str = "",
host: str = "127.0.0.1",
port: int = 8084,
workers: int = 1,
docs: bool = False,
):
"""
Serve an API Endpoint based on an AutoGen Studio workflow json file.
Args:
team (str): Path to the team json file.
host (str, optional): Host to run the UI on. Defaults to 127.0.0.1 (localhost).
port (int, optional): Port to run the UI on. Defaults to 8084
workers (int, optional): Number of workers to run the UI with. Defaults to 1.
reload (bool, optional): Whether to reload the UI on code changes. Defaults to False.
docs (bool, optional): Whether to generate API docs. Defaults to False.
"""
os.environ["AUTOGENSTUDIO_API_DOCS"] = str(docs)
os.environ["AUTOGENSTUDIO_TEAM_FILE"] = team
# validate the team file
if not os.path.exists(team):
raise ValueError(f"Team file not found: {team}")
uvicorn.run(
"autogenstudio.web.serve:app",
host=host,
port=port,
workers=workers,
reload=False,
)
@app.command()
def version():
"""
Print the version of the AutoGen Studio UI CLI.
"""
typer.echo(f"AutoGen Studio CLI version: {VERSION}")
def run():
app()
if __name__ == "__main__":
app()