NoteStack-Test / backend /tests /test_chat_trace.py
HumeraR's picture
feat: add MemoryNotes component to Settings page and enhance UI styles
b74c6ae
Raw History Blame Contribute Delete
12.9 kB
"""The per-message trace block (logger notestack.chat_trace): what Jev chose and what our thresholds changed, which path
ran, what the agent was given, and how long each stage took."""
import logging
import uuid
from contextlib import contextmanager
from types import SimpleNamespace
from app.chat_memory.files import TopicEntry, TopicsFile
from app.config import settings
from app.pipeline import research as r
from app.pipeline.trace import ChatTrace, ms
from app.services import jev
def _memory(*slugs):
tf = TopicsFile(chat_id="c", entries=[TopicEntry(s, s.title(), "2026-10-01", 1, "kw", "gist") for s in slugs])
return SimpleNamespace(tf=tf, latest=lambda: slugs[0] if slugs else None)
def _jev(monkeypatch, **answers):
monkeypatch.setattr(settings, "typesafe_api_key", "sk-test")
monkeypatch.setattr(jev, "decide", lambda *a, **k: answers)
def test_the_trace_shows_every_jev_answer_and_what_a_threshold_changed(monkeypatch):
_jev(monkeypatch,
intent=jev.Choice("archive", 0.91, {"archive": 0.91, "off_topic": 0.09}),
depth=jev.Choice("deep", 0.77, {"deep": 0.77, "quick": 0.23}),
memory=jev.Choice("replay", 0.41, {"replay": 0.41, "compose": 0.39, "lookup": 0.2}),
topic=jev.Choice("evals", 0.3, {"evals": 0.3, "none": 0.7}))
route = r.triage("what did you say about MMLU?", [], "nb", _memory("evals"))
assert (route.kind, route.memory, route.topic, route.engine) == ("archive", "lookup", "", "jev")
assert route.jev["memory"] == {"choice": "replay", "confidence": 0.41,
"probabilities": {"replay": 0.41, "compose": 0.39, "lookup": 0.2}}
assert any("memory 'replay' confidence 0.41" in n and "-> lookup" in n for n in route.notes)
assert any("topic 'evals'" in n and "confidence 0.30" in n for n in route.notes)
trace = ChatTrace(notebook="nb", message="what did you say about MMLU?")
r._trace_triage(trace, route)
text = trace.render()
assert "engine : jev" in text
assert "memory : replay" in text and "confidence 0.41" in text and "compose 0.39" in text
assert "overrides: memory 'replay' confidence 0.41" in text
assert "FINAL : kind=archive depth=deep memory=lookup topic=-" in text
assert "memory=lookup means:" in text
def test_a_trace_without_jev_says_why_and_that_there_are_no_confidences(monkeypatch):
monkeypatch.setattr(settings, "typesafe_api_key", "")
monkeypatch.setattr(r, "triage_llm", lambda *a: r.Triage("archive", "q", "quick", ""))
route = r.triage("hello?", [], "nb")
trace = ChatTrace()
r._trace_triage(trace, route)
assert "llm triage (Jev is not configured" in trace.render() and "no probabilities or confidences" in trace.render()
monkeypatch.setattr(settings, "typesafe_api_key", "sk-test")
def down(*a, **k):
raise jev.JevError("503 busy")
monkeypatch.setattr(jev, "decide", down)
assert "Jev failed: 503 busy" in r.triage("hello?", [], "nb").engine
def test_each_path_says_which_one_ran(monkeypatch, tmp_path):
from app.corpus import Corpus
corpus = Corpus(uuid.uuid4(), cache_dir=str(tmp_path))
allowed = {"a.md": "On Pricing"}
t = ChatTrace()
r.research(corpus, {}, "hi", trace=t)
assert "NO POSTS" in t.render()
t = ChatTrace()
r.research(corpus, allowed, "thanks!", trace=t)
assert "SMALL TALK" in t.render() and "no model, Jev call or R2 read" in t.render()
t = ChatTrace()
monkeypatch.setattr(r, "triage", lambda *a: r.Triage("off_topic", "q", "quick", "", engine="jev"))
r.research(corpus, allowed, "what is the capital of France?", trace=t)
out = t.render()
assert "OFF_TOPIC: triage decided this needs no research" in out and "a fixed template" in out
assert "post reads" in out and "[TIMING] triage" in out
def test_the_archive_trace_lists_what_the_agent_was_given(monkeypatch, tmp_path):
from app.corpus import Corpus
@contextmanager
def no_usage(lm):
yield SimpleNamespace(prompt_tokens=11, completion_tokens=7)
seen = {}
def fake_react(signature, tools, max_iters):
seen["tools"] = [t.__name__ for t in tools]
def run(**inputs):
return SimpleNamespace(answer="x", citations=[], unsupported=True)
return run
monkeypatch.setattr(r.dspy, "ReAct", fake_react)
monkeypatch.setattr(r, "main_lm", lambda: SimpleNamespace(model="glm-test"))
monkeypatch.setattr(r, "track_usage", no_usage)
monkeypatch.setattr(r, "triage", lambda *a: r.Triage("archive", "q", "quick", "", engine="jev"))
corpus = Corpus(uuid.uuid4(), cache_dir=str(tmp_path))
monkeypatch.setattr(corpus, "sync", lambda: {})
monkeypatch.setattr(settings, "chat_memory_read", False)
t = ChatTrace()
t.step("search", "pric(e|ing)")
r.research(corpus, {"a.md": "On Pricing"}, "What did I charge?", profile="- audience: founders\n- tone: brief",
history=[r.Turn("user", "hi there"), r.Turn("assistant", "hello")], trace=t)
text = t.render()
assert "ARCHIVE: the message needs the writer's posts" in text
assert "signature: ResearchArchive | model glm-test | up to 4 reasoning steps (quick)" in text
assert "tools : list_files, search, read" in text
assert "conversation: 2 turn(s)" in text and "writer notes: 2 line(s)" in text
assert "chat memory: OFF (CHAT_MEMORY_READ=false)" in text
assert "the model proposed 0, verified 0" in text and "tokens prompt/completion: triage 0/0, agent 11/7" in text
assert "could not find support" not in text # unsupported=True: the model's own "not covered" answer is kept
assert "search" in text and "[posts, local disk cache]" in text
assert "[TIMING] triage" in text and "sync-posts" in text and "agent" in text and "verify" in text
def test_the_chat_endpoint_writes_one_trace_block_per_message(client, monkeypatch, caplog):
from app.routers import notebooks as nbr
from tests.test_memory import _signup
auth = _signup(client, "trace1@example.com")
nb = client.post("/api/notebooks", json={"title": "My notebook"}, headers=auth).json()
def fake_research(corpus, allowed, question, on_step=None, history=None, notebook_title="", profile="", **kw):
from app.pipeline.trace import current_trace
current_trace().add("route", "ARCHIVE: fake") # research() finds the request's trace through the context
return r.ResearchResult("ok", True, [], kind="archive", question=question)
monkeypatch.setattr(nbr, "research", fake_research)
monkeypatch.setattr(settings, "chat_context_log", True)
monkeypatch.setattr(settings, "chat_memory_read", False)
with caplog.at_level(logging.INFO, logger="notestack.chat_trace"):
client.post(f"/api/notebooks/{nb['id']}/chat", json={"question": "hello there friend"}, headers=auth)
blocks = [rec.getMessage() for rec in caplog.records if rec.name == "notestack.chat_trace"]
assert len(blocks) == 1
text = blocks[0]
assert text.startswith("CHAT TRACE chat=") and "notebook='My notebook'" in text
assert "'hello there friend'" in text # text logging is on
assert "[1 LOAD (Postgres, R2)]" in text and "history: 0 message(s) of this chat [Postgres]" in text
assert "chat memory: OFF (CHAT_MEMORY_READ=false), nothing read from R2" in text
assert "[2 PATH]" in text and "ARCHIVE: fake" in text
assert "answer sent after" in text and "assistant message saved to Postgres" in text
assert "workspace notes: not updated (no LLM_API_KEY)" in text
assert "[TIMING] db-history+notes" in text and "save-answer" in text
def test_the_trace_hides_the_message_when_text_logging_is_off(client, monkeypatch, caplog):
from app.routers import notebooks as nbr
from tests.test_memory import _signup
auth = _signup(client, "trace2@example.com")
nb = client.post("/api/notebooks", json={"title": "N"}, headers=auth).json()
monkeypatch.setattr(nbr, "research", lambda *a, **k: r.ResearchResult("ok", True, [], kind="archive"))
monkeypatch.setattr(settings, "chat_context_log", False)
with caplog.at_level(logging.INFO, logger="notestack.chat_trace"):
client.post(f"/api/notebooks/{nb['id']}/chat", json={"question": "my secret question"}, headers=auth)
assert "CHAT TRACE" in caplog.text and "my secret question" not in caplog.text
assert "message : 18 chars" in caplog.text
def test_the_notes_gate_logs_what_jev_said_and_what_it_did(monkeypatch, caplog):
from app.pipeline import memory as m
monkeypatch.setattr(settings, "typesafe_api_key", "sk-test")
monkeypatch.setattr(jev, "decide", lambda *a, **k: {"remember": jev.Choice("no", 0.93, {"no": 0.93, "yes": 0.07})})
with caplog.at_level(logging.INFO, logger="app.pipeline.memory"):
assert m.worth_remembering("What did I write about pricing?") is False
assert "NOTES GATE Jev says 'no', confidence 0.93 [no 0.93 yes 0.07] -> SKIP the extraction" in caplog.text
caplog.clear()
monkeypatch.setattr(jev, "decide", lambda *a, **k: {"remember": jev.Choice("no", 0.6, {"no": 0.6, "yes": 0.4})})
with caplog.at_level(logging.INFO, logger="app.pipeline.memory"):
assert m.worth_remembering("Please be brief. What did I write?") is True
assert "RUN the extraction model" in caplog.text
def test_durations_read_naturally():
assert ms(12) == "12ms" and ms(1500) == "1.5s"
# A model reply that is not in the shape ReAct needs (it sent only the tool's arguments)
def _malformed():
import dspy
from dspy.utils.exceptions import AdapterParseError
return AdapterParseError("JSONAdapter", dspy.Signature("question -> answer"),
'{"path": "sources/x.md", "start_line": 1, "end_line": 100}')
def _real_react(react, extract):
import dspy
def search(query: str) -> str:
"""Search."""
return "observation"
agent = dspy.ReAct("question -> answer", tools=[search], max_iters=3)
agent.react, agent.extract = react, extract
return agent
def test_react_still_has_the_hook_the_retry_uses():
import dspy
assert hasattr(dspy.ReAct, "_call_with_potential_trajectory_truncation") # fails loudly if dspy is upgraded
def test_a_malformed_step_is_asked_again_once_with_a_different_request():
import dspy
configs = []
def react(**kw):
configs.append(kw.get("config"))
if len(configs) == 1:
raise _malformed()
return dspy.Prediction(next_thought="done", next_tool_name="finish", next_tool_args={})
trace = ChatTrace()
agent = r._retry_malformed_steps(_real_react(react, lambda **kw: dspy.Prediction(answer="ok")), trace)
assert agent(question="q").answer == "ok"
assert configs == [None, {"temperature": r.RETRY_TEMPERATURE}] # the retry is not the cached request
assert "not in the expected shape" in trace.render() and "asking again once" in trace.render()
def test_a_step_that_stays_malformed_ends_the_research_but_still_answers():
import dspy
calls = []
def react(**kw):
calls.append(1)
raise _malformed()
trace = ChatTrace()
agent = r._retry_malformed_steps(_real_react(react, lambda **kw: dspy.Prediction(answer="from what I had")), trace)
assert agent(question="q").answer == "from what I had" # ReAct stopped and wrote the answer from the trajectory
assert len(calls) == 2 and "still malformed" in trace.render() # asked twice, not more
def test_an_answer_step_that_stays_malformed_saves_a_plain_reply_and_is_not_filed(monkeypatch, tmp_path):
from app.chat_memory.organize import is_off_topic
from app.corpus import Corpus
def fake_react(signature, tools, max_iters):
def run(**inputs):
raise _malformed()
return run
@contextmanager
def no_usage(lm):
yield SimpleNamespace(prompt_tokens=0, completion_tokens=0)
monkeypatch.setattr(r.dspy, "ReAct", fake_react)
monkeypatch.setattr(r, "main_lm", lambda: SimpleNamespace(model="m"))
monkeypatch.setattr(r, "track_usage", no_usage)
monkeypatch.setattr(r, "triage", lambda *a: r.Triage("archive", "q", "quick", "", engine="jev"))
corpus = Corpus(uuid.uuid4(), cache_dir=str(tmp_path))
monkeypatch.setattr(corpus, "sync", lambda: {})
t = ChatTrace()
out = r.research(corpus, {"a.md": "On Pricing"}, "What did I charge?", trace=t)
assert out.text == r.FAILED_REPLY and out.kind == "error" and out.unsupported and out.citations == []
assert out.recall["failed"] == "malformed model reply" and "plain 'please send that again'" in t.render()
round_ = SimpleNamespace(assistant=SimpleNamespace(recall_json={"kind": out.kind}, citations=[], content=out.text))
assert is_off_topic(round_) # the failure reply is never filed as a topic