"""Generate summaries using DeepSeek; secrets stay in memory, resumable local cache.""" import argparse import hashlib from http.client import HTTPConnection, HTTPSConnection, HTTPException from io import BytesIO import json import os import re import sqlite3 import tempfile import time import urllib.error import urllib.parse import urllib.request from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import contextmanager from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path from .prepare import ROOT, dump, strict_json from .summary_input import PROMPT_VERSION, chat_text, facts, messages, scenario from src.proxy_auth import ProxyAuthError, proxy_tokens, token_rejected KINDS = ("policy", "finance", "general", "qa") DEFAULTS = {"SUMMARY_CONTEXT_TOKENS": "32768", "SUMMARY_MAX_OUTPUT_TOKENS": "4096", "SUMMARY_TOKEN_MARGIN": "512", "SUMMARY_CONCURRENCY": "8", "SUMMARY_TIMEOUT_SECONDS": "120", "SUMMARY_MAX_ATTEMPTS": "3", "SUMMARY_MAX_OUTPUT_LIMIT": "32768", "SUMMARY_MAX_CONTEXT_LIMIT": "131072"} CONFIG_KEYS = set(DEFAULTS) | {"DEEPSEEK_API_KEY", "DEEPSEEK_BASE_URL", "DEEPSEEK_MODEL"} _JSON_UNREAD = object() def fingerprint(value): return hashlib.sha256(dump(value).encode("utf-8")).hexdigest() def utc_now(): return datetime.now(timezone.utc).isoformat(timespec="seconds") @dataclass class Config: api_key: str = field(repr=False) base_url: str = field(repr=False) model: str context_tokens: int = 32768 max_output_tokens: int = 4096 margin: int = 512 concurrency: int = 8 timeout: int = 120 max_attempts: int = 3 max_output_limit: int = 32768 max_context_limit: int = 131072 proxy_auth: object = field(default=None, repr=False, compare=False) @property def endpoint(self): return chat_endpoint(self.base_url) def chat_endpoint(base_url): """Resolve a base or complete chat URL without rewriting signed query values.""" try: if not isinstance(base_url, str) or any(c.isspace() for c in base_url) or "\\" in base_url: raise ValueError parts = urllib.parse.urlsplit(base_url) if (parts.scheme not in {"http", "https"} or not parts.hostname or parts.username is not None or parts.password is not None or parts.fragment): raise ValueError _ = parts.port except ValueError: raise ValueError("invalid_deepseek_base_url") from None path = parts.path.rstrip("/") if not path.endswith("/chat/completions"): path += "/chat/completions" return urllib.parse.urlunsplit(parts._replace(path=path)) class _BufferedResponse(BytesIO): def __init__(self, body, decoded): super().__init__(body) self.decoded, self.length = decoded, 0 class _PrefixedResponse: """Replay a bounded preflight read without changing JSON or SSE bytes.""" def __init__(self, prefix, response): self.prefix, self.offset, self.response = prefix, 0, response def __getattr__(self, name): return getattr(self.response, name) @property def length(self): remaining = getattr(self.response, 'length', None) return None if remaining is None else remaining + len(self.prefix) - self.offset def read1(self, count=-1): if self.offset < len(self.prefix): end = len(self.prefix) if count < 0 else min(len(self.prefix), self.offset + count) result = self.prefix[self.offset:end] self.offset = end return result return self.response.read1(count) def read(self, count=-1): first = self.read1(count) if self.offset < len(self.prefix) else b'' return first + self.response.read(-1 if count < 0 else count - len(first)) @contextmanager def deepseek_response(payload, config, *, stream=False, check=lambda: None): """One request, with a private connection per gateway call; always close it.""" configured = config.endpoint auth = getattr(config, 'proxy_auth', None) endpoint = auth.url(configured) if auth is not None else configured parts = urllib.parse.urlsplit(endpoint) query = dict(urllib.parse.parse_qsl(parts.query)) gateway = query.get('servertype') == 'deepseek' and bool(query.get('proxyToken')) headers = {'Authorization': 'Bearer ' + config.api_key, 'Content-Type': 'application/json', 'Accept': '*/*' if gateway else ('text/event-stream' if stream else 'application/json')} body = dump(payload).encode('utf-8') if not gateway: request = urllib.request.Request(endpoint, data=body, headers=headers) with urllib.request.build_opener(NoRedirect()).open(request, timeout=config.timeout) as response: yield response return headers['Connection'] = 'keep-alive' for attempt in range(2): check() parts = urllib.parse.urlsplit(endpoint) connection = (HTTPSConnection if parts.scheme == 'https' else HTTPConnection)( parts.hostname, parts.port, timeout=config.timeout) rejected = False try: target = (parts.path or '/') + ('?' + parts.query if parts.query else '') connection.request('POST', target, body=body, headers=headers) with connection.getresponse() as response: buffered, decoded = None, _JSON_UNREAD if not 200 <= response.status < 300 and response.status not in (401, 403): code = response.status raise RequestFailure(f'http_{code}', not stream and (code in {408, 429} or code >= 500)) # read1 validates HTTP chunk framing (peek can swallow broken # chunks). Replay all inspected bytes for normal SSE streams. check() prefix = bytearray(response.read1(4096)) deadline = time.monotonic() + config.timeout while prefix and not prefix.lstrip() and len(prefix) <= 4_000_000: check() if time.monotonic() >= deadline: raise RequestFailure('request_timeout', False) more = response.read1(4096) if not more: break prefix.extend(more) if len(prefix) > 4_000_000: raise RequestFailure('response_too_large', False) response = _PrefixedResponse(bytes(prefix), response) if prefix.lstrip().startswith(b'{'): buffered = response.read(4_000_001) if len(buffered) > 4_000_000: raise RequestFailure('response_too_large', False) if getattr(response, 'length', None) not in (None, 0): raise RequestFailure('response_incomplete', False) try: decoded = strict_json(buffered.decode('utf-8')) rejected = token_rejected(decoded) except (ValueError, UnicodeError): pass if not rejected: if not 200 <= response.status < 300: code = response.status raise RequestFailure(f'http_{code}', not stream and (code in {408, 429} or code >= 500)) if buffered is None: yield response else: with _BufferedResponse(buffered, decoded) as cached: yield cached return finally: connection.close() if rejected: if attempt or auth is None: if auth is not None: auth.reject(endpoint) raise RequestFailure('proxy_token_invalid', False) check() try: auth.refresh(endpoint) except ProxyAuthError: raise RequestFailure('proxy_token_refresh_failed', False) from None endpoint = auth.url(configured) def load_config(path, environ=None): values = dict(DEFAULTS) for line in path.read_text(encoding="utf-8-sig").splitlines(): line = line.strip() if line.startswith("export "): line = line[7:] if not line or line.startswith("#") or "=" not in line: continue key, value = line.split("=", 1) key, value = key.strip(), value.strip() if key not in CONFIG_KEYS: continue if value.startswith(("'", '"')): quote = value[0] end = value.find(quote, 1) if end < 0 or (value[end + 1:].strip() and not value[end + 1:].strip().startswith("#")): raise ValueError("invalid_env_quoting") value = value[1:end] else: value = value.split(" #", 1)[0].strip() values[key] = value env = os.environ if environ is None else environ values.update({key: env[key] for key in CONFIG_KEYS if key in env}) if not all(values.get(key) for key in ("DEEPSEEK_API_KEY", "DEEPSEEK_BASE_URL", "DEEPSEEK_MODEL")): raise ValueError("missing_deepseek_configuration") chat_endpoint(values["DEEPSEEK_BASE_URL"]) try: numbers = {key: int(values[key]) for key in DEFAULTS} except ValueError: raise ValueError("invalid_summary_integer_setting") from None if any(value <= 0 for value in numbers.values()) or numbers["SUMMARY_CONCURRENCY"] > 32 or numbers["SUMMARY_MAX_ATTEMPTS"] > 5: raise ValueError("summary_setting_out_of_range") if numbers["SUMMARY_MAX_OUTPUT_TOKENS"] + numbers["SUMMARY_TOKEN_MARGIN"] >= numbers["SUMMARY_CONTEXT_TOKENS"]: raise ValueError("output_budget_leaves_no_input_space") if numbers["SUMMARY_MAX_OUTPUT_TOKENS"] > numbers["SUMMARY_MAX_OUTPUT_LIMIT"] or numbers["SUMMARY_CONTEXT_TOKENS"] > numbers["SUMMARY_MAX_CONTEXT_LIMIT"]: raise ValueError("initial_budget_exceeds_retry_limit") return Config(values["DEEPSEEK_API_KEY"], values["DEEPSEEK_BASE_URL"], values["DEEPSEEK_MODEL"], *numbers.values(), proxy_auth=proxy_tokens(path, environ)) def make_job(record, config, count, *, workbook_llm=False): if record.get('raw_record', {}).get('workbook_source') and not workbook_llm: # Default to local extraction unless the authorized workflow opts in. summary = '\n'.join(dict.fromkeys(line.strip() for line in record['data']['body'].splitlines() if line.strip() and not re.fullmatch(r'第\d+行', line.strip()))) return {'id': record['id'], 'scenario': scenario(record), 'local_summary': summary, 'key': fingerprint(['local-extractive-v1', record['id'], facts(record)]), 'source_hash': fingerprint(facts(record)), 'estimated_input_tokens': 0} request_messages = messages(record) if request_messages is None: return None estimated = count(chat_text(request_messages)) if estimated + config.max_output_tokens + config.margin > config.context_tokens: raise ValueError("input_exceeds_context_budget:" + record["id"]) payload = {"model": config.model, "messages": request_messages, "max_tokens": config.max_output_tokens, "thinking": {"type": "disabled"}, "response_format": {"type": "json_object"}, "temperature": 0, "stream": False} return {"id": record["id"], "scenario": scenario(record), "payload": payload, "key": fingerprint(["json-retry-v2", record["id"], config.endpoint, PROMPT_VERSION, payload, config.context_tokens, config.margin, config.max_output_limit, config.max_context_limit]), "source_hash": fingerprint(facts(record)), "estimated_input_tokens": estimated} class NoRedirect(urllib.request.HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): return None class RequestFailure(Exception): def __init__(self, code, retryable=True, token_limited=False): super().__init__(code) self.code, self.retryable = code, retryable self.token_limited = token_limited def transport(payload, config): try: with deepseek_response(payload, config) as response: content = response.read(4_000_001) decoded = response.decoded if isinstance(response, _BufferedResponse) else _JSON_UNREAD if getattr(response, 'length', None) not in (None, 0) and len(content) <= 4_000_000: raise RequestFailure('response_incomplete') if len(content) > 4_000_000: raise RequestFailure("response_too_large", False) return strict_json(content.decode("utf-8")) if decoded is _JSON_UNREAD else decoded except urllib.error.HTTPError as exc: code = exc.code exc.close() raise RequestFailure(f"http_{code}", code in {408, 429} or code >= 500) from None except TimeoutError: raise RequestFailure("request_timeout") from None except urllib.error.URLError as exc: raise RequestFailure("request_timeout" if isinstance(exc.reason, TimeoutError) else "network_failure") from None except OSError: raise RequestFailure("network_failure") from None except HTTPException: raise RequestFailure("response_incomplete") from None except (ValueError, UnicodeError): raise RequestFailure("invalid_response_json") from None def validate_response(response): try: choice = response["choices"][0] if choice["finish_reason"] != "stop": raise RequestFailure("incomplete_completion", token_limited=choice["finish_reason"] == "length") result = strict_json(choice["message"]["content"]) if not isinstance(result, dict) or set(result) != {"summary"}: raise RequestFailure("invalid_summary_schema") summary = result["summary"] if not isinstance(summary, str) or not summary.strip(): raise RequestFailure("empty_summary") usage = response.get("usage", {}) usage = {k: v for k, v in usage.items() if isinstance(v, int) and not isinstance(v, bool) and v >= 0} return {"summary": summary.strip(), "summary_result": {"summary": summary.strip()}, "finish_reason": "stop", "usage": usage, "model": response.get("model"), "generated_at": utc_now()} except (KeyError, IndexError, TypeError, ValueError, AttributeError): raise RequestFailure("invalid_response_schema") from None def generate(job, config, send=transport, sleep=time.sleep): if 'local_summary' in job: summary = job['local_summary'] return {'status': 'generated', 'summary': summary, 'summary_result': {'summary': summary}, 'finish_reason': 'stop', 'model': 'local-extractive-v1', 'generated_at': utc_now(), 'attempts': 0, 'attempt_history': [], 'usage': {}} active, history = config, [] for attempt in range(1, config.max_attempts + 1): attempt_info = {"attempt": attempt, "max_output_tokens": active.max_output_tokens, "context_tokens": active.context_tokens} try: payload = {**job["payload"], "max_tokens": active.max_output_tokens} response = send(payload, active) # Only copy known diagnostic fields; never store incomplete content or reasoning text. if isinstance(response, dict): choice = (response.get("choices") or [{}])[0] if isinstance(choice, dict) and choice.get("finish_reason") in {"stop", "length", "content_filter", "tool_calls", "insufficient_system_resource"}: attempt_info["finish_reason"] = choice["finish_reason"] raw_usage = response.get("usage") if isinstance(raw_usage, dict): attempt_info["usage"] = {k: v for k, v in raw_usage.items() if isinstance(v, int) and not isinstance(v, bool) and v >= 0} result = validate_response(response) actual = result["usage"].get("prompt_tokens", 0) + result["usage"].get("completion_tokens", 0) if actual > active.context_tokens: raise RequestFailure("actual_context_budget_exceeded", False) history.append({**attempt_info, "status": "generated"}) return {"status": "generated", **result, "attempts": attempt, "attempt_history": history, "effective_max_output_tokens": active.max_output_tokens, "effective_context_tokens": active.context_tokens} except RequestFailure as exc: history.append({**attempt_info, "status": "failed", "error": exc.code}) if not exc.retryable or attempt == config.max_attempts: return {"status": "failed", "error": exc.code, "attempts": attempt, "attempt_history": history} if exc.token_limited: output = active.max_output_tokens * 2 context = active.context_tokens input_tokens = attempt_info.get("usage", {}).get("prompt_tokens", job["estimated_input_tokens"]) while input_tokens + output + active.margin > context: context *= 2 if output > config.max_output_limit or context > config.max_context_limit: return {"status": "failed", "error": "token_retry_limit_reached", "attempts": attempt, "attempt_history": history} active = replace(active, max_output_tokens=output, context_tokens=context) continue sleep(min(2 ** attempt, 15)) @contextmanager def run_lock(path): """OS lock releases on process exit, including crashes; lock file may remain.""" with path.open("a+b") as stream: if path.stat().st_size == 0: stream.write(b"0") stream.flush() stream.seek(0) try: if os.name == "nt": import msvcrt msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1) else: import fcntl fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB) except OSError: raise ValueError("another_summary_run_is_active") from None try: yield finally: stream.seek(0) if os.name == "nt": msvcrt.locking(stream.fileno(), msvcrt.LK_UNLCK, 1) else: fcntl.flock(stream, fcntl.LOCK_UN) class Cache: def __init__(self, path): self.db = sqlite3.connect(path) self.db.execute("CREATE TABLE IF NOT EXISTS results (key TEXT PRIMARY KEY, value TEXT NOT NULL)") def get(self, key): row = self.db.execute("SELECT value FROM results WHERE key=?", (key,)).fetchone() return json.loads(row[0]) if row else None def put(self, key, value): self.db.execute("INSERT OR REPLACE INTO results VALUES (?,?)", (key, dump(value))) self.db.commit() def close(self): self.db.close() def load_records(prepared): records, hashes, ids = [], {}, set() for kind in KINDS: path = prepared / f"{kind}.jsonl" raw = path.read_bytes() hashes[path.name] = hashlib.sha256(raw).hexdigest() for line in raw.decode("utf-8").splitlines(): record = strict_json(line) if record["kind"] != kind or record["id"] in ids: raise ValueError("duplicate_id_or_kind_mismatch") ids.add(record["id"]) records.append(record) return records, hashes def cache_success(result): return (isinstance(result, dict) and result.get("status") == "generated" and isinstance(result.get("summary"), str) and bool(result["summary"].strip()) and result.get("summary_result") == {"summary": result["summary"]} and result.get("finish_reason") == "stop") def enhanced(record, job, result): output = dict(record) if job is None: qa = record["kind"] == "qa" output.update(summary=None, summary_result=None, summary_status="not_required_qa" if qa else "excluded_catalog", retrieval_text=record["data"].get("question", "") if qa else None, summary_metadata={"scenario": None}) return output result = result or {"status": "pending"} if result.get("status") == "generated" and not cache_success(result): result = {"status": "failed", "error": "invalid_cached_result"} summary = result.get("summary") if result["status"] == "generated" else None output.update(summary=summary, summary_result={"summary": summary} if summary else None, summary_status=result["status"], retrieval_text=record["title"] + "\n" + summary if summary else None, summary_metadata={"scenario": job["scenario"], "request_hash": job["key"], "source_hash": job["source_hash"], "prompt_version": 'local-extractive-v1' if 'local_summary' in job else PROMPT_VERSION, "estimated_input_tokens": job["estimated_input_tokens"], "model": result.get("model"), "generated_at": result.get("generated_at"), "usage": result.get("usage", {}), "error": result.get("error"), "finish_reason": result.get("finish_reason"), "attempt_history": result.get("attempt_history", []), "effective_max_output_tokens": result.get("effective_max_output_tokens"), "effective_context_tokens": result.get("effective_context_tokens"), "review_status": "not_reviewed"}) return output def atomic_write(path, text): content = text.encode("utf-8") if path.exists() and path.read_bytes() == content: return with tempfile.NamedTemporaryFile(dir=path.parent, prefix=path.name, suffix=".tmp", delete=False) as f: tmp = Path(f.name) f.write(content) try: tmp.replace(path) finally: tmp.unlink(missing_ok=True) def materialize(records, jobs, cache, prepared, hashes, config, verify=False): statuses, scenarios, usage, errors = Counter(), Counter(), Counter(), [] retry_errors, retried_records, max_completion = Counter(), 0, 0 outputs = {kind: [] for kind in KINDS} for record in records: job = jobs.get(record["id"]) result = cache.get(job["key"]) if job else None item = enhanced(record, job, result) statuses[item["summary_status"]] += 1 if job: scenarios[job["scenario"]] += 1 if result and result["status"] == "generated": usage.update(result.get("usage", {})) history = result.get("attempt_history", []) retried_records += len(history) > 1 retry_errors.update(attempt["error"] for attempt in history if attempt.get("error")) max_completion = max(max_completion, result.get("usage", {}).get("completion_tokens", 0)) if job and item["summary_status"] != "generated": errors.append({"id": record["id"], "status": item["summary_status"], "error": (result or {}).get("error")}) outputs[record["kind"]].append(dump(item) + "\n") output_hashes = {} for kind, lines in outputs.items(): text = "".join(lines) path = prepared / f"{kind}.with_summary.jsonl" if verify: if path.read_bytes() != text.encode("utf-8"): raise ValueError("enhanced_output_mismatch:" + kind) else: atomic_write(path, text) output_hashes[path.name] = hashlib.sha256(text.encode("utf-8")).hexdigest() report = {"prompt_version": PROMPT_VERSION, "model": config.model, "context_tokens": config.context_tokens, "max_output_tokens": config.max_output_tokens, "margin": config.margin, "thinking": "disabled", "max_output_limit": config.max_output_limit, "max_context_limit": config.max_context_limit, "source_records": len(records), "summary_requests": len(jobs), "statuses": dict(statuses), "scenarios": dict(scenarios), "usage_successful_records": dict(usage), "errors": errors, "retry_stats": {"records_retried": retried_records, "attempt_errors": dict(retry_errors), "max_completion_tokens": max_completion}, "source_sha256": hashes, "output_sha256": output_hashes, "complete": not errors, "review_status": "not_reviewed"} if verify: saved = json.loads((prepared / "summary_report.json").read_text(encoding="utf-8")) if saved != report or errors: raise ValueError("summary_report_incomplete_or_mismatch") else: atomic_write(prepared / "summary_report.json", json.dumps(report, ensure_ascii=False, indent=2) + "\n") return report def execute(args): from tokenizers import Tokenizer config = load_config(args.env_file) workbook_llm = getattr(args, 'llm_workbook', False) if workbook_llm and getattr(args, 'local_only', False): raise ValueError('conflicting_workbook_summary_modes') tokenizer = Tokenizer.from_file(str(args.tokenizer)) tokenizer.no_truncation() tokenizer.no_padding() records, hashes = load_records(args.prepared) # Do not silently replace already-authorized LLM workbook summaries with # the default local fallback on a later ordinary summary run. for kind in KINDS: existing = args.prepared / f'{kind}.with_summary.jsonl' if existing.exists() and not workbook_llm: for line in existing.read_text(encoding='utf-8').splitlines(): item = strict_json(line) if item.get('raw_record', {}).get('workbook_source') and item.get('summary_metadata', {}).get('model') not in (None, 'local-extractive-v1'): raise ValueError('use_refresh_local_llm_workbook_to_preserve_authorized_summary') jobs = {} for record in records: job = make_job(record, config, lambda text: len(tokenizer.encode(text, add_special_tokens=False).ids), workbook_llm=workbook_llm) if job: jobs[record["id"]] = job print(dump({"stage": "preflight", "records": len(records), "requests": len(jobs), "scenarios": dict(Counter(job["scenario"] for job in jobs.values())), "max_input": max((j["estimated_input_tokens"] for j in jobs.values()), default=0), "context": config.context_tokens, "output": config.max_output_tokens}), flush=True) if args.dry_run: return 0 with run_lock(args.prepared / ".summary.lock"): cache = Cache(args.prepared / ".summary_cache.sqlite3") try: if not args.verify_only: pending = [job for job in jobs.values() if not cache_success(cache.get(job["key"]))] if getattr(args, 'local_only', False) and any('local_summary' not in job for job in pending): raise ValueError('remote_summary_required_local_only') if args.sample_per_scenario: selected, counts = [], Counter() for job in pending: if counts[job["scenario"]] < args.sample_per_scenario: selected.append(job) counts[job["scenario"]] += 1 pending = selected print(dump({"stage": "generate", "scheduled": len(pending), "concurrency": config.concurrency}), flush=True) with ThreadPoolExecutor(max_workers=config.concurrency) as pool: futures = {pool.submit(generate, job, config): job for job in pending} done, failed = 0, 0 for future in as_completed(futures): job = futures[future] result = future.result() cache.put(job["key"], result) done += 1 failed += result["status"] != "generated" if done % 10 == 0 or done == len(pending): print(dump({"stage": "progress", "done": done, "scheduled": len(pending), "failed": failed}), flush=True) # Never materialize stale sources if another process edited original JSONL during generation. if any(hashlib.sha256((args.prepared / name).read_bytes()).hexdigest() != value for name, value in hashes.items()): raise ValueError("source_changed_during_generation") report = materialize(records, jobs, cache, args.prepared, hashes, config, args.verify_only) print(dump({"stage": "verified" if args.verify_only else "saved", "statuses": report["statuses"], "complete": report["complete"]}), flush=True) return 0 if report["complete"] or (args.sample_per_scenario and not report["statuses"].get("failed")) else 1 finally: cache.close() def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--env-file", type=Path, default=ROOT / ".env") parser.add_argument("--prepared", type=Path, default=ROOT / "knowledge/prepared") parser.add_argument("--tokenizer", type=Path, default=ROOT / "models/deepseek-tokenizer/tokenizer.json") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--verify-only", action="store_true") parser.add_argument("--local-only", action="store_true", help="Reuse cached summaries and extract workbook text locally; never call remote services") parser.add_argument("--llm-workbook", action="store_true", help="Include workbook records in configured remote LLM summaries (explicit authorization)") parser.add_argument("--sample-per-scenario", type=int, default=0) args = parser.parse_args() if args.sample_per_scenario < 0: parser.error("sample size must not be negative") try: return execute(args) except (ValueError, OSError, sqlite3.Error, KeyError, TypeError) as exc: # Never render arbitrary server responses, request bodies, credentials or connection URLs. print(dump({"stage": "fatal", "error_type": type(exc).__name__}), flush=True) return 1 if __name__ == "__main__": raise SystemExit(main())