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

src/agent_server/chat/runtime_state.py
124 строки4 KB
Vladimir
fixes
10 июл 2026, 21:08
10 июл 2026, 21:086e6ce5e
from dataclasses import dataclass

from chainlit.types import ThreadDict

from agent_server.chat.policy import ChatBehaviorPolicy
from agent_server.chat.types import ChatMessage, ThreadMetadata
from agent_server.chat.session import (
    set_agent_id,
    set_messages,
    set_system_prompt,
    set_thread_metadata_session,
    set_thread_metadata_synced,
)
from agent_server.chat.prompt_cache import cache_system_prompt_for_thread
from agent_server.chat.thread_metadata import (
    get_thread_id,
    get_thread_metadata,
    get_thread_metadata_chat_summary,
    get_thread_metadata_messages,
)
from agent_server.chat.summary import upsert_chat_summary
from agent_server.chat.threads import (
    attach_system_message,
    build_new_chat_system_prompt,
    resolve_resume_system_prompt,
    restore_messages_from_thread,
)
from agent_server.config.agents import AgentDefinition
from agent_server.config.profiles import LlmProfile


@dataclass
class ChatRuntimeState:
    system_prompt: str
    messages: list[ChatMessage]
    metadata: ThreadMetadata
    agent_id: str


def build_new_chat_runtime_state(
    *,
    selected_profile: LlmProfile,
    selected_agent: AgentDefinition,
    mcp_state: dict,
) -> ChatRuntimeState:
    system_prompt = build_new_chat_system_prompt(selected_agent, mcp_state)
    return ChatRuntimeState(
        system_prompt=system_prompt,
        messages=[{"role": "system", "content": system_prompt}],
        metadata={
            "system_prompt": system_prompt,
            "llm_profile_id": selected_profile.id,
            "chat_profile": selected_profile.id,
            "agent_id": selected_agent.id,
        },
        agent_id=selected_agent.id,
    )


def build_resumed_chat_runtime_state(
    *,
    thread: ThreadDict,
    selected_profile: LlmProfile,
    selected_agent: AgentDefinition,
    policy: ChatBehaviorPolicy,
) -> ChatRuntimeState:
    system_prompt = resolve_resume_system_prompt(
        thread,
        selected_agent,
        policy,
    )

    restored = get_thread_metadata_messages(thread)
    chat_summary = get_thread_metadata_chat_summary(thread)
    summary_enabled = policy.chat_summarization_enabled

    if restored:
        messages = restored
        if not messages or messages[0].get("role") != "system":
            messages = attach_system_message(
                messages,
                system_prompt,
                chat_summary,
                include_chat_summary=summary_enabled,
            )
        else:
            messages[0]["content"] = system_prompt
            if summary_enabled:
                upsert_chat_summary(messages, chat_summary or "")
            else:
                upsert_chat_summary(messages, "")
    else:
        messages = attach_system_message(
            restore_messages_from_thread(thread),
            system_prompt,
            chat_summary,
            include_chat_summary=summary_enabled,
        )

    metadata = dict(get_thread_metadata(thread))
    if system_prompt:
        metadata["system_prompt"] = system_prompt
    metadata["llm_profile_id"] = selected_profile.id
    metadata["chat_profile"] = selected_profile.id
    metadata["agent_id"] = selected_agent.id

    return ChatRuntimeState(
        system_prompt=system_prompt,
        messages=messages,
        metadata=metadata,
        agent_id=selected_agent.id,
    )


def apply_runtime_state(state: ChatRuntimeState) -> None:
    set_system_prompt(state.system_prompt)
    set_agent_id(state.agent_id)
    set_thread_metadata_session(state.metadata)
    set_thread_metadata_synced(False)
    set_messages(state.messages)


def cache_runtime_state_for_thread(thread: ThreadDict, state: ChatRuntimeState) -> None:
    cache_system_prompt_for_thread(get_thread_id(thread), state.system_prompt)