import hashlib import json import sqlite3 import time MODE_INSTRUCTIONS = { "history": "Provide HISTORICAL CONTEXT: how this claim or approach emerged in the field, which earlier works set it up, and how thinking shifted over time.", "referee": ( "Act as a RIGOROUS REFEREE grounded in the provided PAPER CONTEXT and RETRIEVED LITERATURE. " "You MUST: (a) ground your critique in a SPECIFIC retrieved work [n] (comparing numbers where possible) " "OR a specific internal inconsistency with the provided paper context; " "(b) name EXACTLY ONE concrete check the authors should run or report; " "(c) NOT restate any caveat already present in the paragraph or its neighbors — " "if the paragraph survives scrutiny, say what evidence would strengthen it instead." ), "literature": "Offer OBSERVATIONS FROM THE RELATED LITERATURE: how the retrieved works support, contradict, or extend what this paragraph claims. Compare numbers where possible.", "explainer": "Write a PLAIN-LANGUAGE EXPLAINER for a beginning graduate student: unpack jargon and explain why this paragraph matters, without dumbing down.", "methods": "Provide a METHODS COMPARISON: contrast the technique in this paragraph with the approaches used in the retrieved works, noting trade-offs.", } _PREAMBLE = ( "You are an erudite astronomical annotator writing a marginal note beside a paragraph of a research paper, " "in the tradition of an annotated encyclopedia. Scholarly, precise, a little wry. " "Write 55–85 words, a single paragraph, no markdown, no preamble. " "Refer to the retrieved works by bracketed number, e.g. [1], where relevant.\n\n" ) def build_prompt( paragraph: str, section: str, mode: str, lit: list[dict], paper_context: dict | None = None, ) -> str: prompt = _PREAMBLE + MODE_INSTRUCTIONS[mode] + "\n\n" # Inject PAPER CONTEXT block for referee mode only if mode == "referee" and paper_context: outline = " → ".join(paper_context.get("sectionOutline") or []) prev = paper_context.get("prevParagraph", "") nxt = paper_context.get("nextParagraph", "") prompt += ( f"PAPER CONTEXT:\n" f"Title: {paper_context.get('title', '')}\n" f"Section outline: {outline}\n" f"Opening paragraph: {paper_context.get('opening', '')}\n" ) if prev: prompt += f"Previous paragraph: {prev}\n" if nxt: prompt += f"Next paragraph: {nxt}\n" prompt += "\n" prompt += f'PARAGRAPH (from section "{section}"):\n{paragraph}\n\n' if lit: block = "\n".join( f'[{i + 1}] {l["short"]} — "{l["title"]}" ({l["journal"]}). {l["abstract"]}' for i, l in enumerate(lit) ) prompt += f"RETRIEVED LITERATURE (via Pathfinder):\n{block}\n\n" return prompt + "Marginal note:" def cache_key( paragraph: str, mode: str, lit_ids: list[str], model: str, paper_context: dict | None = None, ) -> str: if mode == "referee" and paper_context: ctx_hash = hashlib.sha256( json.dumps(paper_context, sort_keys=True).encode() ).hexdigest()[:16] payload = json.dumps([paragraph, mode, sorted(lit_ids), model, ctx_hash]) else: payload = json.dumps([paragraph, mode, sorted(lit_ids), model]) return hashlib.sha256(payload.encode()).hexdigest() class CompletionCache: def __init__(self, path): self.conn = sqlite3.connect(str(path), check_same_thread=False) self.conn.execute( "CREATE TABLE IF NOT EXISTS cache (key TEXT PRIMARY KEY, text TEXT, created REAL)" ) def get(self, key: str) -> str | None: row = self.conn.execute("SELECT text FROM cache WHERE key = ?", (key,)).fetchone() return row[0] if row else None def put(self, key: str, text: str) -> None: self.conn.execute( "INSERT OR REPLACE INTO cache VALUES (?, ?, ?)", (key, text, time.time()) ) self.conn.commit() import httpx class ProviderError(Exception): pass def _parse_or_none(data: str): try: return json.loads(data) except json.JSONDecodeError: return None _client: httpx.AsyncClient | None = None def _http() -> httpx.AsyncClient: global _client if _client is None: _client = httpx.AsyncClient(timeout=httpx.Timeout(60.0, connect=10.0)) return _client async def _raise_provider_error(response: httpx.Response): body = (await response.aread()).decode(errors="replace") raise ProviderError(f"{response.status_code}: {body}"[:300]) def _sse_data_lines(response): async def gen(): async for line in response.aiter_lines(): if line.startswith("data: "): yield line[6:] return gen() async def stream_completion(provider: str, model: str, prompt: str, key: str): """Yields text chunks. Raises ProviderError on non-200 (message truncated to 300 chars). The user's key is used for the request and nothing else — never log it. """ if provider == "anthropic": req = _http().stream( "POST", "https://api.anthropic.com/v1/messages", headers={"x-api-key": key, "anthropic-version": "2023-06-01"}, json={"model": model, "max_tokens": 400, "stream": True, "messages": [{"role": "user", "content": prompt}]}, ) async with req as r: if r.status_code != 200: await _raise_provider_error(r) async for data in _sse_data_lines(r): ev = _parse_or_none(data) if ev is None: continue if ev.get("type") == "error": raise ProviderError(json.dumps(ev.get("error", {}))[:300]) if ev.get("type") == "content_block_delta": yield ev["delta"].get("text", "") elif provider == "openai": req = _http().stream( "POST", "https://api.openai.com/v1/chat/completions", headers={"Authorization": f"Bearer {key}"}, json={"model": model, "max_tokens": 400, "stream": True, "messages": [{"role": "user", "content": prompt}]}, ) async with req as r: if r.status_code != 200: await _raise_provider_error(r) async for data in _sse_data_lines(r): if data.strip() == "[DONE]": break ev = _parse_or_none(data) if ev is None: continue if "error" in ev: raise ProviderError(json.dumps(ev["error"])[:300]) choices = ev.get("choices") or [{}] delta = choices[0].get("delta", {}) if delta.get("content"): yield delta["content"] elif provider == "gemini": url = (f"https://generativelanguage.googleapis.com/v1beta/models/" f"{model}:streamGenerateContent?alt=sse&key={key}") req = _http().stream( "POST", url, json={"contents": [{"parts": [{"text": prompt}]}]}, ) async with req as r: if r.status_code != 200: await _raise_provider_error(r) async for data in _sse_data_lines(r): ev = _parse_or_none(data) if ev is None: continue for cand in ev.get("candidates", []): for part in cand.get("content", {}).get("parts", []): if part.get("text"): yield part["text"] else: raise ProviderError(f"unknown provider: {provider}")