Spaces:
Runtime error
Runtime error
Download backend/tests/test_chat_trace.py from firebird-technologies/NoteStack-Test: direct link, hf CLI and curl.
- Browser
- Download file 12.9 kB
-
https://huggingface.co/spaces/firebird-technologies/NoteStack-Test/resolve/main/backend/tests/test_chat_trace.py
- Command line
-
hf download hf://spaces/firebird-technologies/NoteStack-Test/backend/tests/test_chat_trace.py
-
curl -L -o test_chat_trace.py https://huggingface.co/spaces/firebird-technologies/NoteStack-Test/resolve/main/backend/tests/test_chat_trace.py
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 | |
| 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 | |
| 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 | |