Spaces:
Runtime error
Runtime error
| import os | |
| import pytest | |
| from src import agents, config, pipeline | |
| from src.documents import load_documents | |
| from src.evidence import EvidenceStore | |
| from tests.conftest import WEB, fake_searcher | |
| def run(fake_llm, embedder, **kwargs): | |
| events = list(pipeline.run_research( | |
| "AI in agriculture", searcher=fake_searcher, embedder=embedder, llm_api=fake_llm, **kwargs | |
| )) | |
| return events, events[-1] | |
| # ------------------------------------------------------------ agents | |
| def test_parse_json_handles_fences_and_prose(): | |
| assert agents.parse_json('```json\n{"a": 1}\n```') == {"a": 1} | |
| assert agents.parse_json('Sure! {"a": 1} hope that helps') == {"a": 1} | |
| assert agents.parse_json("not json") is None | |
| def test_planner_falls_back_to_topic_on_garbage(fake_llm): | |
| assert agents.plan("my topic", 3, lambda *a, **k: "garbage") == ["my topic"] | |
| def test_planner_limits_and_dedupes(): | |
| reply = '{"sub_questions": ["a", "a", "b", "c", "d"]}' | |
| assert agents.plan("t", 3, lambda *a, **k: reply) == ["a", "b", "c"] | |
| def test_clean_citations(): | |
| assert agents.clean_citations("A [1] B [9] C [2, 7] D [2026]", 3) == "A [1] B C [2] D [2026]" | |
| assert agents.clean_citations("x [1, 2]", 5) == "x [1][2]" | |
| def test_extract_claims_keeps_citation_with_its_sentence(): | |
| report = ( | |
| "## Heading\n" | |
| "- Solar panels share land with crops. [1] Drones detect disease early [2].\n" | |
| "This sentence has plenty of words but no citation at all in it.\n" | |
| "> quoted verification note [1]\n" | |
| ) | |
| claims, uncited = agents.extract_claims(report) | |
| assert [c["cites"] for c in claims] == [[1], [2]] | |
| assert claims[0]["text"].startswith("Solar panels share land with crops") | |
| assert uncited == 1 | |
| def test_research_web_dedupes_urls_and_collects_errors(): | |
| from src.search import SearchError | |
| def searcher(q, max_results, search_depth): | |
| if q == "bad": | |
| raise SearchError("boom") | |
| return [{"title": "T", "url": "https://x", "content": "text"}] | |
| sources, errors = agents.research_web(["a", "b", "bad"], {"results_per_question": 2, "search_depth": "basic"}, searcher) | |
| assert len(sources) == 1 and errors == ["boom"] | |
| # ---------------------------------------------------------- evidence | |
| def test_evidence_ranks_relevant_sources_first_and_drops_irrelevant(embedder): | |
| sources = [dict(s, kind="web") for s in WEB["solar"] + WEB["drone"]] | |
| store = EvidenceStore(sources, embedder) | |
| selected, numbered = store.gather(["solar panels on farmland", "drone crop monitoring"], 5) | |
| titles = [s["title"] for s in numbered] | |
| assert titles[0] == "Agrivoltaics explained" | |
| assert "Cooking pasta" not in titles | |
| assert len(selected) <= config.MAX_EVIDENCE_CHUNKS | |
| assert all(sum(1 for c in selected if c["source_idx"] == i) <= config.MAX_CHUNKS_PER_SOURCE for i in range(3)) | |
| def test_support_for_returns_passages_from_cited_source_only(embedder): | |
| sources = [dict(s, kind="web") for s in WEB["solar"] + WEB["drone"]] | |
| store = EvidenceStore(sources, embedder) | |
| store.gather(["solar panels on farmland", "drone crop monitoring"], 5) | |
| passages = store.support_for("drones detect plant disease", [2], k=2) | |
| assert passages and all("Drones" in p for p in passages) | |
| assert store.support_for("anything", [99]) == [] | |
| # ----------------------------------------------------------- pipeline | |
| def test_full_agent_loop_catches_and_fixes_unsupported_claim(fake_llm, embedder): | |
| events, last = run(fake_llm, embedder) | |
| assert last["done"] and not last["error"] | |
| result = last["result"] | |
| # Checker flagged the planted claim, Reviser removed it, re-check is clean | |
| m = result["metrics"] | |
| assert m["first_pass"]["unsupported"] == 1 | |
| assert m["final"]["unsupported"] == 0 and m["final"]["support_rate"] == 1.0 | |
| assert m["revised"] is True | |
| assert "cheese" not in result["report"] | |
| # Invented citation [9] removed, sources + verification note appended | |
| assert "[9]" not in result["report"] | |
| assert "## Sources" in result["report"] and "https://a.example/agri" in result["report"] | |
| assert "Verification" in result["report"] | |
| # Trace shows every agent | |
| trace = last["trace"] | |
| for agent in ("Planner", "Researcher", "Retriever", "Writer", "Checker", "Reviser"): | |
| assert agent in trace | |
| assert "Translator" not in trace | |
| def test_draft_streams_progressively(fake_llm, embedder): | |
| events, _ = run(fake_llm, embedder) | |
| drafts = [e["report"] for e in events if e["report"] and e["report"].startswith("*Draft")] | |
| assert len(drafts) >= 3 | |
| def test_urdu_runs_translator_after_verification(fake_llm, embedder): | |
| _, last = run(fake_llm, embedder, language="Urdu") | |
| result = last["result"] | |
| assert result["report"].startswith("[UR]") | |
| assert "cheese" not in result["report_english"] | |
| assert "Translator" in last["trace"] | |
| def test_documents_only_mode(fake_llm, embedder, tmp_path): | |
| doc = tmp_path / "notes.txt" | |
| doc.write_text("Greenhouse farming with solar panels on farmland increases crop yield per hectare. " * 10) | |
| _, last = run(fake_llm, embedder, use_web=False, document_paths=[str(doc)]) | |
| assert last["done"] | |
| sources = last["result"]["sources"] | |
| assert sources and all(s["kind"] == "document" for s in sources) | |
| assert "📄" in last["result"]["report"] | |
| def test_no_web_and_no_documents_is_a_clear_error(fake_llm, embedder): | |
| _, last = run(fake_llm, embedder, use_web=False) | |
| assert last["error"] and "nothing to research" in last["report"] | |
| def test_empty_topic_is_rejected(fake_llm, embedder): | |
| events = list(pipeline.run_research(" ", searcher=fake_searcher, embedder=embedder, llm_api=fake_llm)) | |
| assert events[-1]["error"] | |
| def test_search_failure_without_documents_surfaces_message(fake_llm, embedder): | |
| from src.search import SearchError | |
| def broken(q, max_results=4, search_depth="basic"): | |
| raise SearchError("`TAVILY_API_KEY` is missing") | |
| events = list(pipeline.run_research("topic here", searcher=broken, embedder=embedder, llm_api=fake_llm)) | |
| assert events[-1]["error"] and "TAVILY_API_KEY" in events[-1]["report"] | |
| def test_writer_failure_is_reported_not_raised(embedder, fake_llm): | |
| import types | |
| from src.llm import LLMError | |
| def bad_stream(*a, **k): | |
| raise LLMError("All Groq models failed.") | |
| yield | |
| broken = types.SimpleNamespace(complete=fake_llm.complete, stream=bad_stream) | |
| events = list(pipeline.run_research("topic here", searcher=fake_searcher, embedder=embedder, llm_api=broken)) | |
| assert events[-1]["error"] and "All Groq models failed" in events[-1]["report"] | |
| def test_checker_outage_still_delivers_report(embedder, fake_llm): | |
| import types | |
| from src.llm import LLMError | |
| def complete(messages, role="fast", temperature=0.0, json_mode=False): | |
| if "Planner" in messages[0]["content"]: | |
| return fake_llm.complete(messages, role, temperature, json_mode) | |
| raise LLMError("rate limited") | |
| llm_api = types.SimpleNamespace(complete=complete, stream=fake_llm.stream) | |
| _, last = None, list(pipeline.run_research("topic here", searcher=fake_searcher, embedder=embedder, llm_api=llm_api))[-1] | |
| assert last["done"] and "unavailable" in last["result"]["report"] | |
| def test_reviser_guard_rejects_stub_rewrites(embedder): | |
| report = "word " * 200 | |
| claims = [{"id": 1, "text": "x [1]", "cites": [1]}] | |
| results = {1: {"verdict": "unsupported", "reason": "r"}} | |
| class Store: | |
| def support_for(self, *a, **k): | |
| return ["evidence"] | |
| assert agents.revise(report, claims, results, Store(), lambda *a, **k: "ok") == report | |
| # ---------------------------------------------------------- documents | |
| def test_load_documents_pdf_txt_and_warnings(tmp_path): | |
| from reportlab.lib.pagesizes import LETTER | |
| from reportlab.platypus import Paragraph, SimpleDocTemplate | |
| from reportlab.lib.styles import getSampleStyleSheet | |
| pdf = tmp_path / "paper.pdf" | |
| SimpleDocTemplate(str(pdf), pagesize=LETTER).build( | |
| [Paragraph("Agrivoltaics lets farmers grow crops under solar panels and earn extra income. " * 5, getSampleStyleSheet()["BodyText"])] | |
| ) | |
| txt = tmp_path / "empty.txt" | |
| txt.write_text("hi") | |
| odd = tmp_path / "data.xyz" | |
| odd.write_text("whatever " * 30) | |
| docs, warnings = load_documents([str(pdf), str(txt), str(odd)]) | |
| assert [d["title"] for d in docs] == ["paper.pdf"] | |
| assert "Agrivoltaics" in docs[0]["content"] and docs[0]["kind"] == "document" | |
| assert len(warnings) == 2 | |
| # ---------------------------------------------------------- evaluation | |
| def test_evaluate_aggregates_and_writes_reports(fake_llm, embedder, tmp_path): | |
| import functools | |
| import evaluate | |
| runner = functools.partial(pipeline.run_research, searcher=fake_searcher, embedder=embedder, llm_api=fake_llm) | |
| rows = [evaluate.evaluate_topic("AI in agriculture", "Quick", runner)] | |
| rows.append(evaluate.evaluate_topic("x", "Quick", runner)) # too short -> failure row | |
| assert rows[0]["ok"] and rows[0]["support_first"] < rows[0]["support_final"] == 1.0 | |
| assert not rows[1]["ok"] | |
| summary = evaluate.aggregate(rows) | |
| assert summary["completed"] == 1 and summary["revised_count"] == 1 | |
| evaluate.write_outputs(rows, summary, str(tmp_path), "Quick") | |
| text = (tmp_path / "RESULTS.md").read_text(encoding="utf-8") | |
| assert "first draft" in text and "100%" in text and (tmp_path / "results.json").exists() | |