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

src/agent_server/chat/loop.py
119 строк4 KB
Vladimir
fixes
28 апр 2026, 19:57
28 апр 2026, 19:57f04f1da
from dataclasses import dataclass

import chainlit as cl
from openai import OpenAI

from agent_server.chat.policy import ChatBehaviorPolicy
from agent_server.chat.presenter import (
    apply_tool_results_to_answer,
    create_answer_message,
    finalize_answer_message,
)
from agent_server.chat.types import ChatMessage
from agent_server.chat.session import get_tool_session_state
from agent_server.chat.profile_selection import get_current_bot_name
from agent_server.chat.summary import prepare_messages_for_next_llm_round
from agent_server.chat.tools import dispatch_tool_call, get_selected_agent_local_tools
from agent_server.llm.agent import handle_tool_calls, one_llm_round
from agent_server.message_window import estimated_prompt_tokens
from agent_server.ui_meta import UsageTotals


@dataclass
class ChatLoopResult:
    usage_totals: UsageTotals
    calibrated_char_per_token: float


async def run_chat_loop(
    *,
    client: OpenAI,
    selected_profile,
    selected_agent,
    policy: ChatBehaviorPolicy,
    messages: list[ChatMessage],
    char_per_token: float,
    response_headroom_tokens: int,
) -> ChatLoopResult:
    async with cl.Step(name="Thinking", type="llm") as thinking_step:
        answer_msg = await create_answer_message(get_current_bot_name())

        usage_totals = UsageTotals()
        calibration_actual_prompt_tokens = 0
        calibration_estimated_prompt_tokens = 0
        tools = get_tool_session_state().tools
        if not isinstance(tools, list):
            tools = get_selected_agent_local_tools()

        while True:
            estimated_before_round = estimated_prompt_tokens(messages, char_per_token)

            round_result = await one_llm_round(
                client=client,
                profile=selected_profile,
                messages=messages,
                tools=tools,
                thinking_step=thinking_step,
                answer_msg=answer_msg,
            )

            usage_totals.add(round_result.usage)

            if round_result.usage and getattr(round_result.usage, "prompt_tokens", None):
                calibration_actual_prompt_tokens += int(
                    getattr(round_result.usage, "prompt_tokens", 0) or 0
                )
                calibration_estimated_prompt_tokens += estimated_before_round

            assistant_msg = {
                "role": "assistant",
                "content": round_result.assistant_content,
                "reasoning_content": round_result.reasoning_content,
            }
            if round_result.tool_calls:
                assistant_msg["tool_calls"] = round_result.tool_calls
            messages.append(assistant_msg)

            if not round_result.tool_calls:
                break

            tool_messages, tool_results = await handle_tool_calls(
                tool_calls=round_result.tool_calls,
                call_tool=dispatch_tool_call,
            )
            messages.extend(tool_messages)

            updated = await apply_tool_results_to_answer(answer_msg, tool_results)
            if updated:
                await answer_msg.update()

            await prepare_messages_for_next_llm_round(
                messages,
                client=client,
                profile=selected_profile,
                agent=selected_agent,
                policy=policy,
                char_per_token=char_per_token,
            )

        mcp_state = get_tool_session_state().mcp_state or {}
        await finalize_answer_message(
            answer_msg,
            usage_totals=usage_totals,
            messages=messages,
            char_per_token=char_per_token,
            connected_servers=list((mcp_state.get("servers") or {}).keys()),
            context_size=policy.context_size,
            response_headroom_tokens=response_headroom_tokens,
            estimated_prompt_tokens_fn=estimated_prompt_tokens,
        )

    calibrated_char_per_token = char_per_token
    if calibration_actual_prompt_tokens > 0 and calibration_estimated_prompt_tokens > 0:
        ratio = calibration_actual_prompt_tokens / calibration_estimated_prompt_tokens
        calibrated_char_per_token = max(2.0, min(8.0, char_per_token / ratio))

    return ChatLoopResult(
        usage_totals=usage_totals,
        calibrated_char_per_token=calibrated_char_per_token,
    )