"""MinIO/S3 helpers for the shared kb bucket (formerly `insurance-kb`). Layout in the bucket: inbox/{law,regulation,circular,caselaw,tool}/ processed/{law,regulation,circular,caselaw,tool}/ failed/{law,regulation,circular,caselaw,tool}/ After a successful ingest, the file is moved from inbox/* to processed/*. On failure it is moved to failed/*. """ from __future__ import annotations import logging import os from typing import Iterable import boto3 from botocore.client import Config as BotoConfig logger = logging.getLogger("shira.kb.s3") _KINDS = ("law", "regulation", "circular", "caselaw", "tool") def _client(): endpoint = os.environ.get("KB_S3_ENDPOINT") access = os.environ.get("KB_S3_ACCESS_KEY") secret = os.environ.get("KB_S3_SECRET_KEY") region = os.environ.get("KB_S3_REGION", "us-east-1") if not (endpoint and access and secret): raise RuntimeError("KB_S3_* environment variables are not set") return boto3.client( "s3", endpoint_url=endpoint, aws_access_key_id=access, aws_secret_access_key=secret, region_name=region, config=BotoConfig(signature_version="s3v4"), ) def _bucket() -> str: return os.environ.get("KB_S3_BUCKET", "legal-kb") def list_inbox() -> list[dict]: """Return inbox objects across all kinds, excluding folder markers and sidecars. Sidecar files (.meta.json) carry per-document metadata and are not standalone documents, so they are filtered out here — the scanner fetches them separately when processing their parent. """ s3 = _client() bucket = _bucket() items: list[dict] = [] for kind in _KINDS: prefix = f"inbox/{kind}/" paginator = s3.get_paginator("list_objects_v2") for page in paginator.paginate(Bucket=bucket, Prefix=prefix): for obj in page.get("Contents", []) or []: key = obj["Key"] if key.endswith("/"): continue if key.endswith(".meta.json"): continue filename = key[len(prefix):] if not filename: continue items.append({ "key": key, "kind": kind, "filename": filename, "size": obj["Size"], }) return items def fetch(key: str) -> bytes: s3 = _client() resp = s3.get_object(Bucket=_bucket(), Key=key) return resp["Body"].read() def move(src_key: str, dst_key: str) -> None: s3 = _client() bucket = _bucket() s3.copy_object( Bucket=bucket, Key=dst_key, CopySource={"Bucket": bucket, "Key": src_key}, MetadataDirective="COPY", ) s3.delete_object(Bucket=bucket, Key=src_key) logger.info("[kb.s3] moved %s -> %s", src_key, dst_key) def mark_processed(src_key: str, filename: str, kind: str) -> str: dst = f"processed/{kind}/{filename}" move(src_key, dst) # Best-effort: move the sidecar too if it exists. sidecar_src = src_key + ".meta.json" sidecar_dst = dst + ".meta.json" try: move(sidecar_src, sidecar_dst) except Exception: pass return dst def list_failed() -> list[dict]: s3 = _client() bucket = _bucket() items: list[dict] = [] for kind in _KINDS: prefix = f"failed/{kind}/" paginator = s3.get_paginator("list_objects_v2") for page in paginator.paginate(Bucket=bucket, Prefix=prefix): for obj in page.get("Contents", []) or []: key = obj["Key"] if key.endswith("/"): continue if key.endswith(".meta.json"): continue filename = key[len(prefix):] if not filename: continue items.append({"key": key, "kind": kind, "filename": filename, "size": obj["Size"]}) return items def mark_failed(src_key: str, filename: str, kind: str, reason: str) -> str: s3 = _client() bucket = _bucket() dst = f"failed/{kind}/{filename}" # S3 metadata is encoded as HTTP headers, so only ASCII is safe. safe_reason = reason.encode("ascii", "ignore").decode("ascii")[:800] s3.copy_object( Bucket=bucket, Key=dst, CopySource={"Bucket": bucket, "Key": src_key}, Metadata={"ingest_error": safe_reason}, MetadataDirective="REPLACE", ) s3.delete_object(Bucket=bucket, Key=src_key) # Best-effort: move the sidecar too. sidecar_src = src_key + ".meta.json" sidecar_dst = dst + ".meta.json" try: move(sidecar_src, sidecar_dst) except Exception: pass logger.info("[kb.s3] failed %s -> %s (%s)", src_key, dst, safe_reason[:80]) return dst def kinds() -> Iterable[str]: return _KINDS