summarize.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607
  1. """Generate summaries using DeepSeek; secrets stay in memory, resumable local cache."""
  2. import argparse
  3. import hashlib
  4. from http.client import HTTPConnection, HTTPSConnection, HTTPException
  5. from io import BytesIO
  6. import json
  7. import os
  8. import re
  9. import sqlite3
  10. import tempfile
  11. import time
  12. import urllib.error
  13. import urllib.parse
  14. import urllib.request
  15. from collections import Counter
  16. from concurrent.futures import ThreadPoolExecutor, as_completed
  17. from contextlib import contextmanager
  18. from dataclasses import dataclass, field, replace
  19. from datetime import datetime, timezone
  20. from pathlib import Path
  21. from .prepare import ROOT, dump, strict_json
  22. from .summary_input import PROMPT_VERSION, chat_text, facts, messages, scenario
  23. from src.proxy_auth import ProxyAuthError, proxy_tokens, token_rejected
  24. KINDS = ("policy", "finance", "general", "qa")
  25. DEFAULTS = {"SUMMARY_CONTEXT_TOKENS": "32768", "SUMMARY_MAX_OUTPUT_TOKENS": "4096",
  26. "SUMMARY_TOKEN_MARGIN": "512", "SUMMARY_CONCURRENCY": "8",
  27. "SUMMARY_TIMEOUT_SECONDS": "120", "SUMMARY_MAX_ATTEMPTS": "3",
  28. "SUMMARY_MAX_OUTPUT_LIMIT": "32768", "SUMMARY_MAX_CONTEXT_LIMIT": "131072"}
  29. CONFIG_KEYS = set(DEFAULTS) | {"DEEPSEEK_API_KEY", "DEEPSEEK_BASE_URL", "DEEPSEEK_MODEL"}
  30. _JSON_UNREAD = object()
  31. def fingerprint(value):
  32. return hashlib.sha256(dump(value).encode("utf-8")).hexdigest()
  33. def utc_now():
  34. return datetime.now(timezone.utc).isoformat(timespec="seconds")
  35. @dataclass
  36. class Config:
  37. api_key: str = field(repr=False)
  38. base_url: str = field(repr=False)
  39. model: str
  40. context_tokens: int = 32768
  41. max_output_tokens: int = 4096
  42. margin: int = 512
  43. concurrency: int = 8
  44. timeout: int = 120
  45. max_attempts: int = 3
  46. max_output_limit: int = 32768
  47. max_context_limit: int = 131072
  48. proxy_auth: object = field(default=None, repr=False, compare=False)
  49. @property
  50. def endpoint(self):
  51. return chat_endpoint(self.base_url)
  52. def chat_endpoint(base_url):
  53. """Resolve a base or complete chat URL without rewriting signed query values."""
  54. try:
  55. if not isinstance(base_url, str) or any(c.isspace() for c in base_url) or "\\" in base_url:
  56. raise ValueError
  57. parts = urllib.parse.urlsplit(base_url)
  58. if (parts.scheme not in {"http", "https"} or not parts.hostname
  59. or parts.username is not None or parts.password is not None or parts.fragment):
  60. raise ValueError
  61. _ = parts.port
  62. except ValueError:
  63. raise ValueError("invalid_deepseek_base_url") from None
  64. path = parts.path.rstrip("/")
  65. if not path.endswith("/chat/completions"):
  66. path += "/chat/completions"
  67. return urllib.parse.urlunsplit(parts._replace(path=path))
  68. class _BufferedResponse(BytesIO):
  69. def __init__(self, body, decoded):
  70. super().__init__(body)
  71. self.decoded, self.length = decoded, 0
  72. class _PrefixedResponse:
  73. """Replay a bounded preflight read without changing JSON or SSE bytes."""
  74. def __init__(self, prefix, response):
  75. self.prefix, self.offset, self.response = prefix, 0, response
  76. def __getattr__(self, name):
  77. return getattr(self.response, name)
  78. @property
  79. def length(self):
  80. remaining = getattr(self.response, 'length', None)
  81. return None if remaining is None else remaining + len(self.prefix) - self.offset
  82. def read1(self, count=-1):
  83. if self.offset < len(self.prefix):
  84. end = len(self.prefix) if count < 0 else min(len(self.prefix), self.offset + count)
  85. result = self.prefix[self.offset:end]
  86. self.offset = end
  87. return result
  88. return self.response.read1(count)
  89. def read(self, count=-1):
  90. first = self.read1(count) if self.offset < len(self.prefix) else b''
  91. return first + self.response.read(-1 if count < 0 else count - len(first))
  92. @contextmanager
  93. def deepseek_response(payload, config, *, stream=False, check=lambda: None):
  94. """One request, with a private connection per gateway call; always close it."""
  95. configured = config.endpoint
  96. auth = getattr(config, 'proxy_auth', None)
  97. endpoint = auth.url(configured) if auth is not None else configured
  98. parts = urllib.parse.urlsplit(endpoint)
  99. query = dict(urllib.parse.parse_qsl(parts.query))
  100. gateway = query.get('servertype') == 'deepseek' and bool(query.get('proxyToken'))
  101. headers = {'Authorization': 'Bearer ' + config.api_key, 'Content-Type': 'application/json',
  102. 'Accept': '*/*' if gateway else ('text/event-stream' if stream else 'application/json')}
  103. body = dump(payload).encode('utf-8')
  104. if not gateway:
  105. request = urllib.request.Request(endpoint, data=body, headers=headers)
  106. with urllib.request.build_opener(NoRedirect()).open(request, timeout=config.timeout) as response:
  107. yield response
  108. return
  109. headers['Connection'] = 'keep-alive'
  110. for attempt in range(2):
  111. check()
  112. parts = urllib.parse.urlsplit(endpoint)
  113. connection = (HTTPSConnection if parts.scheme == 'https' else HTTPConnection)(
  114. parts.hostname, parts.port, timeout=config.timeout)
  115. rejected = False
  116. try:
  117. target = (parts.path or '/') + ('?' + parts.query if parts.query else '')
  118. connection.request('POST', target, body=body, headers=headers)
  119. with connection.getresponse() as response:
  120. buffered, decoded = None, _JSON_UNREAD
  121. if not 200 <= response.status < 300 and response.status not in (401, 403):
  122. code = response.status
  123. raise RequestFailure(f'http_{code}', not stream and (code in {408, 429} or code >= 500))
  124. # read1 validates HTTP chunk framing (peek can swallow broken
  125. # chunks). Replay all inspected bytes for normal SSE streams.
  126. check()
  127. prefix = bytearray(response.read1(4096))
  128. deadline = time.monotonic() + config.timeout
  129. while prefix and not prefix.lstrip() and len(prefix) <= 4_000_000:
  130. check()
  131. if time.monotonic() >= deadline:
  132. raise RequestFailure('request_timeout', False)
  133. more = response.read1(4096)
  134. if not more:
  135. break
  136. prefix.extend(more)
  137. if len(prefix) > 4_000_000:
  138. raise RequestFailure('response_too_large', False)
  139. response = _PrefixedResponse(bytes(prefix), response)
  140. if prefix.lstrip().startswith(b'{'):
  141. buffered = response.read(4_000_001)
  142. if len(buffered) > 4_000_000:
  143. raise RequestFailure('response_too_large', False)
  144. if getattr(response, 'length', None) not in (None, 0):
  145. raise RequestFailure('response_incomplete', False)
  146. try:
  147. decoded = strict_json(buffered.decode('utf-8'))
  148. rejected = token_rejected(decoded)
  149. except (ValueError, UnicodeError):
  150. pass
  151. if not rejected:
  152. if not 200 <= response.status < 300:
  153. code = response.status
  154. raise RequestFailure(f'http_{code}', not stream and (code in {408, 429} or code >= 500))
  155. if buffered is None:
  156. yield response
  157. else:
  158. with _BufferedResponse(buffered, decoded) as cached:
  159. yield cached
  160. return
  161. finally:
  162. connection.close()
  163. if rejected:
  164. if attempt or auth is None:
  165. if auth is not None:
  166. auth.reject(endpoint)
  167. raise RequestFailure('proxy_token_invalid', False)
  168. check()
  169. try:
  170. auth.refresh(endpoint)
  171. except ProxyAuthError:
  172. raise RequestFailure('proxy_token_refresh_failed', False) from None
  173. endpoint = auth.url(configured)
  174. def load_config(path, environ=None):
  175. values = dict(DEFAULTS)
  176. for line in path.read_text(encoding="utf-8-sig").splitlines():
  177. line = line.strip()
  178. if line.startswith("export "):
  179. line = line[7:]
  180. if not line or line.startswith("#") or "=" not in line:
  181. continue
  182. key, value = line.split("=", 1)
  183. key, value = key.strip(), value.strip()
  184. if key not in CONFIG_KEYS:
  185. continue
  186. if value.startswith(("'", '"')):
  187. quote = value[0]
  188. end = value.find(quote, 1)
  189. if end < 0 or (value[end + 1:].strip() and not value[end + 1:].strip().startswith("#")):
  190. raise ValueError("invalid_env_quoting")
  191. value = value[1:end]
  192. else:
  193. value = value.split(" #", 1)[0].strip()
  194. values[key] = value
  195. env = os.environ if environ is None else environ
  196. values.update({key: env[key] for key in CONFIG_KEYS if key in env})
  197. if not all(values.get(key) for key in ("DEEPSEEK_API_KEY", "DEEPSEEK_BASE_URL", "DEEPSEEK_MODEL")):
  198. raise ValueError("missing_deepseek_configuration")
  199. chat_endpoint(values["DEEPSEEK_BASE_URL"])
  200. try:
  201. numbers = {key: int(values[key]) for key in DEFAULTS}
  202. except ValueError:
  203. raise ValueError("invalid_summary_integer_setting") from None
  204. if any(value <= 0 for value in numbers.values()) or numbers["SUMMARY_CONCURRENCY"] > 32 or numbers["SUMMARY_MAX_ATTEMPTS"] > 5:
  205. raise ValueError("summary_setting_out_of_range")
  206. if numbers["SUMMARY_MAX_OUTPUT_TOKENS"] + numbers["SUMMARY_TOKEN_MARGIN"] >= numbers["SUMMARY_CONTEXT_TOKENS"]:
  207. raise ValueError("output_budget_leaves_no_input_space")
  208. if numbers["SUMMARY_MAX_OUTPUT_TOKENS"] > numbers["SUMMARY_MAX_OUTPUT_LIMIT"] or numbers["SUMMARY_CONTEXT_TOKENS"] > numbers["SUMMARY_MAX_CONTEXT_LIMIT"]:
  209. raise ValueError("initial_budget_exceeds_retry_limit")
  210. return Config(values["DEEPSEEK_API_KEY"], values["DEEPSEEK_BASE_URL"], values["DEEPSEEK_MODEL"],
  211. *numbers.values(), proxy_auth=proxy_tokens(path, environ))
  212. def make_job(record, config, count, *, workbook_llm=False):
  213. if record.get('raw_record', {}).get('workbook_source') and not workbook_llm:
  214. # Default to local extraction unless the authorized workflow opts in.
  215. summary = '\n'.join(dict.fromkeys(line.strip() for line in record['data']['body'].splitlines()
  216. if line.strip() and not re.fullmatch(r'第\d+行', line.strip())))
  217. return {'id': record['id'], 'scenario': scenario(record), 'local_summary': summary,
  218. 'key': fingerprint(['local-extractive-v1', record['id'], facts(record)]),
  219. 'source_hash': fingerprint(facts(record)), 'estimated_input_tokens': 0}
  220. request_messages = messages(record)
  221. if request_messages is None:
  222. return None
  223. estimated = count(chat_text(request_messages))
  224. if estimated + config.max_output_tokens + config.margin > config.context_tokens:
  225. raise ValueError("input_exceeds_context_budget:" + record["id"])
  226. payload = {"model": config.model, "messages": request_messages,
  227. "max_tokens": config.max_output_tokens, "thinking": {"type": "disabled"},
  228. "response_format": {"type": "json_object"}, "temperature": 0, "stream": False}
  229. return {"id": record["id"], "scenario": scenario(record), "payload": payload,
  230. "key": fingerprint(["json-retry-v2", record["id"], config.endpoint, PROMPT_VERSION, payload,
  231. config.context_tokens, config.margin, config.max_output_limit, config.max_context_limit]),
  232. "source_hash": fingerprint(facts(record)), "estimated_input_tokens": estimated}
  233. class NoRedirect(urllib.request.HTTPRedirectHandler):
  234. def redirect_request(self, req, fp, code, msg, headers, newurl):
  235. return None
  236. class RequestFailure(Exception):
  237. def __init__(self, code, retryable=True, token_limited=False):
  238. super().__init__(code)
  239. self.code, self.retryable = code, retryable
  240. self.token_limited = token_limited
  241. def transport(payload, config):
  242. try:
  243. with deepseek_response(payload, config) as response:
  244. content = response.read(4_000_001)
  245. decoded = response.decoded if isinstance(response, _BufferedResponse) else _JSON_UNREAD
  246. if getattr(response, 'length', None) not in (None, 0) and len(content) <= 4_000_000:
  247. raise RequestFailure('response_incomplete')
  248. if len(content) > 4_000_000:
  249. raise RequestFailure("response_too_large", False)
  250. return strict_json(content.decode("utf-8")) if decoded is _JSON_UNREAD else decoded
  251. except urllib.error.HTTPError as exc:
  252. code = exc.code
  253. exc.close()
  254. raise RequestFailure(f"http_{code}", code in {408, 429} or code >= 500) from None
  255. except TimeoutError:
  256. raise RequestFailure("request_timeout") from None
  257. except urllib.error.URLError as exc:
  258. raise RequestFailure("request_timeout" if isinstance(exc.reason, TimeoutError) else "network_failure") from None
  259. except OSError:
  260. raise RequestFailure("network_failure") from None
  261. except HTTPException:
  262. raise RequestFailure("response_incomplete") from None
  263. except (ValueError, UnicodeError):
  264. raise RequestFailure("invalid_response_json") from None
  265. def validate_response(response):
  266. try:
  267. choice = response["choices"][0]
  268. if choice["finish_reason"] != "stop":
  269. raise RequestFailure("incomplete_completion", token_limited=choice["finish_reason"] == "length")
  270. result = strict_json(choice["message"]["content"])
  271. if not isinstance(result, dict) or set(result) != {"summary"}:
  272. raise RequestFailure("invalid_summary_schema")
  273. summary = result["summary"]
  274. if not isinstance(summary, str) or not summary.strip():
  275. raise RequestFailure("empty_summary")
  276. usage = response.get("usage", {})
  277. usage = {k: v for k, v in usage.items() if isinstance(v, int) and not isinstance(v, bool) and v >= 0}
  278. return {"summary": summary.strip(), "summary_result": {"summary": summary.strip()},
  279. "finish_reason": "stop", "usage": usage, "model": response.get("model"), "generated_at": utc_now()}
  280. except (KeyError, IndexError, TypeError, ValueError, AttributeError):
  281. raise RequestFailure("invalid_response_schema") from None
  282. def generate(job, config, send=transport, sleep=time.sleep):
  283. if 'local_summary' in job:
  284. summary = job['local_summary']
  285. return {'status': 'generated', 'summary': summary, 'summary_result': {'summary': summary},
  286. 'finish_reason': 'stop', 'model': 'local-extractive-v1', 'generated_at': utc_now(),
  287. 'attempts': 0, 'attempt_history': [], 'usage': {}}
  288. active, history = config, []
  289. for attempt in range(1, config.max_attempts + 1):
  290. attempt_info = {"attempt": attempt, "max_output_tokens": active.max_output_tokens,
  291. "context_tokens": active.context_tokens}
  292. try:
  293. payload = {**job["payload"], "max_tokens": active.max_output_tokens}
  294. response = send(payload, active)
  295. # Only copy known diagnostic fields; never store incomplete content or reasoning text.
  296. if isinstance(response, dict):
  297. choice = (response.get("choices") or [{}])[0]
  298. if isinstance(choice, dict) and choice.get("finish_reason") in {"stop", "length", "content_filter", "tool_calls", "insufficient_system_resource"}:
  299. attempt_info["finish_reason"] = choice["finish_reason"]
  300. raw_usage = response.get("usage")
  301. if isinstance(raw_usage, dict):
  302. attempt_info["usage"] = {k: v for k, v in raw_usage.items() if isinstance(v, int) and not isinstance(v, bool) and v >= 0}
  303. result = validate_response(response)
  304. actual = result["usage"].get("prompt_tokens", 0) + result["usage"].get("completion_tokens", 0)
  305. if actual > active.context_tokens:
  306. raise RequestFailure("actual_context_budget_exceeded", False)
  307. history.append({**attempt_info, "status": "generated"})
  308. return {"status": "generated", **result, "attempts": attempt, "attempt_history": history,
  309. "effective_max_output_tokens": active.max_output_tokens, "effective_context_tokens": active.context_tokens}
  310. except RequestFailure as exc:
  311. history.append({**attempt_info, "status": "failed", "error": exc.code})
  312. if not exc.retryable or attempt == config.max_attempts:
  313. return {"status": "failed", "error": exc.code, "attempts": attempt, "attempt_history": history}
  314. if exc.token_limited:
  315. output = active.max_output_tokens * 2
  316. context = active.context_tokens
  317. input_tokens = attempt_info.get("usage", {}).get("prompt_tokens", job["estimated_input_tokens"])
  318. while input_tokens + output + active.margin > context:
  319. context *= 2
  320. if output > config.max_output_limit or context > config.max_context_limit:
  321. return {"status": "failed", "error": "token_retry_limit_reached", "attempts": attempt, "attempt_history": history}
  322. active = replace(active, max_output_tokens=output, context_tokens=context)
  323. continue
  324. sleep(min(2 ** attempt, 15))
  325. @contextmanager
  326. def run_lock(path):
  327. """OS lock releases on process exit, including crashes; lock file may remain."""
  328. with path.open("a+b") as stream:
  329. if path.stat().st_size == 0:
  330. stream.write(b"0")
  331. stream.flush()
  332. stream.seek(0)
  333. try:
  334. if os.name == "nt":
  335. import msvcrt
  336. msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1)
  337. else:
  338. import fcntl
  339. fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
  340. except OSError:
  341. raise ValueError("another_summary_run_is_active") from None
  342. try:
  343. yield
  344. finally:
  345. stream.seek(0)
  346. if os.name == "nt":
  347. msvcrt.locking(stream.fileno(), msvcrt.LK_UNLCK, 1)
  348. else:
  349. fcntl.flock(stream, fcntl.LOCK_UN)
  350. class Cache:
  351. def __init__(self, path):
  352. self.db = sqlite3.connect(path)
  353. self.db.execute("CREATE TABLE IF NOT EXISTS results (key TEXT PRIMARY KEY, value TEXT NOT NULL)")
  354. def get(self, key):
  355. row = self.db.execute("SELECT value FROM results WHERE key=?", (key,)).fetchone()
  356. return json.loads(row[0]) if row else None
  357. def put(self, key, value):
  358. self.db.execute("INSERT OR REPLACE INTO results VALUES (?,?)", (key, dump(value)))
  359. self.db.commit()
  360. def close(self):
  361. self.db.close()
  362. def load_records(prepared):
  363. records, hashes, ids = [], {}, set()
  364. for kind in KINDS:
  365. path = prepared / f"{kind}.jsonl"
  366. raw = path.read_bytes()
  367. hashes[path.name] = hashlib.sha256(raw).hexdigest()
  368. for line in raw.decode("utf-8").splitlines():
  369. record = strict_json(line)
  370. if record["kind"] != kind or record["id"] in ids:
  371. raise ValueError("duplicate_id_or_kind_mismatch")
  372. ids.add(record["id"])
  373. records.append(record)
  374. return records, hashes
  375. def cache_success(result):
  376. return (isinstance(result, dict) and result.get("status") == "generated"
  377. and isinstance(result.get("summary"), str) and bool(result["summary"].strip())
  378. and result.get("summary_result") == {"summary": result["summary"]}
  379. and result.get("finish_reason") == "stop")
  380. def enhanced(record, job, result):
  381. output = dict(record)
  382. if job is None:
  383. qa = record["kind"] == "qa"
  384. output.update(summary=None, summary_result=None, summary_status="not_required_qa" if qa else "excluded_catalog",
  385. retrieval_text=record["data"].get("question", "") if qa else None,
  386. summary_metadata={"scenario": None})
  387. return output
  388. result = result or {"status": "pending"}
  389. if result.get("status") == "generated" and not cache_success(result):
  390. result = {"status": "failed", "error": "invalid_cached_result"}
  391. summary = result.get("summary") if result["status"] == "generated" else None
  392. output.update(summary=summary, summary_result={"summary": summary} if summary else None, summary_status=result["status"],
  393. retrieval_text=record["title"] + "\n" + summary if summary else None,
  394. summary_metadata={"scenario": job["scenario"], "request_hash": job["key"],
  395. "source_hash": job["source_hash"], "prompt_version": 'local-extractive-v1' if 'local_summary' in job else PROMPT_VERSION,
  396. "estimated_input_tokens": job["estimated_input_tokens"],
  397. "model": result.get("model"), "generated_at": result.get("generated_at"),
  398. "usage": result.get("usage", {}), "error": result.get("error"),
  399. "finish_reason": result.get("finish_reason"),
  400. "attempt_history": result.get("attempt_history", []),
  401. "effective_max_output_tokens": result.get("effective_max_output_tokens"),
  402. "effective_context_tokens": result.get("effective_context_tokens"),
  403. "review_status": "not_reviewed"})
  404. return output
  405. def atomic_write(path, text):
  406. content = text.encode("utf-8")
  407. if path.exists() and path.read_bytes() == content:
  408. return
  409. with tempfile.NamedTemporaryFile(dir=path.parent, prefix=path.name, suffix=".tmp", delete=False) as f:
  410. tmp = Path(f.name)
  411. f.write(content)
  412. try:
  413. tmp.replace(path)
  414. finally:
  415. tmp.unlink(missing_ok=True)
  416. def materialize(records, jobs, cache, prepared, hashes, config, verify=False):
  417. statuses, scenarios, usage, errors = Counter(), Counter(), Counter(), []
  418. retry_errors, retried_records, max_completion = Counter(), 0, 0
  419. outputs = {kind: [] for kind in KINDS}
  420. for record in records:
  421. job = jobs.get(record["id"])
  422. result = cache.get(job["key"]) if job else None
  423. item = enhanced(record, job, result)
  424. statuses[item["summary_status"]] += 1
  425. if job:
  426. scenarios[job["scenario"]] += 1
  427. if result and result["status"] == "generated":
  428. usage.update(result.get("usage", {}))
  429. history = result.get("attempt_history", [])
  430. retried_records += len(history) > 1
  431. retry_errors.update(attempt["error"] for attempt in history if attempt.get("error"))
  432. max_completion = max(max_completion, result.get("usage", {}).get("completion_tokens", 0))
  433. if job and item["summary_status"] != "generated":
  434. errors.append({"id": record["id"], "status": item["summary_status"], "error": (result or {}).get("error")})
  435. outputs[record["kind"]].append(dump(item) + "\n")
  436. output_hashes = {}
  437. for kind, lines in outputs.items():
  438. text = "".join(lines)
  439. path = prepared / f"{kind}.with_summary.jsonl"
  440. if verify:
  441. if path.read_bytes() != text.encode("utf-8"):
  442. raise ValueError("enhanced_output_mismatch:" + kind)
  443. else:
  444. atomic_write(path, text)
  445. output_hashes[path.name] = hashlib.sha256(text.encode("utf-8")).hexdigest()
  446. report = {"prompt_version": PROMPT_VERSION, "model": config.model, "context_tokens": config.context_tokens,
  447. "max_output_tokens": config.max_output_tokens, "margin": config.margin, "thinking": "disabled",
  448. "max_output_limit": config.max_output_limit, "max_context_limit": config.max_context_limit,
  449. "source_records": len(records), "summary_requests": len(jobs), "statuses": dict(statuses),
  450. "scenarios": dict(scenarios), "usage_successful_records": dict(usage), "errors": errors,
  451. "retry_stats": {"records_retried": retried_records, "attempt_errors": dict(retry_errors),
  452. "max_completion_tokens": max_completion},
  453. "source_sha256": hashes, "output_sha256": output_hashes,
  454. "complete": not errors, "review_status": "not_reviewed"}
  455. if verify:
  456. saved = json.loads((prepared / "summary_report.json").read_text(encoding="utf-8"))
  457. if saved != report or errors:
  458. raise ValueError("summary_report_incomplete_or_mismatch")
  459. else:
  460. atomic_write(prepared / "summary_report.json", json.dumps(report, ensure_ascii=False, indent=2) + "\n")
  461. return report
  462. def execute(args):
  463. from tokenizers import Tokenizer
  464. config = load_config(args.env_file)
  465. workbook_llm = getattr(args, 'llm_workbook', False)
  466. if workbook_llm and getattr(args, 'local_only', False):
  467. raise ValueError('conflicting_workbook_summary_modes')
  468. tokenizer = Tokenizer.from_file(str(args.tokenizer))
  469. tokenizer.no_truncation()
  470. tokenizer.no_padding()
  471. records, hashes = load_records(args.prepared)
  472. # Do not silently replace already-authorized LLM workbook summaries with
  473. # the default local fallback on a later ordinary summary run.
  474. for kind in KINDS:
  475. existing = args.prepared / f'{kind}.with_summary.jsonl'
  476. if existing.exists() and not workbook_llm:
  477. for line in existing.read_text(encoding='utf-8').splitlines():
  478. item = strict_json(line)
  479. if item.get('raw_record', {}).get('workbook_source') and item.get('summary_metadata', {}).get('model') not in (None, 'local-extractive-v1'):
  480. raise ValueError('use_refresh_local_llm_workbook_to_preserve_authorized_summary')
  481. jobs = {}
  482. for record in records:
  483. job = make_job(record, config, lambda text: len(tokenizer.encode(text, add_special_tokens=False).ids),
  484. workbook_llm=workbook_llm)
  485. if job:
  486. jobs[record["id"]] = job
  487. print(dump({"stage": "preflight", "records": len(records), "requests": len(jobs),
  488. "scenarios": dict(Counter(job["scenario"] for job in jobs.values())),
  489. "max_input": max((j["estimated_input_tokens"] for j in jobs.values()), default=0),
  490. "context": config.context_tokens, "output": config.max_output_tokens}), flush=True)
  491. if args.dry_run:
  492. return 0
  493. with run_lock(args.prepared / ".summary.lock"):
  494. cache = Cache(args.prepared / ".summary_cache.sqlite3")
  495. try:
  496. if not args.verify_only:
  497. pending = [job for job in jobs.values() if not cache_success(cache.get(job["key"]))]
  498. if getattr(args, 'local_only', False) and any('local_summary' not in job for job in pending):
  499. raise ValueError('remote_summary_required_local_only')
  500. if args.sample_per_scenario:
  501. selected, counts = [], Counter()
  502. for job in pending:
  503. if counts[job["scenario"]] < args.sample_per_scenario:
  504. selected.append(job)
  505. counts[job["scenario"]] += 1
  506. pending = selected
  507. print(dump({"stage": "generate", "scheduled": len(pending), "concurrency": config.concurrency}), flush=True)
  508. with ThreadPoolExecutor(max_workers=config.concurrency) as pool:
  509. futures = {pool.submit(generate, job, config): job for job in pending}
  510. done, failed = 0, 0
  511. for future in as_completed(futures):
  512. job = futures[future]
  513. result = future.result()
  514. cache.put(job["key"], result)
  515. done += 1
  516. failed += result["status"] != "generated"
  517. if done % 10 == 0 or done == len(pending):
  518. print(dump({"stage": "progress", "done": done, "scheduled": len(pending), "failed": failed}), flush=True)
  519. # Never materialize stale sources if another process edited original JSONL during generation.
  520. if any(hashlib.sha256((args.prepared / name).read_bytes()).hexdigest() != value for name, value in hashes.items()):
  521. raise ValueError("source_changed_during_generation")
  522. report = materialize(records, jobs, cache, args.prepared, hashes, config, args.verify_only)
  523. print(dump({"stage": "verified" if args.verify_only else "saved", "statuses": report["statuses"], "complete": report["complete"]}), flush=True)
  524. return 0 if report["complete"] or (args.sample_per_scenario and not report["statuses"].get("failed")) else 1
  525. finally:
  526. cache.close()
  527. def main():
  528. parser = argparse.ArgumentParser(description=__doc__)
  529. parser.add_argument("--env-file", type=Path, default=ROOT / ".env")
  530. parser.add_argument("--prepared", type=Path, default=ROOT / "knowledge/prepared")
  531. parser.add_argument("--tokenizer", type=Path, default=ROOT / "models/deepseek-tokenizer/tokenizer.json")
  532. parser.add_argument("--dry-run", action="store_true")
  533. parser.add_argument("--verify-only", action="store_true")
  534. parser.add_argument("--local-only", action="store_true", help="Reuse cached summaries and extract workbook text locally; never call remote services")
  535. parser.add_argument("--llm-workbook", action="store_true", help="Include workbook records in configured remote LLM summaries (explicit authorization)")
  536. parser.add_argument("--sample-per-scenario", type=int, default=0)
  537. args = parser.parse_args()
  538. if args.sample_per_scenario < 0:
  539. parser.error("sample size must not be negative")
  540. try:
  541. return execute(args)
  542. except (ValueError, OSError, sqlite3.Error, KeyError, TypeError) as exc:
  543. # Never render arbitrary server responses, request bodies, credentials or connection URLs.
  544. print(dump({"stage": "fatal", "error_type": type(exc).__name__}), flush=True)
  545. return 1
  546. if __name__ == "__main__":
  547. raise SystemExit(main())