mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
fe0ff240e7
Rich is a more mature library for displaying progress bars, logging, and console output. This should help properly align progress bars within the terminal. Side note: "We're Rich!" Signed-off-by: kingbri <bdashore3@proton.me>
98 lines
2.2 KiB
Python
98 lines
2.2 KiB
Python
"""Common utility functions"""
|
|
|
|
import traceback
|
|
from pydantic import BaseModel
|
|
from rich.progress import (
|
|
Progress,
|
|
TextColumn,
|
|
BarColumn,
|
|
TimeRemainingColumn,
|
|
TaskProgressColumn,
|
|
MofNCompleteColumn,
|
|
)
|
|
from typing import Optional
|
|
|
|
from common.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def load_progress(module, modules):
|
|
"""Wrapper callback for load progress."""
|
|
yield module, modules
|
|
|
|
|
|
class TabbyRequestErrorMessage(BaseModel):
|
|
"""Common request error type."""
|
|
|
|
message: str
|
|
trace: Optional[str] = None
|
|
|
|
|
|
class TabbyRequestError(BaseModel):
|
|
"""Common request error type."""
|
|
|
|
error: TabbyRequestErrorMessage
|
|
|
|
|
|
def get_generator_error(message: str, exc_info: bool = True):
|
|
"""Get a generator error."""
|
|
|
|
generator_error = handle_request_error(message)
|
|
|
|
return get_sse_packet(generator_error.model_dump_json())
|
|
|
|
|
|
def handle_request_error(message: str, exc_info: bool = True):
|
|
"""Log a request error to the console."""
|
|
|
|
error_message = TabbyRequestErrorMessage(
|
|
message=message, trace=traceback.format_exc()
|
|
)
|
|
|
|
request_error = TabbyRequestError(error=error_message)
|
|
|
|
# Log the error and provided message to the console
|
|
if error_message.trace and exc_info:
|
|
logger.error(error_message.trace)
|
|
|
|
logger.error(f"Sent to request: {message}")
|
|
|
|
return request_error
|
|
|
|
|
|
def get_sse_packet(json_data: str):
|
|
"""Get an SSE packet."""
|
|
return f"data: {json_data}\n\n"
|
|
|
|
|
|
def get_loading_progress_bar():
|
|
"""Gets a pre-made progress bar for loading tasks."""
|
|
|
|
return Progress(
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
MofNCompleteColumn(),
|
|
TimeRemainingColumn(),
|
|
)
|
|
|
|
|
|
def unwrap(wrapped, default=None):
|
|
"""Unwrap function for Optionals."""
|
|
if wrapped is None:
|
|
return default
|
|
|
|
return wrapped
|
|
|
|
|
|
def coalesce(*args):
|
|
"""Coalesce function for multiple unwraps."""
|
|
return next((arg for arg in args if arg is not None), None)
|
|
|
|
|
|
def prune_dict(input_dict):
|
|
"""Trim out instances of None from a dictionary"""
|
|
|
|
return {k: v for k, v in input_dict.items() if v is not None}
|