from pydantic import BaseModel

from llm_server.config.global_config import GlobalConfig


def cluster_worker_count():
    count = 0
    for item in GlobalConfig.get().cluster:
        count += item.concurrent_gens
    return count


class ModeUINameStr(BaseModel):
    name: str
    api_name: str
    streaming_name: str


MODE_UI_NAMES = {
    'ooba': ModeUINameStr(name='Text Gen WebUI (ooba)', api_name='Blocking API url', streaming_name='Streaming API url'),
    'vllm': ModeUINameStr(name='Text Gen WebUI (ooba)', api_name='Blocking API url', streaming_name='Streaming API url'),
}