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

src/agent_server/chat/threads.py
131 строка4 KB
Vladimir
fixes
25 апр 2026, 18:59
25 апр 2026, 18:5953a2633
from datetime import datetime
from zoneinfo import ZoneInfo

from chainlit.types import ThreadDict

from agent_server.chat.policy import ChatBehaviorPolicy
from agent_server.chat.types import ChatMessage, McpState
from agent_server.chat.session import (
    get_mcp_state,
    get_messages,
    get_system_prompt,
)
from agent_server.chat.prompt_cache import get_cached_system_prompt_for_thread
from agent_server.chat.thread_metadata import get_thread_metadata_system_prompt
from agent_server.chat.summary import upsert_chat_summary
from agent_server.mcp_client.client import all_prompts_from_state
from agent_server.config.agents import AgentDefinition, load_system_prompt
from agent_server.ui_meta import strip_meta_block


def now_iso() -> str:
    return datetime.now(ZoneInfo("Europe/Moscow")).isoformat()


def with_timestamp(content: str) -> str:
    return f"[timestamp: {now_iso()}]\n{content}"


def build_new_chat_system_prompt(agent: AgentDefinition, mcp_state: McpState | None) -> str:
    return (
        load_system_prompt(agent.system_prompt_file, skills_dir=agent.skills_dir)
        + all_prompts_from_state(mcp_state or {})
    )


def serialize_messages_for_thread(messages: list[ChatMessage]) -> list[ChatMessage]:
    result: list[ChatMessage] = []

    for msg in messages:
        item: ChatMessage = {
            "role": msg.get("role"),
            "content": msg.get("content"),
        }

        if msg.get("reasoning_content") is not None:
            item["reasoning_content"] = msg.get("reasoning_content")

        if msg.get("tool_call_id") is not None:
            item["tool_call_id"] = msg.get("tool_call_id")

        if msg.get("tool_calls") is not None:
            item["tool_calls"] = msg.get("tool_calls")

        if msg.get("name") is not None:
            item["name"] = msg.get("name")

        result.append(item)

    return result


def restore_messages_from_thread(thread: ThreadDict) -> list[ChatMessage]:
    messages: list[ChatMessage] = []

    steps = thread.get("steps") or []
    for step in steps:
        step_type = step.get("type")
        output = step.get("output") or ""

        if not output:
            continue

        clean_output = strip_meta_block(output)

        if step_type == "user_message":
            messages.append({"role": "user", "content": clean_output})
        elif step_type == "assistant_message":
            messages.append({"role": "assistant", "content": clean_output})

    return messages


def attach_system_message(
    messages: list[ChatMessage] | None,
    system_prompt: str,
    chat_summary: str | None = None,
    *,
    include_chat_summary: bool = True,
) -> list[ChatMessage]:
    result = list(messages or [])

    if result and result[0]["role"] == "system":
        if include_chat_summary and chat_summary:
            upsert_chat_summary(result, chat_summary)
        elif not include_chat_summary:
            upsert_chat_summary(result, "")
        return result

    result = [{"role": "system", "content": system_prompt}, *result]
    if include_chat_summary and chat_summary:
        upsert_chat_summary(result, chat_summary)
    return result


def resolve_resume_system_prompt(
    thread: ThreadDict | None,
    agent: AgentDefinition,
    policy: ChatBehaviorPolicy,
) -> str:
    if policy.rebuild_system_prompt_on_chat_resume:
        return build_new_chat_system_prompt(agent, get_mcp_state())

    metadata_prompt = get_thread_metadata_system_prompt(thread)
    if metadata_prompt:
        return metadata_prompt

    current_messages = get_messages()
    if current_messages and current_messages[0]["role"] == "system":
        current_prompt = current_messages[0].get("content") or ""
        if current_prompt.strip():
            return current_prompt

    cached_prompt = get_cached_system_prompt_for_thread(thread)
    if cached_prompt:
        return cached_prompt

    stored_prompt = get_system_prompt()
    if isinstance(stored_prompt, str) and stored_prompt.strip():
        return stored_prompt

    return build_new_chat_system_prompt(agent, get_mcp_state())