cba19c7bb3
Submit up to 1000 unbriefed documents to the provider Batch API in one
shot instead of individual synchronous LLM calls. Results are polled
every 30 minutes via a new Celery beat task and imported automatically.
- New worker: llm_batch_processor.py
- submit_llm_batch: guards against duplicate batches, builds JSONL
(OpenAI) or request list (Anthropic), stores state in AppSetting
- poll_llm_batch_results: checks batch status, imports completed
results with idempotency, emits notifications + triggers news fetch
- celery_app: register worker, route to llm queue, beat every 30 min
- admin API: POST /submit-llm-batch + GET /llm-batch-status endpoints
- Frontend: submitLlmBatch + getLlmBatchStatus in adminAPI; settings
page shows batch control row (openai/anthropic only) with live
progress line while batch is processing
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
402 lines
14 KiB
Python
402 lines
14 KiB
Python
"""
|
|
LLM Batch processor — submits and polls OpenAI/Anthropic Batch API jobs.
|
|
50% cheaper than synchronous calls; 24-hour processing window.
|
|
New bills still use the synchronous llm_processor task.
|
|
"""
|
|
import io
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import text
|
|
|
|
from app.config import settings
|
|
from app.database import get_sync_db
|
|
from app.models import Bill, BillBrief, BillDocument, Member
|
|
from app.models.setting import AppSetting
|
|
from app.services.llm_service import (
|
|
AMENDMENT_SYSTEM_PROMPT,
|
|
MAX_TOKENS_DEFAULT,
|
|
SYSTEM_PROMPT,
|
|
build_amendment_prompt,
|
|
build_prompt,
|
|
parse_brief_json,
|
|
)
|
|
from app.workers.celery_app import celery_app
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_BATCH_SETTING_KEY = "llm_active_batch"
|
|
|
|
|
|
# ── State helpers ──────────────────────────────────────────────────────────────
|
|
|
|
def _save_batch_state(db, state: dict):
|
|
row = db.get(AppSetting, _BATCH_SETTING_KEY)
|
|
if row:
|
|
row.value = json.dumps(state)
|
|
else:
|
|
row = AppSetting(key=_BATCH_SETTING_KEY, value=json.dumps(state))
|
|
db.add(row)
|
|
db.commit()
|
|
|
|
|
|
def _clear_batch_state(db):
|
|
row = db.get(AppSetting, _BATCH_SETTING_KEY)
|
|
if row:
|
|
db.delete(row)
|
|
db.commit()
|
|
|
|
|
|
# ── Request builder ────────────────────────────────────────────────────────────
|
|
|
|
def _build_request_data(db, doc_id: int, bill_id: str) -> tuple[str, str, str]:
|
|
"""Returns (custom_id, system_prompt, user_prompt) for a document."""
|
|
doc = db.get(BillDocument, doc_id)
|
|
if not doc or not doc.raw_text:
|
|
raise ValueError(f"Document {doc_id} missing or has no text")
|
|
|
|
bill = db.get(Bill, bill_id)
|
|
if not bill:
|
|
raise ValueError(f"Bill {bill_id} not found")
|
|
|
|
sponsor = db.get(Member, bill.sponsor_id) if bill.sponsor_id else None
|
|
|
|
bill_metadata = {
|
|
"title": bill.title or "Unknown Title",
|
|
"sponsor_name": sponsor.name if sponsor else "Unknown",
|
|
"party": sponsor.party if sponsor else "Unknown",
|
|
"state": sponsor.state if sponsor else "Unknown",
|
|
"chamber": bill.chamber or "Unknown",
|
|
"introduced_date": str(bill.introduced_date) if bill.introduced_date else "Unknown",
|
|
"latest_action_text": bill.latest_action_text or "None",
|
|
"latest_action_date": str(bill.latest_action_date) if bill.latest_action_date else "Unknown",
|
|
}
|
|
|
|
previous_full_brief = (
|
|
db.query(BillBrief)
|
|
.filter_by(bill_id=bill_id, brief_type="full")
|
|
.order_by(BillBrief.created_at.desc())
|
|
.first()
|
|
)
|
|
|
|
if previous_full_brief and previous_full_brief.document_id:
|
|
previous_doc = db.get(BillDocument, previous_full_brief.document_id)
|
|
if previous_doc and previous_doc.raw_text:
|
|
brief_type = "amendment"
|
|
prompt = build_amendment_prompt(doc.raw_text, previous_doc.raw_text, bill_metadata, MAX_TOKENS_DEFAULT)
|
|
system_prompt = AMENDMENT_SYSTEM_PROMPT + "\n\nIMPORTANT: Respond with ONLY valid JSON. No other text."
|
|
else:
|
|
brief_type = "full"
|
|
prompt = build_prompt(doc.raw_text, bill_metadata, MAX_TOKENS_DEFAULT)
|
|
system_prompt = SYSTEM_PROMPT
|
|
else:
|
|
brief_type = "full"
|
|
prompt = build_prompt(doc.raw_text, bill_metadata, MAX_TOKENS_DEFAULT)
|
|
system_prompt = SYSTEM_PROMPT
|
|
|
|
custom_id = f"doc-{doc_id}-{brief_type}"
|
|
return custom_id, system_prompt, prompt
|
|
|
|
|
|
# ── Submit task ────────────────────────────────────────────────────────────────
|
|
|
|
@celery_app.task(bind=True, name="app.workers.llm_batch_processor.submit_llm_batch")
|
|
def submit_llm_batch(self):
|
|
"""Submit all unbriefed documents to the OpenAI or Anthropic Batch API."""
|
|
db = get_sync_db()
|
|
try:
|
|
prov_row = db.get(AppSetting, "llm_provider")
|
|
model_row = db.get(AppSetting, "llm_model")
|
|
provider_name = ((prov_row.value if prov_row else None) or settings.LLM_PROVIDER).lower()
|
|
|
|
if provider_name not in ("openai", "anthropic"):
|
|
return {"status": "unsupported", "provider": provider_name}
|
|
|
|
# Check for already-active batch
|
|
active_row = db.get(AppSetting, _BATCH_SETTING_KEY)
|
|
if active_row:
|
|
try:
|
|
active = json.loads(active_row.value)
|
|
if active.get("status") == "processing":
|
|
return {"status": "already_active", "batch_id": active.get("batch_id")}
|
|
except Exception:
|
|
pass
|
|
|
|
# Find docs with text but no brief
|
|
rows = db.execute(text("""
|
|
SELECT bd.id AS doc_id, bd.bill_id, bd.govinfo_url
|
|
FROM bill_documents bd
|
|
LEFT JOIN bill_briefs bb ON bb.document_id = bd.id
|
|
WHERE bd.raw_text IS NOT NULL AND bb.id IS NULL
|
|
LIMIT 1000
|
|
""")).fetchall()
|
|
|
|
if not rows:
|
|
return {"status": "nothing_to_process"}
|
|
|
|
doc_ids = [r.doc_id for r in rows]
|
|
|
|
if provider_name == "openai":
|
|
model = (model_row.value if model_row else None) or settings.OPENAI_MODEL
|
|
batch_id = _submit_openai_batch(db, rows, model)
|
|
else:
|
|
model = (model_row.value if model_row else None) or settings.ANTHROPIC_MODEL
|
|
batch_id = _submit_anthropic_batch(db, rows, model)
|
|
|
|
state = {
|
|
"batch_id": batch_id,
|
|
"provider": provider_name,
|
|
"model": model,
|
|
"doc_ids": doc_ids,
|
|
"doc_count": len(doc_ids),
|
|
"submitted_at": datetime.utcnow().isoformat(),
|
|
"status": "processing",
|
|
}
|
|
_save_batch_state(db, state)
|
|
logger.info(f"Submitted {len(doc_ids)}-doc batch to {provider_name}: {batch_id}")
|
|
return {"status": "submitted", "batch_id": batch_id, "doc_count": len(doc_ids)}
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def _submit_openai_batch(db, rows, model: str) -> str:
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=settings.OPENAI_API_KEY)
|
|
|
|
lines = []
|
|
for row in rows:
|
|
try:
|
|
custom_id, system_prompt, prompt = _build_request_data(db, row.doc_id, row.bill_id)
|
|
except Exception as exc:
|
|
logger.warning(f"Skipping doc {row.doc_id}: {exc}")
|
|
continue
|
|
lines.append(json.dumps({
|
|
"custom_id": custom_id,
|
|
"method": "POST",
|
|
"url": "/v1/chat/completions",
|
|
"body": {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
"temperature": 0.1,
|
|
"max_tokens": MAX_TOKENS_DEFAULT,
|
|
},
|
|
}))
|
|
|
|
jsonl_bytes = "\n".join(lines).encode()
|
|
file_obj = client.files.create(
|
|
file=("batch.jsonl", io.BytesIO(jsonl_bytes), "application/jsonl"),
|
|
purpose="batch",
|
|
)
|
|
batch = client.batches.create(
|
|
input_file_id=file_obj.id,
|
|
endpoint="/v1/chat/completions",
|
|
completion_window="24h",
|
|
)
|
|
return batch.id
|
|
|
|
|
|
def _submit_anthropic_batch(db, rows, model: str) -> str:
|
|
import anthropic
|
|
client = anthropic.Anthropic(api_key=settings.ANTHROPIC_API_KEY)
|
|
|
|
requests = []
|
|
for row in rows:
|
|
try:
|
|
custom_id, system_prompt, prompt = _build_request_data(db, row.doc_id, row.bill_id)
|
|
except Exception as exc:
|
|
logger.warning(f"Skipping doc {row.doc_id}: {exc}")
|
|
continue
|
|
requests.append({
|
|
"custom_id": custom_id,
|
|
"params": {
|
|
"model": model,
|
|
"max_tokens": 4096,
|
|
"system": [{"type": "text", "text": system_prompt, "cache_control": {"type": "ephemeral"}}],
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
},
|
|
})
|
|
|
|
batch = client.messages.batches.create(requests=requests)
|
|
return batch.id
|
|
|
|
|
|
# ── Poll task ──────────────────────────────────────────────────────────────────
|
|
|
|
@celery_app.task(bind=True, name="app.workers.llm_batch_processor.poll_llm_batch_results")
|
|
def poll_llm_batch_results(self):
|
|
"""Check active batch status and import completed results (runs every 30 min via beat)."""
|
|
db = get_sync_db()
|
|
try:
|
|
active_row = db.get(AppSetting, _BATCH_SETTING_KEY)
|
|
if not active_row:
|
|
return {"status": "no_active_batch"}
|
|
|
|
try:
|
|
state = json.loads(active_row.value)
|
|
except Exception:
|
|
_clear_batch_state(db)
|
|
return {"status": "invalid_state"}
|
|
|
|
batch_id = state["batch_id"]
|
|
provider_name = state["provider"]
|
|
model = state["model"]
|
|
|
|
if provider_name == "openai":
|
|
return _poll_openai(db, state, batch_id, model)
|
|
elif provider_name == "anthropic":
|
|
return _poll_anthropic(db, state, batch_id, model)
|
|
else:
|
|
_clear_batch_state(db)
|
|
return {"status": "unknown_provider"}
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
# ── Result processing helpers ──────────────────────────────────────────────────
|
|
|
|
def _save_brief(db, doc_id: int, bill_id: str, brief, brief_type: str, govinfo_url) -> bool:
|
|
"""Idempotency check + save. Returns True if saved, False if already exists."""
|
|
if db.query(BillBrief).filter_by(document_id=doc_id).first():
|
|
return False
|
|
|
|
db_brief = BillBrief(
|
|
bill_id=bill_id,
|
|
document_id=doc_id,
|
|
brief_type=brief_type,
|
|
summary=brief.summary,
|
|
key_points=brief.key_points,
|
|
risks=brief.risks,
|
|
deadlines=brief.deadlines,
|
|
topic_tags=brief.topic_tags,
|
|
llm_provider=brief.llm_provider,
|
|
llm_model=brief.llm_model,
|
|
govinfo_url=govinfo_url,
|
|
)
|
|
db.add(db_brief)
|
|
db.commit()
|
|
db.refresh(db_brief)
|
|
return True
|
|
|
|
|
|
def _emit_notifications_and_news(db, bill_id: str, brief, brief_type: str):
|
|
bill = db.get(Bill, bill_id)
|
|
if not bill:
|
|
return
|
|
from app.workers.notification_utils import (
|
|
emit_bill_notification,
|
|
emit_member_follow_notifications,
|
|
emit_topic_follow_notifications,
|
|
)
|
|
event_type = "new_amendment" if brief_type == "amendment" else "new_document"
|
|
emit_bill_notification(db, bill, event_type, brief.summary)
|
|
emit_member_follow_notifications(db, bill, event_type, brief.summary)
|
|
emit_topic_follow_notifications(db, bill, event_type, brief.summary, brief.topic_tags or [])
|
|
|
|
from app.workers.news_fetcher import fetch_news_for_bill
|
|
fetch_news_for_bill.delay(bill_id)
|
|
|
|
|
|
def _parse_custom_id(custom_id: str) -> tuple[int, str]:
|
|
"""Parse 'doc-{doc_id}-{brief_type}' → (doc_id, brief_type)."""
|
|
parts = custom_id.split("-")
|
|
return int(parts[1]), parts[2]
|
|
|
|
|
|
def _poll_openai(db, state: dict, batch_id: str, model: str) -> dict:
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=settings.OPENAI_API_KEY)
|
|
|
|
batch = client.batches.retrieve(batch_id)
|
|
logger.info(f"OpenAI batch {batch_id} status: {batch.status}")
|
|
|
|
if batch.status in ("failed", "cancelled", "expired"):
|
|
_clear_batch_state(db)
|
|
return {"status": batch.status}
|
|
|
|
if batch.status != "completed":
|
|
return {"status": "processing", "batch_status": batch.status}
|
|
|
|
content = client.files.content(batch.output_file_id).read().decode()
|
|
saved = failed = 0
|
|
|
|
for line in content.strip().split("\n"):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
item = json.loads(line)
|
|
custom_id = item["custom_id"]
|
|
doc_id, brief_type = _parse_custom_id(custom_id)
|
|
|
|
if item.get("error"):
|
|
logger.warning(f"Batch result error for {custom_id}: {item['error']}")
|
|
failed += 1
|
|
continue
|
|
|
|
raw = item["response"]["body"]["choices"][0]["message"]["content"]
|
|
brief = parse_brief_json(raw, "openai", model)
|
|
|
|
doc = db.get(BillDocument, doc_id)
|
|
if not doc:
|
|
failed += 1
|
|
continue
|
|
|
|
if _save_brief(db, doc_id, doc.bill_id, brief, brief_type, doc.govinfo_url):
|
|
_emit_notifications_and_news(db, doc.bill_id, brief, brief_type)
|
|
saved += 1
|
|
except Exception as exc:
|
|
logger.warning(f"Failed to process OpenAI batch result line: {exc}")
|
|
failed += 1
|
|
|
|
_clear_batch_state(db)
|
|
logger.info(f"OpenAI batch {batch_id} complete: {saved} saved, {failed} failed")
|
|
return {"status": "completed", "saved": saved, "failed": failed}
|
|
|
|
|
|
def _poll_anthropic(db, state: dict, batch_id: str, model: str) -> dict:
|
|
import anthropic
|
|
client = anthropic.Anthropic(api_key=settings.ANTHROPIC_API_KEY)
|
|
|
|
batch = client.messages.batches.retrieve(batch_id)
|
|
logger.info(f"Anthropic batch {batch_id} processing_status: {batch.processing_status}")
|
|
|
|
if batch.processing_status != "ended":
|
|
return {"status": "processing", "batch_status": batch.processing_status}
|
|
|
|
saved = failed = 0
|
|
|
|
for result in client.messages.batches.results(batch_id):
|
|
try:
|
|
custom_id = result.custom_id
|
|
doc_id, brief_type = _parse_custom_id(custom_id)
|
|
|
|
if result.result.type != "succeeded":
|
|
logger.warning(f"Batch result {custom_id} type: {result.result.type}")
|
|
failed += 1
|
|
continue
|
|
|
|
raw = result.result.message.content[0].text
|
|
brief = parse_brief_json(raw, "anthropic", model)
|
|
|
|
doc = db.get(BillDocument, doc_id)
|
|
if not doc:
|
|
failed += 1
|
|
continue
|
|
|
|
if _save_brief(db, doc_id, doc.bill_id, brief, brief_type, doc.govinfo_url):
|
|
_emit_notifications_and_news(db, doc.bill_id, brief, brief_type)
|
|
saved += 1
|
|
except Exception as exc:
|
|
logger.warning(f"Failed to process Anthropic batch result: {exc}")
|
|
failed += 1
|
|
|
|
_clear_batch_state(db)
|
|
logger.info(f"Anthropic batch {batch_id} complete: {saved} saved, {failed} failed")
|
|
return {"status": "completed", "saved": saved, "failed": failed}
|