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

src/agent_server/chat/session.py
216 строк6 KB
Vladimir
fixes
10 июл 2026, 21:08
10 июл 2026, 21:086e6ce5e
import asyncio
from dataclasses import dataclass, field

import chainlit as cl

from agent_server.chat.types import ChatMessage, McpState, OpenAIToolSpec, ThreadMetadata


@dataclass
class RuntimeHistoryState:
    messages: list[ChatMessage] = field(default_factory=list)
    char_per_token: float | None = None


@dataclass
class ToolSessionState:
    agent_id: str | None = None
    profile_id: str | None = None
    mcp_state: McpState | None = None
    tools: list[OpenAIToolSpec] | None = None


def get_session_value(key: str, default=None):
    value = cl.user_session.get(key)
    return default if value is None else value


def set_session_value(key: str, value) -> None:
    cl.user_session.set(key, value)


def get_session_lock() -> asyncio.Lock:
    value = cl.user_session.get("session_lock")
    if isinstance(value, asyncio.Lock):
        return value

    # Chainlit session state is shared across chat lifecycle handlers.
    # Use one lock per session so concurrent handlers do not interleave writes.
    lock = asyncio.Lock()
    cl.user_session.set("session_lock", lock)
    return lock


def get_profiles():
    return get_session_value("profiles")


def set_profiles(profiles) -> None:
    set_session_value("profiles", profiles)


def get_profiles_map() -> dict:
    value = get_session_value("profiles_map", {})
    return value if isinstance(value, dict) else {}


def set_profiles_map(profiles_map: dict) -> None:
    set_session_value("profiles_map", profiles_map)


def get_chat_profile_id() -> str | None:
    value = get_session_value("chat_profile")
    return value if isinstance(value, str) and value else None


def _sync_chainlit_chat_profile(profile_id: str) -> None:
    try:
        context = getattr(cl, "context", None)
        session = getattr(context, "session", None)
    except Exception:
        return

    if session is None:
        return

    try:
        setattr(session, "chat_profile", profile_id)
        if hasattr(session, "config"):
            session.config = None
            session.config = session.get_config()
    except Exception:
        return


def set_chat_profile_id(profile_id: str) -> None:
    _sync_chainlit_chat_profile(profile_id)
    set_session_value("chat_profile", profile_id)


def get_llm_profile_id() -> str | None:
    value = get_session_value("llm_profile_id")
    return value if isinstance(value, str) and value else None


def set_llm_profile_id(profile_id: str) -> None:
    set_session_value("llm_profile_id", profile_id)


def set_llm_model(model: str) -> None:
    set_session_value("llm_model", model)


def get_agent_id() -> str | None:
    value = get_session_value("agent_id")
    return value if isinstance(value, str) and value else None


def set_agent_id(agent_id: str | None) -> None:
    set_session_value("agent_id", agent_id)


def get_mcp_state() -> McpState | None:
    return get_session_value("mcp_state")


def set_mcp_state(state: McpState | None) -> None:
    set_session_value("mcp_state", state)


def get_tool_profile_id() -> str | None:
    value = get_session_value("tool_profile_id")
    return value if isinstance(value, str) and value else None


def set_tool_profile_id(profile_id: str | None) -> None:
    set_session_value("tool_profile_id", profile_id)


def get_tools() -> list[OpenAIToolSpec] | None:
    return get_session_value("tools")


def set_tools(tools: list[OpenAIToolSpec] | None) -> None:
    set_session_value("tools", tools)


def get_messages() -> list[ChatMessage]:
    value = get_session_value("messages", [])
    return value if isinstance(value, list) else []


def set_messages(messages: list[ChatMessage]) -> None:
    set_session_value("messages", messages)


def get_system_prompt() -> str | None:
    value = get_session_value("system_prompt")
    return value if isinstance(value, str) and value.strip() else None


def set_system_prompt(system_prompt: str) -> None:
    set_session_value("system_prompt", system_prompt)


def get_thread_metadata_session() -> ThreadMetadata:
    value = get_session_value("thread_metadata", {})
    return value if isinstance(value, dict) else {}


def set_thread_metadata_session(metadata: ThreadMetadata) -> None:
    set_session_value("thread_metadata", metadata)


def set_thread_metadata_synced(is_synced: bool) -> None:
    set_session_value("thread_metadata_synced", is_synced)


def get_char_per_token(default: float | None = None) -> float | None:
    value = get_session_value("char_per_token", default)
    return value


def set_char_per_token(value: float) -> None:
    set_session_value("char_per_token", value)


def get_thread_system_prompts() -> dict[str, str]:
    value = get_session_value("thread_system_prompts", {})
    return value if isinstance(value, dict) else {}


def set_thread_system_prompts(cache: dict[str, str]) -> None:
    set_session_value("thread_system_prompts", cache)


def get_runtime_history_state(default_char_per_token: float | None = None) -> RuntimeHistoryState:
    return RuntimeHistoryState(
        messages=get_messages(),
        char_per_token=get_char_per_token(default_char_per_token),
    )


def apply_runtime_history_state(state: RuntimeHistoryState) -> None:
    set_messages(state.messages)
    if state.char_per_token is not None:
        set_char_per_token(state.char_per_token)


def get_tool_session_state() -> ToolSessionState:
    return ToolSessionState(
        agent_id=get_agent_id(),
        profile_id=get_tool_profile_id(),
        mcp_state=get_mcp_state(),
        tools=get_tools(),
    )


def apply_tool_session_state(state: ToolSessionState) -> None:
    set_agent_id(state.agent_id)
    set_tool_profile_id(state.profile_id)
    set_mcp_state(state.mcp_state)
    set_tools(state.tools)


def clear_tool_session_state() -> None:
    apply_tool_session_state(ToolSessionState())