| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | from typing import Any |
| 5 | |
| 6 | from agent import Agent |
| 7 | from helpers import persist_chat, plugins, tokens |
| 8 | from helpers.state_monitor_integration import mark_dirty_all |
| 9 | |
| 10 | |
| 11 | PLUGIN_NAME = "_chat_naming" |
| 12 | MODE_ONCE = "once" |
| 13 | MODE_ALWAYS = "always" |
| 14 | RECENT_USER_MESSAGES = 4 |
| 15 | GENERATED_NAME_LIMIT = 40 |
| 16 | UTILITY_CONTEXT_INPUT_RATIO = 0.7 |
| 17 | |
| 18 | |
| 19 | def get_config(agent: Agent) -> dict[str, Any]: |
| 20 | config = plugins.get_plugin_config(PLUGIN_NAME, agent=agent) or {} |
| 21 | mode = str(config.get("automatic_naming_mode", MODE_ONCE) or MODE_ONCE) |
| 22 | if mode not in {MODE_ONCE, MODE_ALWAYS}: |
| 23 | mode = MODE_ONCE |
| 24 | return { |
| 25 | "automatic_naming": bool(config.get("automatic_naming", True)), |
| 26 | "automatic_naming_mode": mode, |
| 27 | } |
| 28 | |
| 29 | |
| 30 | def get_user_messages(agent: Agent, *, limit: int | None = RECENT_USER_MESSAGES) -> list[str]: |
| 31 | messages: list[str] = [] |
| 32 | for message in agent.history.all_messages(): |
| 33 | if message.ai: |
| 34 | continue |
| 35 | text = _user_message_text(message.content) |
| 36 | if text: |
| 37 | messages.append(text) |
| 38 | return messages[-limit:] if limit else messages |
| 39 | |
| 40 | |
| 41 | def latest_user_sequence(agent: Agent) -> int: |
| 42 | for message in reversed(agent.history.all_messages()): |
| 43 | if not message.ai and _user_message_text(message.content): |
| 44 | return int(message.sequence or 0) |
| 45 | return 0 |
| 46 | |
| 47 | |
| 48 | async def generate_name( |
| 49 | agent: Agent, |
| 50 | *, |
| 51 | user_messages: list[str] | None = None, |
| 52 | current_name: str | None = None, |
| 53 | ) -> str: |
| 54 | selected_messages = user_messages or get_user_messages(agent) |
| 55 | if not selected_messages: |
| 56 | raise ValueError("This chat has no user messages to name yet.") |
| 57 | |
| 58 | from plugins._model_config.helpers.model_config import get_utility_model_config |
| 59 | |
| 60 | utility_config = get_utility_model_config(agent) |
| 61 | context_length = int(utility_config.get("ctx_length", 128000) or 128000) |
| 62 | if context_length <= 0: |
| 63 | context_length = 128000 |
| 64 | input_budget = max(int(context_length * UTILITY_CONTEXT_INPUT_RATIO), 1) |
| 65 | |
| 66 | system_prompt = agent.read_prompt("fw.chat_naming.system.md") |
| 67 | resolved_name = ( |
| 68 | current_name if current_name is not None else agent.context.name |
| 69 | ) or "(unnamed)" |
| 70 | prompt_without_messages = agent.read_prompt( |
| 71 | "fw.chat_naming.message.md", |
| 72 | current_name=resolved_name, |
| 73 | user_messages="", |
| 74 | ) |
| 75 | fixed_tokens = _estimated_input_tokens(system_prompt, prompt_without_messages) |
| 76 | message_budget = input_budget - fixed_tokens |
| 77 | if message_budget <= 0: |
| 78 | raise ValueError("The Utility Model context window is too small for chat naming.") |
| 79 | |
| 80 | message_text = _fit_user_messages(selected_messages, message_budget) |
| 81 | user_prompt = agent.read_prompt( |
| 82 | "fw.chat_naming.message.md", |
| 83 | current_name=resolved_name, |
| 84 | user_messages=message_text, |
| 85 | ) |
| 86 | |
| 87 | estimated_tokens = _estimated_input_tokens(system_prompt, user_prompt) |
| 88 | while estimated_tokens > input_budget and message_text: |
| 89 | excess = estimated_tokens - input_budget |
| 90 | reduced_budget = max( |
| 91 | tokens.approximate_tokens(message_text) - excess - 1, |
| 92 | 1, |
| 93 | ) |
| 94 | trimmed = _trim_to_estimated_tokens(message_text, reduced_budget) |
| 95 | if trimmed == message_text: |
| 96 | break |
| 97 | message_text = trimmed |
| 98 | user_prompt = agent.read_prompt( |
| 99 | "fw.chat_naming.message.md", |
| 100 | current_name=resolved_name, |
| 101 | user_messages=message_text, |
| 102 | ) |
| 103 | estimated_tokens = _estimated_input_tokens(system_prompt, user_prompt) |
| 104 | |
| 105 | if estimated_tokens > input_budget: |
| 106 | raise ValueError("Chat naming input exceeds the Utility Model context budget.") |
| 107 | |
| 108 | response = await agent.call_utility_model( |
| 109 | system=system_prompt, |
| 110 | message=user_prompt, |
| 111 | background=True, |
| 112 | ) |
| 113 | name = normalize_generated_name(response) |
| 114 | if not name: |
| 115 | raise ValueError("The Utility Model did not return a chat name.") |
| 116 | return name |
| 117 | |
| 118 | |
| 119 | def save_context_name(agent: Agent, name: str) -> str: |
| 120 | normalized = normalize_manual_name(name) |
| 121 | agent.context.name = normalized |
| 122 | if "parent_context_label" in agent.context.output_data: |
| 123 | agent.context.output_data["parent_context_label"] = normalized |
| 124 | persist_chat.save_tmp_chat(agent.context) |
| 125 | mark_dirty_all(reason="plugins._chat_naming.save_context_name") |
| 126 | return normalized |
| 127 | |
| 128 | |
| 129 | def normalize_generated_name(value: object) -> str: |
| 130 | name = " ".join(str(value or "").split()).strip(" \"'`#") |
| 131 | if len(name) > GENERATED_NAME_LIMIT: |
| 132 | name = name[: GENERATED_NAME_LIMIT - 3].rstrip() + "..." |
| 133 | return name |
| 134 | |
| 135 | |
| 136 | def normalize_manual_name(value: object) -> str: |
| 137 | name = " ".join(str(value or "").split()) |
| 138 | if not name: |
| 139 | raise ValueError("Name is required.") |
| 140 | if len(name) > 200: |
| 141 | raise ValueError("Name must be 200 characters or fewer.") |
| 142 | return name |
| 143 | |
| 144 | |
| 145 | def _user_message_text(content: object) -> str: |
| 146 | if not isinstance(content, dict): |
| 147 | return "" |
| 148 | for key in ("user_message", "user_intervention"): |
| 149 | value = content.get(key) |
| 150 | if isinstance(value, str): |
| 151 | return value.strip() |
| 152 | if value: |
| 153 | return json.dumps(value, ensure_ascii=False) |
| 154 | return "" |
| 155 | |
| 156 | |
| 157 | def _fit_user_messages(messages: list[str], token_budget: int) -> str: |
| 158 | selected: list[str] = [] |
| 159 | for message in reversed(messages): |
| 160 | candidate = [message, *selected] |
| 161 | text = _format_user_messages(candidate) |
| 162 | if tokens.approximate_tokens(text) <= token_budget: |
| 163 | selected = candidate |
| 164 | continue |
| 165 | return _trim_to_estimated_tokens(text, token_budget) |
| 166 | return _format_user_messages(selected) |
| 167 | |
| 168 | |
| 169 | def _estimated_input_tokens(system_prompt: str, user_prompt: str) -> int: |
| 170 | return tokens.approximate_tokens(system_prompt) + tokens.approximate_tokens( |
| 171 | user_prompt |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | def _format_user_messages(messages: list[str]) -> str: |
| 176 | return "\n\n".join( |
| 177 | f"{index}. {text}" for index, text in enumerate(messages, start=1) |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | def _trim_to_estimated_tokens(text: str, token_budget: int) -> str: |
| 182 | if tokens.approximate_tokens(text) <= token_budget: |
| 183 | return text |
| 184 | |
| 185 | exact_budget = max(int(token_budget / tokens.APPROX_BUFFER) - 1, 1) |
| 186 | trimmed = tokens.trim_to_tokens(text, exact_budget, "end") |
| 187 | while tokens.approximate_tokens(trimmed) > token_budget and exact_budget > 1: |
| 188 | exact_budget = max(int(exact_budget * 0.8), 1) |
| 189 | trimmed = tokens.trim_to_tokens(text, exact_budget, "end") |
| 190 | return trimmed |