返回 JoyAI-Echo
helpers.py
1 """Utility functions for nanobot."""
2
3 import base64
4 import json
5 import re
6 import shutil
7 import time
8 import uuid
9 from datetime import datetime
10 from pathlib import Path
11 from typing import Any
12
13 import tiktoken
14 from loguru import logger
15
16
17 def strip_think(text: str) -> str:
18 """Remove thinking blocks, unclosed trailing tags, and tokenizer-level
19 template leaks occasionally emitted by some models (notably Gemma 4's
20 Ollama renderer).
21
22 Covers:
23 1. Well-formed `<think>...</think>` and `<thought>...</thought>` blocks.
24 2. Streaming prefixes where the block is never closed.
25 3. *Malformed* opening tags missing the `>` — e.g. `<think广场…`. The
26 model sometimes emits the tag name directly followed by user-facing
27 content with no delimiter; without this step the literal `<think`
28 leaks into the rendered message.
29 4. Harmony-style channel markers like `<channel|>` / `<|channel|>`
30 **at the start of the text** — conservative to avoid eating
31 explanatory prose that mentions these tokens.
32 5. Orphan closing tags `</think>` / `</thought>` **at the very start
33 or end of the text** only, for the same reason.
34
35 Since this is also applied before persisting to history (memory.py),
36 the edge-only stripping of (4) and (5) is deliberate: stripping those
37 tokens mid-text would silently rewrite any message where a user or the
38 assistant discusses the tokens themselves.
39 """
40 # Well-formed blocks first.
41 text = re.sub(r"<think>[\s\S]*?</think>", "", text)
42 text = re.sub(r"^\s*<think>[\s\S]*$", "", text)
43 text = re.sub(r"<thought>[\s\S]*?</thought>", "", text)
44 text = re.sub(r"^\s*<thought>[\s\S]*$", "", text)
45 # Malformed opening tags: `<think` / `<thought` where the next char is
46 # NOT one that could continue a valid tag / identifier name. Explicitly
47 # listing ASCII tag-name chars (letters, digits, `_`, `-`, `:`) plus
48 # `>` / `/` — we can't use `\w` here because in Python's default
49 # Unicode regex mode it matches CJK characters too, which would defeat
50 # the primary fix for `<think广场…` leaks.
51 text = re.sub(r"<think(?![A-Za-z0-9_\-:>/])", "", text)
52 text = re.sub(r"<thought(?![A-Za-z0-9_\-:>/])", "", text)
53 # Edge-only orphan closing tags (start or end of text).
54 text = re.sub(r"^\s*</think>\s*", "", text)
55 text = re.sub(r"\s*</think>\s*$", "", text)
56 text = re.sub(r"^\s*</thought>\s*", "", text)
57 text = re.sub(r"\s*</thought>\s*$", "", text)
58 # Edge-only channel markers (harmony / Gemma 4 variant leaks).
59 text = re.sub(r"^\s*<\|?channel\|?>\s*", "", text)
60 return text.strip()
61
62
63 def detect_image_mime(data: bytes) -> str | None:
64 """Detect image MIME type from magic bytes, ignoring file extension."""
65 if data[:8] == b"\x89PNG\r\n\x1a\n":
66 return "image/png"
67 if data[:3] == b"\xff\xd8\xff":
68 return "image/jpeg"
69 if data[:6] in (b"GIF87a", b"GIF89a"):
70 return "image/gif"
71 if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
72 return "image/webp"
73 return None
74
75
76 def build_image_content_blocks(
77 raw: bytes, mime: str, path: str, label: str
78 ) -> list[dict[str, Any]]:
79 """Build native image blocks plus a short text label."""
80 b64 = base64.b64encode(raw).decode()
81 return [
82 {
83 "type": "image_url",
84 "image_url": {"url": f"data:{mime};base64,{b64}"},
85 "_meta": {"path": path},
86 },
87 {"type": "text", "text": label},
88 ]
89
90
91 def ensure_dir(path: Path) -> Path:
92 """Ensure directory exists, return it."""
93 path.mkdir(parents=True, exist_ok=True)
94 return path
95
96
97 def timestamp() -> str:
98 """Current ISO timestamp."""
99 return datetime.now().isoformat()
100
101
102 def current_time_str(timezone: str | None = None) -> str:
103 """Return the current time string."""
104 from zoneinfo import ZoneInfo
105
106 try:
107 tz = ZoneInfo(timezone) if timezone else None
108 except (KeyError, Exception):
109 tz = None
110
111 now = datetime.now(tz=tz) if tz else datetime.now().astimezone()
112 offset = now.strftime("%z")
113 offset_fmt = f"{offset[:3]}:{offset[3:]}" if len(offset) == 5 else offset
114 tz_name = timezone or (time.strftime("%Z") or "UTC")
115 return f"{now.strftime('%Y-%m-%d %H:%M (%A)')} ({tz_name}, UTC{offset_fmt})"
116
117
118 _UNSAFE_CHARS = re.compile(r'[<>:"/\\|?*]')
119 _TOOL_RESULT_PREVIEW_CHARS = 1200
120 _TOOL_RESULTS_DIR = ".nanobot/tool-results"
121 _TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60
122 _TOOL_RESULT_MAX_BUCKETS = 32
123
124
125 def safe_filename(name: str) -> str:
126 """Replace unsafe path characters with underscores."""
127 return _UNSAFE_CHARS.sub("_", name).strip()
128
129
130 def image_placeholder_text(path: str | None, *, empty: str = "[image]") -> str:
131 """Build an image placeholder string."""
132 return f"[image: {path}]" if path else empty
133
134
135 def truncate_text(text: str, max_chars: int) -> str:
136 """Truncate text with a stable suffix."""
137 if max_chars <= 0 or len(text) <= max_chars:
138 return text
139 return text[:max_chars] + "\n... (truncated)"
140
141
142 def find_legal_message_start(messages: list[dict[str, Any]]) -> int:
143 """Find the first index whose tool results have matching assistant calls."""
144 declared: set[str] = set()
145 start = 0
146 for i, msg in enumerate(messages):
147 role = msg.get("role")
148 if role == "assistant":
149 for tc in msg.get("tool_calls") or []:
150 if isinstance(tc, dict) and tc.get("id"):
151 declared.add(str(tc["id"]))
152 elif role == "tool":
153 tid = msg.get("tool_call_id")
154 if tid and str(tid) not in declared:
155 start = i + 1
156 declared.clear()
157 for prev in messages[start : i + 1]:
158 if prev.get("role") == "assistant":
159 for tc in prev.get("tool_calls") or []:
160 if isinstance(tc, dict) and tc.get("id"):
161 declared.add(str(tc["id"]))
162 return start
163
164
165 def stringify_text_blocks(content: list[dict[str, Any]]) -> str | None:
166 parts: list[str] = []
167 for block in content:
168 if not isinstance(block, dict):
169 return None
170 if block.get("type") != "text":
171 return None
172 text = block.get("text")
173 if not isinstance(text, str):
174 return None
175 parts.append(text)
176 return "\n".join(parts)
177
178
179 def _render_tool_result_reference(
180 filepath: Path,
181 *,
182 original_size: int,
183 preview: str,
184 truncated_preview: bool,
185 ) -> str:
186 result = (
187 f"[tool output persisted]\n"
188 f"Full output saved to: {filepath}\n"
189 f"Original size: {original_size} chars\n"
190 f"Preview:\n{preview}"
191 )
192 if truncated_preview:
193 result += "\n...\n(Read the saved file if you need the full output.)"
194 return result
195
196
197 def _bucket_mtime(path: Path) -> float:
198 try:
199 return path.stat().st_mtime
200 except OSError:
201 return 0.0
202
203
204 def _cleanup_tool_result_buckets(root: Path, current_bucket: Path) -> None:
205 siblings = [path for path in root.iterdir() if path.is_dir() and path != current_bucket]
206 cutoff = time.time() - _TOOL_RESULT_RETENTION_SECS
207 for path in siblings:
208 if _bucket_mtime(path) < cutoff:
209 shutil.rmtree(path, ignore_errors=True)
210 keep = max(_TOOL_RESULT_MAX_BUCKETS - 1, 0)
211 siblings = [path for path in siblings if path.exists()]
212 if len(siblings) <= keep:
213 return
214 siblings.sort(key=_bucket_mtime, reverse=True)
215 for path in siblings[keep:]:
216 shutil.rmtree(path, ignore_errors=True)
217
218
219 def _write_text_atomic(path: Path, content: str) -> None:
220 tmp = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
221 try:
222 tmp.write_text(content, encoding="utf-8")
223 tmp.replace(path)
224 finally:
225 if tmp.exists():
226 tmp.unlink(missing_ok=True)
227
228
229 def write_json_atomic(path: Path, value: Any) -> None:
230 """Write JSON so concurrent readers never see a truncated file."""
231 path.parent.mkdir(parents=True, exist_ok=True)
232 _write_text_atomic(
233 path,
234 json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
235 )
236
237
238 def maybe_persist_tool_result(
239 workspace: Path | None,
240 session_key: str | None,
241 tool_call_id: str,
242 content: Any,
243 *,
244 max_chars: int,
245 ) -> Any:
246 """Persist oversized tool output and replace it with a stable reference string."""
247 if workspace is None or max_chars <= 0:
248 return content
249
250 text_payload: str | None = None
251 suffix = "txt"
252 if isinstance(content, str):
253 text_payload = content
254 elif isinstance(content, list):
255 text_payload = stringify_text_blocks(content)
256 if text_payload is None:
257 return content
258 suffix = "json"
259 else:
260 return content
261
262 if len(text_payload) <= max_chars:
263 return content
264
265 root = ensure_dir(workspace / _TOOL_RESULTS_DIR)
266 bucket = ensure_dir(root / safe_filename(session_key or "default"))
267 try:
268 _cleanup_tool_result_buckets(root, bucket)
269 except Exception as exc:
270 logger.warning("Failed to clean stale tool result buckets in {}: {}", root, exc)
271 path = bucket / f"{safe_filename(tool_call_id)}.{suffix}"
272 if not path.exists():
273 if suffix == "json" and isinstance(content, list):
274 _write_text_atomic(path, json.dumps(content, ensure_ascii=False, indent=2))
275 else:
276 _write_text_atomic(path, text_payload)
277
278 preview = text_payload[:_TOOL_RESULT_PREVIEW_CHARS]
279 return _render_tool_result_reference(
280 path,
281 original_size=len(text_payload),
282 preview=preview,
283 truncated_preview=len(text_payload) > _TOOL_RESULT_PREVIEW_CHARS,
284 )
285
286
287 def split_message(content: str, max_len: int = 2000) -> list[str]:
288 """
289 Split content into chunks within max_len, preferring line breaks.
290
291 Args:
292 content: The text content to split.
293 max_len: Maximum length per chunk (default 2000 for Discord compatibility).
294
295 Returns:
296 List of message chunks, each within max_len.
297 """
298 if not content:
299 return []
300 if len(content) <= max_len:
301 return [content]
302 chunks: list[str] = []
303 while content:
304 if len(content) <= max_len:
305 chunks.append(content)
306 break
307 cut = content[:max_len]
308 # Try to break at newline first, then space, then hard break
309 pos = cut.rfind("\n")
310 if pos <= 0:
311 pos = cut.rfind(" ")
312 if pos <= 0:
313 pos = max_len
314 chunks.append(content[:pos])
315 content = content[pos:].lstrip()
316 return chunks
317
318
319 def build_assistant_message(
320 content: str | None,
321 tool_calls: list[dict[str, Any]] | None = None,
322 reasoning_content: str | None = None,
323 thinking_blocks: list[dict] | None = None,
324 ) -> dict[str, Any]:
325 """Build a provider-safe assistant message with optional reasoning fields."""
326 msg: dict[str, Any] = {"role": "assistant", "content": content or ""}
327 if tool_calls:
328 msg["tool_calls"] = tool_calls
329 if reasoning_content is not None or thinking_blocks:
330 msg["reasoning_content"] = reasoning_content if reasoning_content is not None else ""
331 if thinking_blocks:
332 msg["thinking_blocks"] = thinking_blocks
333 return msg
334
335
336 def estimate_prompt_tokens(
337 messages: list[dict[str, Any]],
338 tools: list[dict[str, Any]] | None = None,
339 ) -> int:
340 """Estimate prompt tokens with tiktoken.
341
342 Counts all fields that providers send to the LLM: content, tool_calls,
343 reasoning_content, tool_call_id, name, plus per-message framing overhead.
344 """
345 try:
346 enc = tiktoken.get_encoding("cl100k_base")
347 parts: list[str] = []
348 for msg in messages:
349 content = msg.get("content")
350 if isinstance(content, str):
351 parts.append(content)
352 elif isinstance(content, list):
353 for part in content:
354 if isinstance(part, dict) and part.get("type") == "text":
355 txt = part.get("text", "")
356 if txt:
357 parts.append(txt)
358
359 tc = msg.get("tool_calls")
360 if tc:
361 parts.append(json.dumps(tc, ensure_ascii=False))
362
363 rc = msg.get("reasoning_content")
364 if isinstance(rc, str) and rc:
365 parts.append(rc)
366
367 for key in ("name", "tool_call_id"):
368 value = msg.get(key)
369 if isinstance(value, str) and value:
370 parts.append(value)
371
372 if tools:
373 parts.append(json.dumps(tools, ensure_ascii=False))
374
375 per_message_overhead = len(messages) * 4
376 return len(enc.encode("\n".join(parts))) + per_message_overhead
377 except Exception:
378 return 0
379
380
381 def estimate_message_tokens(message: dict[str, Any]) -> int:
382 """Estimate prompt tokens contributed by one persisted message."""
383 content = message.get("content")
384 parts: list[str] = []
385 if isinstance(content, str):
386 parts.append(content)
387 elif isinstance(content, list):
388 for part in content:
389 if isinstance(part, dict) and part.get("type") == "text":
390 text = part.get("text", "")
391 if text:
392 parts.append(text)
393 else:
394 parts.append(json.dumps(part, ensure_ascii=False))
395 elif content is not None:
396 parts.append(json.dumps(content, ensure_ascii=False))
397
398 for key in ("name", "tool_call_id"):
399 value = message.get(key)
400 if isinstance(value, str) and value:
401 parts.append(value)
402 if message.get("tool_calls"):
403 parts.append(json.dumps(message["tool_calls"], ensure_ascii=False))
404
405 rc = message.get("reasoning_content")
406 if isinstance(rc, str) and rc:
407 parts.append(rc)
408
409 payload = "\n".join(parts)
410 if not payload:
411 return 4
412 try:
413 enc = tiktoken.get_encoding("cl100k_base")
414 return max(4, len(enc.encode(payload)) + 4)
415 except Exception:
416 return max(4, len(payload) // 4 + 4)
417
418
419 def estimate_prompt_tokens_chain(
420 provider: Any,
421 model: str | None,
422 messages: list[dict[str, Any]],
423 tools: list[dict[str, Any]] | None = None,
424 ) -> tuple[int, str]:
425 """Estimate prompt tokens via provider counter first, then tiktoken fallback."""
426 provider_counter = getattr(provider, "estimate_prompt_tokens", None)
427 if callable(provider_counter):
428 try:
429 tokens, source = provider_counter(messages, tools, model)
430 if isinstance(tokens, (int, float)) and tokens > 0:
431 return int(tokens), str(source or "provider_counter")
432 except Exception:
433 pass
434
435 estimated = estimate_prompt_tokens(messages, tools)
436 if estimated > 0:
437 return int(estimated), "tiktoken"
438 return 0, "none"
439
440
441 def build_status_content(
442 *,
443 version: str,
444 model: str,
445 start_time: float,
446 last_usage: dict[str, int],
447 context_window_tokens: int,
448 session_msg_count: int,
449 context_tokens_estimate: int,
450 search_usage_text: str | None = None,
451 active_task_count: int = 0,
452 max_completion_tokens: int = 8192,
453 ) -> str:
454 """Build a human-readable runtime status snapshot.
455
456 Args:
457 search_usage_text: Optional pre-formatted web search usage string
458 (produced by SearchUsageInfo.format()). When provided
459 it is appended as an extra section.
460 """
461 uptime_s = int(time.time() - start_time)
462 uptime = (
463 f"{uptime_s // 3600}h {(uptime_s % 3600) // 60}m"
464 if uptime_s >= 3600
465 else f"{uptime_s // 60}m {uptime_s % 60}s"
466 )
467 last_in = last_usage.get("prompt_tokens", 0)
468 last_out = last_usage.get("completion_tokens", 0)
469 cached = last_usage.get("cached_tokens", 0)
470 ctx_total = max(context_window_tokens, 0)
471 # Budget mirrors Consolidator formula: ctx_window - max_completion - _SAFETY_BUFFER
472 ctx_budget = max(ctx_total - int(max_completion_tokens) - 1024, 1)
473 ctx_pct = min(int((context_tokens_estimate / ctx_budget) * 100), 999) if ctx_budget > 0 else 0
474 ctx_used_str = (
475 f"{context_tokens_estimate // 1000}k"
476 if context_tokens_estimate >= 1000
477 else str(context_tokens_estimate)
478 )
479 ctx_total_str = f"{ctx_total // 1000}k" if ctx_total > 0 else "n/a"
480 token_line = f"\U0001f4ca Tokens: {last_in} in / {last_out} out"
481 if cached and last_in:
482 token_line += f" ({cached * 100 // last_in}% cached)"
483 lines = [
484 f"\U0001f408 nanobot v{version}",
485 f"\U0001f9e0 Model: {model}",
486 token_line,
487 f"\U0001f4da Context: {ctx_used_str}/{ctx_total_str} ({ctx_pct}% of input budget)",
488 f"\U0001f4ac Session: {session_msg_count} messages",
489 f"\u23f1 Uptime: {uptime}",
490 f"\u26a1 Tasks: {active_task_count} active",
491 ]
492 if search_usage_text:
493 lines.append(search_usage_text)
494 return "\n".join(lines)
495
496
497 def sync_workspace_templates(workspace: Path, silent: bool = False) -> list[str]:
498 """Sync bundled templates to workspace. Only creates missing files."""
499 from importlib.resources import files as pkg_files
500
501 try:
502 tpl = pkg_files("nanobot") / "templates"
503 except Exception:
504 return []
505 if not tpl.is_dir():
506 return []
507
508 added: list[str] = []
509
510 def _write(src, dest: Path):
511 if dest.exists():
512 return
513 dest.parent.mkdir(parents=True, exist_ok=True)
514 dest.write_text(src.read_text(encoding="utf-8") if src else "", encoding="utf-8")
515 added.append(str(dest.relative_to(workspace)))
516
517 for item in tpl.iterdir():
518 if item.name.endswith(".md") and not item.name.startswith("."):
519 _write(item, workspace / item.name)
520 _write(tpl / "memory" / "MEMORY.md", workspace / "memory" / "MEMORY.md")
521 _write(None, workspace / "memory" / "history.jsonl")
522 (workspace / "skills").mkdir(exist_ok=True)
523
524 if added and not silent:
525 from rich.console import Console
526
527 for name in added:
528 Console().print(f" [dim]Created {name}[/dim]")
529
530 # Initialize git for memory version control
531 try:
532 from nanobot.utils.gitstore import GitStore
533
534 gs = GitStore(
535 workspace,
536 tracked_files=[
537 "SOUL.md",
538 "USER.md",
539 "memory/MEMORY.md",
540 ],
541 )
542 gs.init()
543 except Exception:
544 logger.warning("Failed to initialize git store for {}", workspace)
545
546 return added
547
547 lines PYTHON