Spaces:

kiyer
/

beacon

Sleeping

feat: in-text citations — visual styling, per-unit splitting, ADS search fallback

6ac1f78 4 months ago

9.04 kB

import json
import numpy as np
import pytest
from fastapi.testclient import TestClient
from pathlib import Path

import astroparse_api.main as main_mod
from astroparse_api.main import app
from astroparse_api.annotate import CompletionCache

FIXTURE = Path(__file__).parent / "fixtures" / "mowla_iyer_2024.pdf"


class FakeCorpus:
    class _DS:
        def search(self, col, q, k):
            class R:
                indices = [0]
                scores = np.array([0.5])
            return R()

        def __getitem__(self, idx):
            import datetime
            return {
                "title": ["T"],
                "abstract": ["A"],
                "authors": [["Iyer, K."]],
                "date": [datetime.date(2022, 1, 1)],
                "keywords": [["galaxy"]],
                "bibcode": ["2022ApJ...1I"],
                "arxiv_id": ["2201.00001"],
            }

    ds = _DS()

    @property
    def arxiv_ids(self) -> set:
        return {"2201.00001"}

    @property
    def bibcodes(self) -> set:
        return {"2022ApJ...1I"}

    def search(self, emb, k=1000):
        return [0], np.array([2.0])


@pytest.fixture
def client(monkeypatch, tmp_path):
    monkeypatch.setenv("ASTROPARSE_SKIP_CORPUS", "1")
    monkeypatch.setenv("ASTROPARSE_FIGURES_DIR", str(tmp_path / "figures"))
    monkeypatch.setattr(main_mod, "get_corpus", lambda: FakeCorpus())
    monkeypatch.setattr(
        main_mod,
        "embed_batch",
        lambda texts, api_key: [np.zeros(2, dtype=np.float32) for _ in texts],
    )
    main_mod.cache = CompletionCache(tmp_path / "cache.sqlite")
    with TestClient(app) as c:
        yield c


def _collect_sse(resp):
    events, cur = [], {}
    for line in resp.iter_lines():
        if line.startswith("event:"):
            cur["event"] = line.split(":", 1)[1].strip()
        elif line.startswith("data:"):
            cur["data"] = json.loads(line.split(":", 1)[1])
            events.append(cur)
            cur = {}
    return events


def test_parse_job_flow(client):
    r = client.post(
        "/api/parse",
        files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
    )
    assert r.status_code == 202
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    names = [e["event"] for e in events]
    assert names.count("result") == 1
    stage_events = [e for e in events if e["event"] == "stage"]
    assert {s["data"]["index"] for s in stage_events} == set(range(7))
    result = next(e for e in events if e["event"] == "result")["data"]
    assert result["paper"]["paragraphs"] and result["litByPara"]
    assert "pdfHash" in result
    assert "figures" in result
    assert "litByFig" in result


