src/agent_server/ui_meta.py - master - man4j/agent-server

src/agent_server/ui_meta.py
71 строка2 KB
Vladimir
fixes
28 апр 2026, 19:57
28 апр 2026, 19:57f04f1da
from dataclasses import dataclass
from typing import Callable

from agent_server.chat.types import ChatMessage

META_SEPARATOR = "\n&nbsp;\n\n---\n&nbsp;\n"


@dataclass
class UsageTotals:
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

    def add(self, usage) -> None:
        if not usage:
            return
        self.prompt_tokens += int(getattr(usage, "prompt_tokens", 0) or 0)
        self.completion_tokens += int(getattr(usage, "completion_tokens", 0) or 0)
        self.total_tokens += int(getattr(usage, "total_tokens", 0) or 0)

    @property
    def has_any(self) -> bool:
        return any((self.prompt_tokens, self.completion_tokens, self.total_tokens))


def strip_meta_block(content: str | None) -> str:
    if not content:
        return ""

    pos = content.find(META_SEPARATOR)
    if pos == -1:
        return content

    return content[:pos].rstrip()


def apply_meta_block(content: str | None, meta_block: str) -> str:
    base = strip_meta_block(content)
    if not meta_block:
        return base
    return f"{base}{meta_block}" if base else meta_block.lstrip("\n")


def build_meta_block(
    *,
    usage_totals: UsageTotals | None,
    messages: list[ChatMessage],
    char_per_token: float,
    estimated_prompt_tokens_fn: Callable[[list[ChatMessage], float], int],
    connected_servers: list[str] | None = None,
    context_size: int,
    response_headroom_tokens: int,
) -> str:
    meta_lines = [META_SEPARATOR]

    prompt_limit = max(1, context_size - response_headroom_tokens)
    current_prompt_tokens = estimated_prompt_tokens_fn(messages, char_per_token)
    current_prompt_pct = min(100.0, (current_prompt_tokens / prompt_limit) * 100.0)
    context_message_count = sum(1 for msg in messages if msg.get("role") != "system")
    user_message_count = sum(1 for msg in messages if msg.get("role") == "user")

    meta_lines.append(f"###### Prompt tokens: `{current_prompt_tokens}` / `{prompt_limit}`")
    meta_lines.append(f"###### Prompt usage: `{current_prompt_pct:.1f}%`")
    meta_lines.append(f"###### User messages: `{user_message_count}`")
    meta_lines.append(f"###### Context messages: `{context_message_count}`")

    if connected_servers:
        meta_lines.append(f"###### MCP: `{', '.join(sorted(connected_servers))}`")

    return "\n".join(meta_lines)