shira-hermes/api/services/kb/s3.py

"""MinIO/S3 helpers for the shared kb bucket (formerly `insurance-kb`).

Layout in the bucket:
    inbox/{law,regulation,circular,caselaw,tool}/<file>
    processed/{law,regulation,circular,caselaw,tool}/<file>
    failed/{law,regulation,circular,caselaw,tool}/<file>

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 (<name>.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