| 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" |
|
|
| |
| 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}") |
|
|
