src/agent_server/chat/profile_selection.py - master - man4j/agent-server

src/agent_server/chat/profile_selection.py
218 строк7 KB
Vladimir
initial
18 апр 2026, 19:36
18 апр 2026, 19:36a95f341
from chainlit.types import ThreadDict

import chainlit as cl

from agent_server.chat.chainlit_context import (
    get_auth_users_map_cached,
    get_current_auth_user,
    get_current_chainlit_user,
    has_chainlit_context,
)
from agent_server.chat.session import (
    get_agent_id,
    get_chat_profile_id,
    get_llm_profile_id,
    get_profiles,
    get_profiles_map,
    set_agent_id,
    set_chat_profile_id,
    set_llm_model,
    set_llm_profile_id,
    set_profiles,
    set_profiles_map,
)
from agent_server.chat.thread_metadata import get_thread_metadata_profile_id
from agent_server.config.agents import AgentDefinition, load_agent_definition
from agent_server.config.profiles import LlmProfile, load_chat_profiles


_PROFILES_CACHE: list[LlmProfile] | None = None


def get_profiles_cached() -> list[LlmProfile]:
    global _PROFILES_CACHE
    if _PROFILES_CACHE is None:
        _PROFILES_CACHE = load_chat_profiles()
    return _PROFILES_CACHE


def normalize_profiles(raw_profiles) -> list[LlmProfile]:
    if not isinstance(raw_profiles, list):
        return []

    result: list[LlmProfile] = []
    for item in raw_profiles:
        if isinstance(item, LlmProfile) and getattr(item, "id", None):
            result.append(item)

    return result


def build_profiles_map(profiles: list[LlmProfile]) -> dict[str, LlmProfile]:
    result: dict[str, LlmProfile] = {}
    for profile in profiles:
        if isinstance(profile, LlmProfile) and profile.id:
            result[profile.id] = profile
    return result


def get_first_valid_profile(profiles: list[LlmProfile]) -> LlmProfile:
    for profile in profiles:
        if isinstance(profile, LlmProfile) and profile.id:
            return profile
    raise ValueError("Не найден ни один валидный LLM profile")


def ensure_profiles_loaded() -> list[LlmProfile]:
    if not has_chainlit_context():
        profiles = normalize_profiles(get_profiles_cached())
        if not profiles:
            raise ValueError("Список chat profiles пуст: load_chat_profiles() не вернул валидных профилей")
        return profiles

    session_profiles = normalize_profiles(get_profiles())
    session_profiles_map = get_profiles_map()

    if session_profiles and isinstance(session_profiles_map, dict):
        valid_map = {
            profile_id: profile
            for profile_id, profile in session_profiles_map.items()
            if isinstance(profile_id, str)
            and isinstance(profile, LlmProfile)
            and profile.id == profile_id
        }
        if valid_map:
            set_profiles(session_profiles)
            set_profiles_map(valid_map)
            return session_profiles

    profiles = normalize_profiles(get_profiles_cached())
    if not profiles:
        raise ValueError("Список chat profiles пуст: load_chat_profiles() не вернул валидных профилей")

    profiles_map = build_profiles_map(profiles)
    set_profiles(profiles)
    set_profiles_map(profiles_map)
    return profiles


def filter_profiles_for_user(
    profiles: list[LlmProfile],
    current_user: cl.User | None,
) -> list[LlmProfile]:
    if current_user is None:
        return []

    username = str(getattr(current_user, "identifier", "") or "").strip()
    if not username:
        return []

    auth_user = get_auth_users_map_cached().get(username)
    if auth_user is None:
        return []

    allowed_agents = set(auth_user.allowed_agents or ())
    return [profile for profile in profiles if profile.agent_id in allowed_agents]


def is_agent_allowed_for_current_user(agent_id: str) -> bool:
    auth_user = get_current_auth_user()
    if auth_user is None:
        return False
    return agent_id in set(auth_user.allowed_agents or ())


def _build_allowed_profiles_map() -> tuple[list[LlmProfile], dict[str, LlmProfile]]:
    profiles = ensure_profiles_loaded()
    profiles = filter_profiles_for_user(profiles, get_current_chainlit_user())
    if not profiles:
        raise ValueError("Для текущего пользователя нет доступных LLM profiles")

    profiles_map_all: dict[str, LlmProfile] = get_profiles_map()
    allowed_profile_ids = {profile.id for profile in profiles}
    profiles_map = {
        profile_id: profile
        for profile_id, profile in profiles_map_all.items()
        if profile_id in allowed_profile_ids
    }
    return profiles, profiles_map


def get_selected_profile() -> LlmProfile:
    profiles, profiles_map = _build_allowed_profiles_map()

    selected = get_chat_profile_id()
    if isinstance(selected, str) and selected:
        selected_profile = profiles_map.get(selected)
        if isinstance(selected_profile, LlmProfile) and selected_profile.id:
            return selected_profile

    llm_profile_id = get_llm_profile_id()
    if isinstance(llm_profile_id, str) and llm_profile_id:
        selected_profile = profiles_map.get(llm_profile_id)
        if isinstance(selected_profile, LlmProfile) and selected_profile.id:
            return selected_profile

    return get_first_valid_profile(profiles)


def get_selected_agent() -> AgentDefinition:
    profile = get_selected_profile()

    if not is_agent_allowed_for_current_user(profile.agent_id):
        raise ValueError(f"Agent is not allowed for current user: {profile.agent_id}")

    return load_agent_definition(profile.agent_id)


def set_selected_profile(profile: LlmProfile) -> LlmProfile:
    if profile is None or not isinstance(profile, LlmProfile) or not profile.id:
        raise ValueError("Попытка сохранить невалидный LLM profile")

    if not is_agent_allowed_for_current_user(profile.agent_id):
        raise ValueError(f"Профиль недоступен текущему пользователю: {profile.id}")

    set_chat_profile_id(profile.id)
    set_llm_profile_id(profile.id)
    set_llm_model(profile.model)
    agent = load_agent_definition(profile.agent_id)
    set_agent_id(agent.id)
    return profile


def ensure_profile_selected() -> LlmProfile:
    return set_selected_profile(get_selected_profile())


def restore_profile_for_thread(thread: ThreadDict | None) -> LlmProfile:
    profiles, profiles_map = _build_allowed_profiles_map()

    candidates = [
        get_thread_metadata_profile_id(thread),
        get_chat_profile_id(),
        get_llm_profile_id(),
    ]

    for candidate in candidates:
        if not isinstance(candidate, str) or not candidate:
            continue
        selected_profile = profiles_map.get(candidate)
        if isinstance(selected_profile, LlmProfile) and selected_profile.id:
            return set_selected_profile(selected_profile)

    return set_selected_profile(get_first_valid_profile(profiles))


def get_current_bot_name() -> str:
    agent_id = get_agent_id()
    if isinstance(agent_id, str) and agent_id.strip():
        return agent_id

    try:
        agent = get_selected_agent()
        if isinstance(agent.id, str) and agent.id.strip():
            return agent.id
    except Exception:
        pass

    return "bot"