ai-station/.venv/lib/python3.12/site-packages/chainlit/cli/__init__.py

245 lines
6.2 KiB
Python

import asyncio
import logging
import os
import sys
import click
import nest_asyncio
import uvicorn
# Not sure if it is necessary to call nest_asyncio.apply() before the other imports
nest_asyncio.apply()
# ruff: noqa: E402
from chainlit.auth import ensure_jwt_secret
from chainlit.cache import init_lc_cache
from chainlit.config import (
BACKEND_ROOT,
DEFAULT_HOST,
DEFAULT_PORT,
DEFAULT_ROOT_PATH,
config,
init_config,
lint_translations,
load_module,
)
from chainlit.logger import logger
from chainlit.markdown import init_markdown
from chainlit.secret import random_secret
from chainlit.utils import check_file
logging.basicConfig(
level=logging.INFO,
stream=sys.stdout,
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
def assert_app():
if (
not config.code.on_chat_start
and not config.code.on_message
and not config.code.on_audio_chunk
):
raise Exception(
"You need to configure at least one of on_chat_start, on_message or on_audio_chunk callback"
)
# Create the main command group for Chainlit CLI
@click.group(context_settings={"auto_envvar_prefix": "CHAINLIT"})
@click.version_option(prog_name="Chainlit")
def cli():
return
# Define the function to run Chainlit with provided options
def run_chainlit(target: str):
host = os.environ.get("CHAINLIT_HOST", DEFAULT_HOST)
port = int(os.environ.get("CHAINLIT_PORT", DEFAULT_PORT))
root_path = os.environ.get("CHAINLIT_ROOT_PATH", DEFAULT_ROOT_PATH)
ssl_certfile = os.environ.get("CHAINLIT_SSL_CERT", None)
ssl_keyfile = os.environ.get("CHAINLIT_SSL_KEY", None)
ws_per_message_deflate_env = os.environ.get(
"UVICORN_WS_PER_MESSAGE_DEFLATE", "true"
)
ws_per_message_deflate = ws_per_message_deflate_env.lower() in [
"true",
"1",
"yes",
] # Convert to boolean
ws_protocol = os.environ.get("UVICORN_WS_PROTOCOL", "auto")
config.run.host = host
config.run.port = port
config.run.root_path = root_path
from chainlit.server import app
check_file(target)
# Load the module provided by the user
config.run.module_name = target
load_module(config.run.module_name)
ensure_jwt_secret()
assert_app()
# Create the chainlit.md file if it doesn't exist
init_markdown(config.root)
# Initialize the LangChain cache if installed and enabled
init_lc_cache()
log_level = "debug" if config.run.debug else "error"
# Start the server
async def start():
config = uvicorn.Config(
app,
host=host,
port=port,
ws=ws_protocol,
log_level=log_level,
ws_per_message_deflate=ws_per_message_deflate,
ssl_keyfile=ssl_keyfile,
ssl_certfile=ssl_certfile,
)
server = uvicorn.Server(config)
await server.serve()
# Run the asyncio event loop instead of uvloop to enable re entrance
asyncio.run(start())
# uvicorn.run(app, host=host, port=port, log_level=log_level)
# Define the "run" command for Chainlit CLI
@cli.command("run")
@click.argument("target", required=True, envvar="RUN_TARGET")
@click.option(
"-w",
"--watch",
default=False,
is_flag=True,
envvar="WATCH",
help="Reload the app when the module changes",
)
@click.option(
"-h",
"--headless",
default=False,
is_flag=True,
envvar="HEADLESS",
help="Will prevent to auto open the app in the browser",
)
@click.option(
"-d",
"--debug",
default=False,
is_flag=True,
envvar="DEBUG",
help="Set the log level to debug",
)
@click.option(
"-c",
"--ci",
default=False,
is_flag=True,
envvar="CI",
help="Flag to run in CI mode",
)
@click.option(
"--no-cache",
default=False,
is_flag=True,
envvar="NO_CACHE",
help="Useful to disable third parties cache, such as langchain.",
)
@click.option(
"--ssl-cert",
default=None,
envvar="CHAINLIT_SSL_CERT",
help="Specify the file path for the SSL certificate.",
)
@click.option(
"--ssl-key",
default=None,
envvar="CHAINLIT_SSL_KEY",
help="Specify the file path for the SSL key",
)
@click.option("--host", help="Specify a different host to run the server on")
@click.option("--port", help="Specify a different port to run the server on")
@click.option("--root-path", help="Specify a different root path to run the server on")
def chainlit_run(
target,
watch,
headless,
debug,
ci,
no_cache,
ssl_cert,
ssl_key,
host,
port,
root_path,
):
if host:
os.environ["CHAINLIT_HOST"] = host
if port:
os.environ["CHAINLIT_PORT"] = port
if bool(ssl_cert) != bool(ssl_key):
raise click.UsageError(
"Both --ssl-cert and --ssl-key must be provided together."
)
if ssl_cert:
os.environ["CHAINLIT_SSL_CERT"] = ssl_cert
os.environ["CHAINLIT_SSL_KEY"] = ssl_key
if root_path:
os.environ["CHAINLIT_ROOT_PATH"] = root_path
if ci:
logger.info("Running in CI mode")
no_cache = True
# This is required to have OpenAI LLM providers available for the CI run
os.environ["OPENAI_API_KEY"] = "sk-FAKE-OPENAI-API-KEY"
config.run.headless = headless
config.run.debug = debug
config.run.no_cache = no_cache
config.run.ci = ci
config.run.watch = watch
config.run.ssl_cert = ssl_cert
config.run.ssl_key = ssl_key
run_chainlit(target)
@cli.command("hello")
@click.argument("args", nargs=-1)
def chainlit_hello(args=None, **kwargs):
hello_path = os.path.join(BACKEND_ROOT, "sample", "hello.py")
run_chainlit(hello_path)
@cli.command("init")
@click.argument("args", nargs=-1)
def chainlit_init(args=None, **kwargs):
init_config(log=True)
@cli.command("create-secret")
@click.argument("args", nargs=-1)
def chainlit_create_secret(args=None, **kwargs):
print(
f'Copy the following secret into your .env file. Once it is set, changing it will logout all users with active sessions.\nCHAINLIT_AUTH_SECRET="{random_secret()}"'
)
@cli.command("lint-translations")
@click.argument("args", nargs=-1)
def chainlit_lint_translations(args=None, **kwargs):
lint_translations()