def test_annotate_cached_round_trip(client, monkeypatch):
    async def fake_stream(provider, model, prompt, key):
        for chunk in ["A fine ", "note."]:
            yield chunk

    monkeypatch.setattr(main_mod, "stream_completion", fake_stream)
    body = {
        "paragraphId": "p1",
        "paragraph": "text",
        "section": "S",
        "mode": "referee",
        "lit": [],
        "provider": "anthropic",
        "model": "claude-haiku-4-5",
        "key": "sk-x",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    done = next(e for e in events if e["event"] == "done")["data"]
    assert done["text"] == "A fine note." and done["cached"] is False

    with client.stream("POST", "/api/annotate", json=body) as resp:  # second call: cache hit
        events2 = _collect_sse(resp)
    done2 = next(e for e in events2 if e["event"] == "done")["data"]
    assert done2["cached"] is True and done2["text"] == "A fine note."


def test_annotate_missing_key(client):
    body = {
        "paragraphId": "p1",
        "paragraph": "t",
        "section": "S",
        "mode": "referee",
        "lit": [],
        "provider": "anthropic",
        "model": "m",
        "key": "",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    assert events[-1]["event"] == "error"
    assert "no API key set" in events[-1]["data"]["message"]


def test_figures_endpoint_serves_png(client, tmp_path, monkeypatch):
    """After a parse, the figure PNG can be fetched from the figures endpoint."""
    # Run a parse to generate figures
    r = client.post(
        "/api/parse",
        files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
    )
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    result = next(e for e in events if e["event"] == "result")["data"]
    pdf_hash = result["pdfHash"]
    figures = result["figures"]

    # Find a figure with an image
    with_image = [f for f in figures if f["hasImage"]]
    assert with_image, "Expected at least one figure with an image"
    fig_id = with_image[0]["id"]

    resp = client.get(f"/api/figures/{pdf_hash}/{fig_id}.png")
    assert resp.status_code == 200
    assert resp.headers["content-type"] == "image/png"
    assert len(resp.content) > 5_000


def test_figures_endpoint_invalid_hash(client):
    """An invalid (non-hex-64) hash returns 404."""
    r = client.get("/api/figures/invalid_hash/f1.png")
    assert r.status_code == 404


def test_figures_endpoint_invalid_fig_id(client):
    """An invalid fig_id like 'fx1' returns 404."""
    good_hash = "a" * 64
    r = client.get(f"/api/figures/{good_hash}/fx1.png")
    assert r.status_code == 404


def test_figures_endpoint_traversal_attempt(client):
    """Traversal payloads in either path segment are rejected by validation."""
    good_hash = "a" * 64
    # traversal in the hash segment (encoded so the path reaches the route)
    r = client.get(f"/api/figures/..%2f..%2f{good_hash[:58]}/f1.png")
    assert r.status_code == 404
    # traversal in the fig_id segment
    r = client.get(f"/api/figures/{good_hash}/..%2f..%2fsecret.png")
    assert r.status_code == 404
    # valid hash + valid fig_id but no file on disk → also 404 (no info leak)
    r = client.get(f"/api/figures/{good_hash}/f1.png")
    assert r.status_code == 404


def test_jobs_stage_count():
    """STAGES list must have exactly 7 entries."""
    from astroparse_api.jobs import STAGES
    assert len(STAGES) == 7
    assert STAGES[6] == "Tipping in the plates"


# ---------------------------------------------------------------------------
# Feature: plaintext/markdown manuscript input
# ---------------------------------------------------------------------------

_MD_UPLOAD = (
    "# Test Manuscript\n\n"
    "## Introduction\n\n"
    + ("Galaxy formation and evolution in the early universe is a key topic. " * 20)
    + "\n\n"
    + ("Star formation rates and quenching mechanisms drive cosmic evolution. " * 20)
    + "\n\n"
    "## References\n\n"
    "[1] Someone et al. 2020\n"
).encode()


def test_markdown_upload_all_stages_emitted(client):
    """Uploading a .md file triggers all 7 stages and returns figures=[]."""
    r = client.post(
        "/api/parse",
        files={"file": ("notes.md", _MD_UPLOAD, "text/markdown")},
    )
    assert r.status_code == 202
    job_id = r.json()["jobId"]
    with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
        events = _collect_sse(resp)
    # All 7 stages must be emitted
    stage_events = [e for e in events if e["event"] == "stage"]
    assert {s["data"]["index"] for s in stage_events} == set(range(7)), (
        f"Missing stages: got indices {[s['data']['index'] for s in stage_events]}"
    )
    # Result must have paragraphs, figures=[], litByPara set
    result = next(e for e in events if e["event"] == "result")["data"]
    assert result["paper"]["paragraphs"], "Expected paragraphs in result"
    assert result["figures"] == [], f"Expected figures=[], got {result['figures']}"
    assert result["litByPara"], "Expected litByPara non-empty"
    assert result["litByFig"] == {}, f"Expected litByFig={{}}, got {result['litByFig']}"


def test_annotate_network_error_sanitized(client, monkeypatch):
    """Non-ProviderError failures (e.g. httpx ConnectError) emit a sanitized
    error event — never the raw exception, which can carry a key-bearing URL."""
    async def broken_stream(provider, model, prompt, key):
        raise RuntimeError(f"https://example/?key={key}")
        yield  # pragma: no cover — makes this an async generator
    monkeypatch.setattr(main_mod, "stream_completion", broken_stream)
    body = {
        "paragraphId": "p1", "paragraph": "t", "section": "S", "mode": "referee",
        "lit": [], "provider": "gemini", "model": "m", "key": "SECRETKEY",
    }
    with client.stream("POST", "/api/annotate", json=body) as resp:
        events = _collect_sse(resp)
    assert events[-1]["event"] == "error"
    msg = events[-1]["data"]["message"]
    assert "SECRETKEY" not in msg and "could not reach gemini" in msg

	import json
	import numpy as np
	import pytest
	from fastapi.testclient import TestClient
	from pathlib import Path

	import astroparse_api.main as main_mod
	from astroparse_api.main import app
	from astroparse_api.annotate import CompletionCache

	FIXTURE = Path(__file__).parent / "fixtures" / "mowla_iyer_2024.pdf"


	class FakeCorpus:
	class _DS:
	def search(self, col, q, k):
	class R:
	indices = [0]
	scores = np.array([0.5])
	return R()

	def __getitem__(self, idx):
	import datetime
	return {
	"title": ["T"],
	"abstract": ["A"],
	"authors": [["Iyer, K."]],
	"date": [datetime.date(2022, 1, 1)],
	"keywords": [["galaxy"]],
	"bibcode": ["2022ApJ...1I"],
	"arxiv_id": ["2201.00001"],
	}

	ds = _DS()

	@property
	def arxiv_ids(self) -> set:
	return {"2201.00001"}

	@property
	def bibcodes(self) -> set:
	return {"2022ApJ...1I"}

	def search(self, emb, k=1000):
	return [0], np.array([2.0])


	@pytest.fixture
	def client(monkeypatch, tmp_path):
	monkeypatch.setenv("ASTROPARSE_SKIP_CORPUS", "1")
	monkeypatch.setenv("ASTROPARSE_FIGURES_DIR", str(tmp_path / "figures"))
	monkeypatch.setattr(main_mod, "get_corpus", lambda: FakeCorpus())
	monkeypatch.setattr(
	main_mod,
	"embed_batch",
	lambda texts, api_key: [np.zeros(2, dtype=np.float32) for _ in texts],
	)
	main_mod.cache = CompletionCache(tmp_path / "cache.sqlite")
	with TestClient(app) as c:
	yield c


	def _collect_sse(resp):
	events, cur = [], {}
	for line in resp.iter_lines():
	if line.startswith("event:"):
	cur["event"] = line.split(":", 1)[1].strip()
	elif line.startswith("data:"):
	cur["data"] = json.loads(line.split(":", 1)[1])
	events.append(cur)
	cur = {}
	return events


	def test_parse_job_flow(client):
	r = client.post(
	"/api/parse",
	files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
	)
	assert r.status_code == 202
	job_id = r.json()["jobId"]
	with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
	events = _collect_sse(resp)
	names = [e["event"] for e in events]
	assert names.count("result") == 1
	stage_events = [e for e in events if e["event"] == "stage"]
	assert {s["data"]["index"] for s in stage_events} == set(range(7))
	result = next(e for e in events if e["event"] == "result")["data"]
	assert result["paper"]["paragraphs"] and result["litByPara"]
	assert "pdfHash" in result
	assert "figures" in result
	assert "litByFig" in result


	def test_annotate_cached_round_trip(client, monkeypatch):
	async def fake_stream(provider, model, prompt, key):
	for chunk in ["A fine ", "note."]:
	yield chunk

	monkeypatch.setattr(main_mod, "stream_completion", fake_stream)
	body = {
	"paragraphId": "p1",
	"paragraph": "text",
	"section": "S",
	"mode": "referee",
	"lit": [],
	"provider": "anthropic",
	"model": "claude-haiku-4-5",
	"key": "sk-x",
	}
	with client.stream("POST", "/api/annotate", json=body) as resp:
	events = _collect_sse(resp)
	done = next(e for e in events if e["event"] == "done")["data"]
	assert done["text"] == "A fine note." and done["cached"] is False

	with client.stream("POST", "/api/annotate", json=body) as resp: # second call: cache hit
	events2 = _collect_sse(resp)
	done2 = next(e for e in events2 if e["event"] == "done")["data"]
	assert done2["cached"] is True and done2["text"] == "A fine note."


	def test_annotate_missing_key(client):
	body = {
	"paragraphId": "p1",
	"paragraph": "t",
	"section": "S",
	"mode": "referee",
	"lit": [],
	"provider": "anthropic",
	"model": "m",
	"key": "",
	}
	with client.stream("POST", "/api/annotate", json=body) as resp:
	events = _collect_sse(resp)
	assert events[-1]["event"] == "error"
	assert "no API key set" in events[-1]["data"]["message"]


	def test_figures_endpoint_serves_png(client, tmp_path, monkeypatch):
	"""After a parse, the figure PNG can be fetched from the figures endpoint."""
	# Run a parse to generate figures
	r = client.post(
	"/api/parse",
	files={"file": ("m.pdf", FIXTURE.read_bytes(), "application/pdf")},
	)
	job_id = r.json()["jobId"]
	with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
	events = _collect_sse(resp)
	result = next(e for e in events if e["event"] == "result")["data"]
	pdf_hash = result["pdfHash"]
	figures = result["figures"]

	# Find a figure with an image
	with_image = [f for f in figures if f["hasImage"]]
	assert with_image, "Expected at least one figure with an image"
	fig_id = with_image[0]["id"]

	resp = client.get(f"/api/figures/{pdf_hash}/{fig_id}.png")
	assert resp.status_code == 200
	assert resp.headers["content-type"] == "image/png"
	assert len(resp.content) > 5_000


	def test_figures_endpoint_invalid_hash(client):
	"""An invalid (non-hex-64) hash returns 404."""
	r = client.get("/api/figures/invalid_hash/f1.png")
	assert r.status_code == 404


	def test_figures_endpoint_invalid_fig_id(client):
	"""An invalid fig_id like 'fx1' returns 404."""
	good_hash = "a" * 64
	r = client.get(f"/api/figures/{good_hash}/fx1.png")
	assert r.status_code == 404


	def test_figures_endpoint_traversal_attempt(client):
	"""Traversal payloads in either path segment are rejected by validation."""
	good_hash = "a" * 64
	# traversal in the hash segment (encoded so the path reaches the route)
	r = client.get(f"/api/figures/..%2f..%2f{good_hash[:58]}/f1.png")
	assert r.status_code == 404
	# traversal in the fig_id segment
	r = client.get(f"/api/figures/{good_hash}/..%2f..%2fsecret.png")
	assert r.status_code == 404
	# valid hash + valid fig_id but no file on disk → also 404 (no info leak)
	r = client.get(f"/api/figures/{good_hash}/f1.png")
	assert r.status_code == 404


	def test_jobs_stage_count():
	"""STAGES list must have exactly 7 entries."""
	from astroparse_api.jobs import STAGES
	assert len(STAGES) == 7
	assert STAGES[6] == "Tipping in the plates"


	# ---------------------------------------------------------------------------
	# Feature: plaintext/markdown manuscript input
	# ---------------------------------------------------------------------------

	_MD_UPLOAD = (
	"# Test Manuscript\n\n"
	"## Introduction\n\n"
	+ ("Galaxy formation and evolution in the early universe is a key topic. " * 20)
	+ "\n\n"
	+ ("Star formation rates and quenching mechanisms drive cosmic evolution. " * 20)
	+ "\n\n"
	"## References\n\n"
	"[1] Someone et al. 2020\n"
	).encode()


	def test_markdown_upload_all_stages_emitted(client):
	"""Uploading a .md file triggers all 7 stages and returns figures=[]."""
	r = client.post(
	"/api/parse",
	files={"file": ("notes.md", _MD_UPLOAD, "text/markdown")},
	)
	assert r.status_code == 202
	job_id = r.json()["jobId"]
	with client.stream("GET", f"/api/jobs/{job_id}/events") as resp:
	events = _collect_sse(resp)
	# All 7 stages must be emitted
	stage_events = [e for e in events if e["event"] == "stage"]
	assert {s["data"]["index"] for s in stage_events} == set(range(7)), (
	f"Missing stages: got indices {[s['data']['index'] for s in stage_events]}"
	)
	# Result must have paragraphs, figures=[], litByPara set
	result = next(e for e in events if e["event"] == "result")["data"]
	assert result["paper"]["paragraphs"], "Expected paragraphs in result"
	assert result["figures"] == [], f"Expected figures=[], got {result['figures']}"
	assert result["litByPara"], "Expected litByPara non-empty"
	assert result["litByFig"] == {}, f"Expected litByFig={{}}, got {result['litByFig']}"


	def test_annotate_network_error_sanitized(client, monkeypatch):
	"""Non-ProviderError failures (e.g. httpx ConnectError) emit a sanitized
	error event — never the raw exception, which can carry a key-bearing URL."""
	async def broken_stream(provider, model, prompt, key):
	raise RuntimeError(f"https://example/?key={key}")
	yield # pragma: no cover — makes this an async generator
	monkeypatch.setattr(main_mod, "stream_completion", broken_stream)
	body = {
	"paragraphId": "p1", "paragraph": "t", "section": "S", "mode": "referee",
	"lit": [], "provider": "gemini", "model": "m", "key": "SECRETKEY",
	}
	with client.stream("POST", "/api/annotate", json=body) as resp:
	events = _collect_sse(resp)
	assert events[-1]["event"] == "error"
	msg = events[-1]["data"]["message"]
	assert "SECRETKEY" not in msg and "could not reach gemini" in msg