Spaces:
Runtime error
Runtime error
| """One readable block in the server log for every chat message: what ran, in what order, what Jev or the triage model | |
| chose (and why a threshold changed it), what each model was given, which tools the agent called, and how long every | |
| stage took. | |
| Logger `notestack.chat_trace`. The structure is always logged. Writers' text (the message, the note values) is added | |
| only when text logging is on (the same switch as `notestack.chat_context`: CHAT_CONTEXT_LOG, or development). | |
| `notestack.chat_context` still holds the full text of everything the agent was given, for when you need to read it.""" | |
| import logging | |
| import time | |
| from contextlib import contextmanager | |
| from contextvars import ContextVar | |
| from app.chat_memory.logs import text_enabled | |
| trace_log = logging.getLogger("notestack.chat_trace") | |
| # What each memory level hands to the research agent (see chat_memory/context.py build_context). | |
| MEMORY_GIVES = { | |
| "none": "nothing extra: the agent only sees the last messages of this chat", | |
| "lookup": "the topic's summary (R2 file) + its last messages word for word (Postgres). No chat tools", | |
| "replay": "the topic's summary + its saved exact quotes (R2 file) + the 3 chat tools", | |
| "compose": "the notebook's topic index, first 20 lines (R2 file) + the 3 chat tools", | |
| } | |
| # Where each tool reads from, for the step list. | |
| STEP_SOURCE = {"list": "posts, local disk cache", "search": "posts, local disk cache", | |
| "read": "post, local disk cache", "recall": "earlier chats, local disk cache", "think": ""} | |
| SECTIONS = ("load", "route", "triage", "context", "steps", "citations", "result", "save", "notes") | |
| TITLES = {"load": "1 LOAD (Postgres, R2)", "route": "2 PATH", "triage": "3 TRIAGE", | |
| "context": "4 WHAT THE AGENT WAS GIVEN", "steps": "5 AGENT STEPS", "citations": "6 CITATIONS", | |
| "result": "7 RESULT", "save": "8 SAVED", "notes": "9 AFTER THE ANSWER"} | |
| def ms(value: float) -> str: | |
| return f"{value:.0f}ms" if value < 1000 else f"{value / 1000:.1f}s" | |
| class ChatTrace: | |
| def __init__(self, chat_id: str = "", notebook: str = "", message: str = ""): | |
| self.chat_id, self.notebook, self.message = chat_id, notebook, message | |
| self.t0 = time.perf_counter() | |
| self.stages: list[tuple[str, float]] = [] | |
| self.sections: dict[str, list[str]] = {} | |
| self.marks: dict[str, float] = {} | |
| self.agent_started: float | None = None | |
| self.error = "" | |
| # Recording | |
| def stage(self, name: str): | |
| start = time.perf_counter() | |
| try: | |
| yield | |
| finally: | |
| self.stages.append((name, (time.perf_counter() - start) * 1000)) | |
| def mark(self, name: str) -> None: | |
| """A moment on the clock, in ms since the message arrived (e.g. when the answer was sent).""" | |
| self.marks[name] = (time.perf_counter() - self.t0) * 1000 | |
| def add(self, section: str, *lines: str) -> None: | |
| self.sections.setdefault(section, []).extend(lines) | |
| def start_agent(self) -> None: | |
| self.agent_started = time.perf_counter() | |
| def step(self, kind: str, detail: str) -> None: | |
| offset = (time.perf_counter() - (self.agent_started or self.t0)) * 1000 | |
| source = STEP_SOURCE.get(kind, "") | |
| self.add("steps", f"+{ms(offset):>6} {kind:<6} {detail}" + (f" [{source}]" if source else "")) | |
| # Output | |
| def render(self) -> str: | |
| total = (time.perf_counter() - self.t0) * 1000 | |
| answered = self.marks.get("answer") | |
| head = f"CHAT TRACE chat={self.chat_id[:8] or '-'} notebook={self.notebook!r}" | |
| shown = f" {self.message[:140]!r}" if text_enabled() else "" | |
| lines = [head, f" message : {len(self.message)} chars{shown}"] | |
| if answered is not None: | |
| lines.append(f" clock : answer sent after {ms(answered)}, stream closed after {ms(total)}") | |
| else: | |
| lines.append(f" clock : stopped after {ms(total)}" + (f" ERROR: {self.error}" if self.error else "")) | |
| for key in SECTIONS: | |
| body = self.sections.get(key) | |
| if not body: | |
| continue | |
| lines.append(f" [{TITLES[key]}]") | |
| lines += [f" {ln}" for ln in body] | |
| if self.stages: | |
| lines.append(" [TIMING] " + " | ".join(f"{name} {ms(took)}" for name, took in self.stages)) | |
| return "\n".join(lines) | |
| def emit(self) -> None: | |
| try: | |
| trace_log.info(self.render()) | |
| except Exception: # a log line must never break a chat | |
| logging.getLogger(__name__).warning("Could not write the chat trace", exc_info=True) | |
| # research() runs in a worker thread (asyncio.to_thread copies the context), so the request hands its trace over | |
| # this way instead of through research()'s arguments. | |
| _current: ContextVar["ChatTrace | None"] = ContextVar("chat_trace", default=None) | |
| def use_trace(trace: "ChatTrace") -> None: | |
| _current.set(trace) | |
| def current_trace() -> "ChatTrace | None": | |
| return _current.get() | |
| def probs_text(probabilities: dict[str, float], limit: int = 5) -> str: | |
| top = sorted(probabilities.items(), key=lambda kv: kv[1], reverse=True)[:limit] | |
| return " ".join(f"{name} {p:.2f}" for name, p in top) | |