Spaces:
Runtime error
Runtime error
File size: 7,898 Bytes
94b5741 528cb2d 94b5741 528cb2d 94b5741 b74c6ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | import uuid
from datetime import UTC, datetime
import pytest
from app.corpus import INDEX, Corpus, CorpusError, ScopedTools, index_line, post_path, render_post
from app.llm.signatures import FileCitation
from app.pipeline.research import verify_citations
class MemoryStore:
"""Stands in for R2."""
def __init__(self):
self.objects: dict[str, bytes] = {}
def put_text(self, key, text, content_type=None):
self.objects[key] = text.encode()
return key
def get_bytes(self, key):
return self.objects[key]
def exists(self, key):
return key in self.objects
def delete(self, key):
self.objects.pop(key, None)
PRICING = [("Why I raised prices", "I moved the paid tier from $5 to $8 a month.\n\nChurn stayed flat."),
("What readers said", "Most readers told me the archive alone was worth it.")]
HABITS = [(None, "I write every morning before email.")]
@pytest.fixture()
def corpus_pair(tmp_path):
store = MemoryStore()
ws = uuid.uuid4()
writer = Corpus(ws, store=store, cache_dir=str(tmp_path / "worker"))
d1 = datetime(2026, 3, 1, tzinfo=UTC)
d2 = datetime(2025, 11, 9, tzinfo=UTC)
p1 = post_path("ada-substack-com", "On pricing", "https://ada.substack.com/p/on-pricing", d1)
p2 = post_path("ada-substack-com", "Habits", "https://ada.substack.com/p/habits", d2)
writer.write_files({
p1: render_post(
title="On pricing", url="https://x/p/on-pricing", published_at=d1, source="ada", sections=PRICING
),
p2: render_post(title="Habits", url="https://x/p/habits", published_at=d2, source="ada", sections=HABITS),
INDEX: "\n".join([index_line(p1, "On pricing", d1, ["Why I raised prices"]), index_line(p2, "Habits", d2, [])]),
})
# A different process (the API container) with an empty cache
reader = Corpus(ws, store=store, cache_dir=str(tmp_path / "api"))
reader.sync()
return reader, p1, p2
def test_paths_are_stable_and_readable():
p = post_path("ada-substack-com", "Any Title", "https://ada.substack.com/p/on-pricing/", datetime(2026, 3, 1))
assert p == "sources/ada-substack-com/2026-03-01-on-pricing.md"
def test_sync_mirrors_r2_into_a_fresh_cache(corpus_pair):
reader, p1, _ = corpus_pair
assert reader.exists(p1) and reader.exists(INDEX)
assert "# On pricing" in reader.read_lines(p1)
def test_tools_search_read_and_scope(corpus_pair):
reader, p1, p2 = corpus_pair
steps = []
tools = ScopedTools(reader, {p1: "On pricing"}, on_step=lambda k, d: steps.append(k))
listing = tools.list_files()
assert p1 in listing and p2 not in listing # notebook scope hides other posts
hits = tools.search(r"paid tier|price")
assert f"{p1}:" in hits and "$5 to $8" in hits
assert "Habits" not in hits
body = tools.read(p1, 1, 5)
assert "1| ---" in body
with pytest.raises(CorpusError):
tools.read(p2) # outside the notebook
with pytest.raises(CorpusError):
tools.read("../../etc/passwd")
assert tools.search("(unclosed") # bad regex falls back to a literal search
assert steps[:3] == ["list", "search", "read"]
def test_verify_citations_drops_fakes_and_renumbers(corpus_pair):
reader, p1, p2 = corpus_pair
tools = ScopedTools(reader, {p1: "On pricing"})
line = next(i for i, t in enumerate(reader.read_lines(p1), 1) if "$5 to $8" in t)
cites = [
FileCitation(marker=1, path=p2, line_start=1, line_end=3), # out of scope
FileCitation(marker=2, path=p1, line_start=line, line_end=line), # real
FileCitation(marker=3, path=p1, line_start=9999, line_end=10000), # past end of file
]
text, verified = verify_citations(tools, "Habits matter [1]. Prices went to $8 [2]. Churn fell [3].", cites)
assert [c.marker for c in verified] == [1]
assert "$5 to $8" in verified[0].quote
assert text == "Habits matter. Prices went to $8 [1]. Churn fell."
def test_react_agent_end_to_end_with_scripted_model(corpus_pair, monkeypatch):
"""Drives the real dspy.ReAct loop with a fake LM: search, read, then answer with a citation."""
import dspy
from dspy.utils.dummies import DummyLM
import app.pipeline.research as research_mod
reader, p1, _ = corpus_pair
line = next(i for i, t in enumerate(reader.read_lines(p1), 1) if "$5 to $8" in t)
lm = DummyLM([
{"next_thought": "Search for pricing.", "next_tool_name": "search", "next_tool_args": {"pattern": "paid tier"}},
{"next_thought": "Read the post.", "next_tool_name": "read",
"next_tool_args": {"path": p1, "start_line": 1, "end_line": 30}},
{"next_thought": "I have what I need.", "next_tool_name": "finish", "next_tool_args": {}},
{"reasoning": "The post states the new price.",
"answer": "You raised the paid tier from $5 to $8 a month [1].",
"citations": [{"marker": 1, "path": p1, "line_start": line, "line_end": line}],
"unsupported": False},
])
monkeypatch.setattr(research_mod, "main_lm", lambda: lm)
monkeypatch.setattr(research_mod, "track_usage", _no_usage)
monkeypatch.setattr(research_mod, "triage", lambda *a: research_mod.Triage("archive", "What did I charge?",
"deep", ""))
seen = []
with dspy.context(adapter=dspy.ChatAdapter()):
result = research_mod.research(
reader, {p1: "On pricing"}, "What did I charge?", on_step=lambda k, d: seen.append(k)
)
assert seen == ["think", "search", "read"]
assert result.unsupported is False
assert result.text.endswith("[1].")
assert result.citations[0].path == p1 and "$5 to $8" in result.citations[0].quote
class _Usage:
prompt_tokens = completion_tokens = 0
from contextlib import contextmanager # noqa: E402
@contextmanager
def _no_usage(lm):
yield _Usage()
def test_long_messages_keep_their_start_and_end():
from app.pipeline.research import HISTORY_CHARS, Turn, conversation_text, shorten
assert shorten("short and sweet") == "short and sweet"
text = "OPENING " + "middle words " * 200 + "CLOSING?"
out = shorten(text)
assert out.startswith("OPENING") and out.endswith("CLOSING?") and " ... " in out
assert "middle words" in out and len(out) < len(text) and len(out) <= HISTORY_CHARS + 5
assert shorten("x" * 5000).count("...") == 1 # no spaces to cut on: still bounded
assert len(shorten("x" * 5000)) <= HISTORY_CHARS + 5
long_answer = "First [1] " + "point " * 300 + "Want more?"
history = [Turn("user", "Compare my pricing posts"), Turn("assistant", long_answer)]
lines = conversation_text(history).splitlines()
assert lines[0] == "Writer: Compare my pricing posts" # short messages are untouched
assert lines[1].startswith("Assistant: First") and lines[1].endswith("Want more?") and "[1]" not in lines[1]
assert "..." not in lines[1] # the newest two messages are kept whole
def test_only_the_newest_two_messages_stay_whole():
from app.pipeline.research import HISTORY_CHARS, RECENT_CHARS, Turn, conversation_text
body = "start " + "filler " * 300 + "end"
history = [Turn("user" if i % 2 == 0 else "assistant", f"m{i} {body}") for i in range(6)]
lines = conversation_text(history).splitlines()
assert len(lines) == 6
assert all("..." in line and len(line) <= HISTORY_CHARS + 20 for line in lines[:4]) # older: head and tail
assert all("..." not in line and line.endswith("end") for line in lines[4:]) # newest two: whole
assert len(conversation_text(history, limit=3).splitlines()) == 3 # the window still applies
huge = [Turn("user", "q"), Turn("assistant", "word " * 5000)]
assert len(conversation_text(huge).splitlines()[1]) <= RECENT_CHARS + 30 # even the newest are bounded
|