Multi-Agent Compliance Orchestration Platform

from pydantic import BaseModel
import hashlib, json, boto3

CLAIM_CHECK_THRESHOLD = 50_000  # bytes; 50KB keeps inline state well under 400KB DDB cap

class S3Ref(BaseModel):
    """Pointer that replaces oversized fields in DynamoDB state."""
    kind: str = "s3_ref"
    bucket: str
    key: str
    sha256: str
    size_bytes: int

def serialize_field(value, run_id: str, field_name: str) -> dict | S3Ref:
    payload = json.dumps(value).encode("utf-8")
    if len(payload) < CLAIM_CHECK_THRESHOLD:
        return value  # stays inline in DynamoDB

    digest = hashlib.sha256(payload).hexdigest()
    key = f"state/{run_id}/{field_name}/{digest}.json"
    s3.put_object(Bucket=STATE_BUCKET, Key=key, Body=payload)

    return S3Ref(
        bucket=STATE_BUCKET,
        key=key,
        sha256=digest,
        size_bytes=len(payload),
    )

def hydrate_field(value):
    if isinstance(value, dict) and value.get("kind") == "s3_ref":
        obj = s3.get_object(Bucket=value["bucket"], Key=value["key"])
        return json.loads(obj["Body"].read())
    return value

from langgraph.graph import StateGraph, END

CONFIDENCE_THRESHOLD = 0.75

def reviewer_node(state: ComplianceState) -> ComplianceState:
    score = grade_finding(state.latest_finding)  # deterministic rubric, no LLM
    state.review = ReviewResult(
        finding_id=state.latest_finding.id,
        confidence=score.confidence,
        missing_evidence=score.missing_evidence,
        schema_ok=score.schema_ok,
    )
    return state

def route_after_review(state: ComplianceState) -> str:
    review = state.review
    if state.revision_count >= state.max_revisions:
        return "human_interrupt"              # terminal, no tickets
    if not review.schema_ok:
        return "evidence_enricher"
    if review.confidence < CONFIDENCE_THRESHOLD:
        return "evidence_enricher"
    return "remediation_drafter"

def evidence_enricher_node(state: ComplianceState) -> ComplianceState:
    for gap in state.review.missing_evidence:
        state.raw_evidence.enrichments.append(fetch_targeted(gap))
    state.revision_count += 1
    return state

def human_interrupt_node(state: ComplianceState) -> ComplianceState:
    # Terminal. No LLM, no Jira, no external calls.
    unresolved = [c for c in state.controls if c.status != "complete"]
    log.warning(
        "ESCALATION REQUIRED thread_id=%s client=%s framework=%s "
        "revision_count=%d controls=%s",
        state.thread_id, state.client, state.framework,
        state.revision_count, [c.id for c in unresolved],
    )
    state.human_approval_status = "pending"
    return state  # graph ends here

graph = StateGraph(ComplianceState)
graph.add_node("reviewer", reviewer_node)
graph.add_node("evidence_enricher", evidence_enricher_node)
graph.add_node("human_interrupt", human_interrupt_node)
graph.add_conditional_edges("reviewer", route_after_review, {
    "evidence_enricher": "evidence_enricher",
    "remediation_drafter": "remediation_drafter",
    "human_interrupt": "human_interrupt",
})
graph.add_edge("evidence_enricher", "gap_analyzer")  # loop back with new evidence
graph.add_edge("human_interrupt", END)               # terminal

from abc import ABC, abstractmethod
from enum import Enum

class Framework(str, Enum):
    FEDRAMP_20X = "FEDRAMP_20X"
    SOC2_TYPE2  = "SOC2_TYPE2"
    ISO_27001   = "ISO_27001"   # stub
    HIPAA       = "HIPAA"       # stub

class FrameworkProfile(ABC):
    framework: Framework

    @abstractmethod
    def analyst_backstory(self) -> str: ...
    @abstractmethod
    def standards_corpus(self) -> list[str]: ...
    @abstractmethod
    def assessor_persona(self) -> str: ...
    @abstractmethod
    def writer_citation_style(self) -> str: ...
    @abstractmethod
    def reporting_vocabulary(self) -> dict[str, str]: ...

class FedRAMP20xProfile(FrameworkProfile):
    framework = Framework.FEDRAMP_20X
    def analyst_backstory(self) -> str:
        return "FedRAMP 20x 3PAO assessor specializing in KSI, ADS, and CCM families."
    def writer_citation_style(self) -> str:
        return "KSI-<family>-<number>"  # e.g. KSI-IAM-02

class SOC2TypeIIProfile(FrameworkProfile):
    framework = Framework.SOC2_TYPE2
    def analyst_backstory(self) -> str:
        return "AICPA SOC 2 Type II practitioner focused on Trust Services Criteria."
    def writer_citation_style(self) -> str:
        return "TSC <criterion>.<point>"  # e.g. CC6.1

class ISO27001Profile(FrameworkProfile):    # stub — wired for future
    framework = Framework.ISO_27001

class HIPAAProfile(FrameworkProfile):        # stub — wired for future
    framework = Framework.HIPAA

def build_graph(profile: FrameworkProfile) -> StateGraph:
    graph = StateGraph(ComplianceState)
    graph.add_node("gap_analyzer", gap_analyzer_factory(profile))
    graph.add_node("reviewer", reviewer_factory(profile))
    # ...same graph topology for every framework
    return graph.compile(checkpointer=DynamoDBSaver())

from functools import wraps

def idempotent(per_revision: bool = False):
    """Record completion only on normal return. Raises retry on resume."""
    def decorator(node_fn):
        @wraps(node_fn)
        def wrapped(state: ComplianceState) -> ComplianceState:
            key = _ledger_key(state, node_fn.__name__, per_revision)
            if ledger.get(key):
                log.info("skip %s: already completed", node_fn.__name__)
                return state

            result = node_fn(state)       # may raise
            ledger.put(key)               # only reached on normal return
            return result
        return wrapped
    return decorator

def _ledger_key(state, node_name, per_revision):
    base = f"{state.thread_id}:{node_name}"
    return f"{base}:r{state.revision_count}" if per_revision else base

# --- loopable nodes: one execution per revision is legitimate ---
@idempotent(per_revision=True)
def gap_analyzer(state): ...

@idempotent(per_revision=True)
def evidence_enricher(state): ...

@idempotent(per_revision=True)
def evidence_merger(state): ...

@idempotent(per_revision=True)
def custom_test_runner(state): ...

@idempotent(per_revision=True)
def reviewer(state): ...

# --- terminal / single-shot nodes: replay would produce duplicates ---
@idempotent(per_revision=False)
def remediation_drafter(state): ...

@idempotent(per_revision=False)
def ui_mapper(state):
    status = state.ui_upload_status
    if status not in {"complete", "skipped", "not_applicable"}:
        # raise -> partial return discarded, ledger NOT written,
        # graph re-pauses at interrupt_before on resume
        raise ValueError(f"ui_upload_status={status!r} not allowed")
    return state

@idempotent(per_revision=False)
def jira_publisher(state): ...

@idempotent(per_revision=False)
def human_interrupt(state): ...