Files
googlecloudplatform--genera…/agents/adk/contract-compliance-pipeline/python-extraction-agent/app/fast_api_app.py
T
2026-07-13 13:30:30 +08:00

575 lines
19 KiB
Python

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""FastAPI application for the Contract Compliance Pipeline.
Provides REST API endpoints for contract upload, status checking, and
result retrieval. The live cockpit uses deterministic extraction and a real
ADK RemoteA2aAgent handoff to the Go compliance service so the demo remains
stable while still exercising the current A2A SendMessage protocol.
"""
import json
import os
import secrets
import shutil
import logging
import uuid
import time
import asyncio
import urllib.error
import urllib.request
from pathlib import Path
from typing import Optional
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import PlainTextResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from a2a.types import DataPart as A2ADataPart
from google.adk.a2a.agent.config import A2aRemoteAgentConfig, RequestInterceptor
from google.adk.a2a.converters.part_converter import (
A2A_DATA_PART_END_TAG,
A2A_DATA_PART_START_TAG,
A2A_DATA_PART_TEXT_MIME_TYPE,
)
from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
from google.adk.apps import App as ADKApp
from google.adk.cli.fast_api import get_fast_api_app
from google.adk.runners import Runner
from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.genai import types as genai_types
from app.state_schema import ComplianceStep
from app.tools import (
SANDBOX_DIR,
UPLOAD_DIR,
classify_contract_risk,
extract_contract_details_from_text,
)
from app.app_utils.telemetry import setup_telemetry
from app.live_compliance import (
create_compliance_case,
get_case,
latest_case_payload,
sync_case_with_session_state,
case_payload,
artifact_response,
_event,
)
setup_telemetry()
logger = logging.getLogger(__name__)
GO_AGENT_CARD_URL = os.environ.get(
"GO_AGENT_CARD_URL",
"http://localhost:8888/.well-known/agent.json",
)
PROJECT_ROOT = Path(SANDBOX_DIR)
SAMPLE_CONTRACTS_DIR = PROJECT_ROOT / "sample-contracts"
def _go_jsonrpc_url() -> str:
if GO_AGENT_CARD_URL.endswith("/.well-known/agent.json"):
return GO_AGENT_CARD_URL.removesuffix("/.well-known/agent.json")
return GO_AGENT_CARD_URL.rstrip("/")
def build_contract_handoff_data(case_id: str, details: dict, policy: dict | None) -> dict:
data = {
"schema_version": "contract-compliance.a2a.v1",
"case_id": case_id,
"contract": details,
}
if policy:
data["policy"] = policy
return data
def build_go_message_payload(case_id: str, details: dict, policy: dict | None) -> dict:
return {
"jsonrpc": "2.0",
"id": f"case-{case_id}",
"method": "SendMessage",
"params": {
"metadata": {"task_id": case_id},
"message": {
"messageId": f"case-{case_id}-request",
"taskId": case_id,
"role": "ROLE_USER",
"parts": [{
"data": build_contract_handoff_data(case_id, details, policy),
"mediaType": "application/json"
}],
},
},
}
def extract_verdict_from_go_response(response: dict) -> dict:
result = response.get("result", response)
parts = (
result
.get("status", {})
.get("message", {})
.get("parts", [])
)
for part in parts:
if "data" in part:
return part["data"]
if part.get("kind") == "data" or part.get("type") == "data":
return part.get("data", {})
raise RuntimeError("Go compliance service returned no verdict data")
def _a2a_data_part_as_genai_part(data: dict) -> genai_types.Part:
"""Build a GenAI Part that ADK converts back into an A2A DataPart."""
data_part = A2ADataPart(data=data)
return genai_types.Part(
inline_data=genai_types.Blob(
data=(
A2A_DATA_PART_START_TAG
+ data_part.model_dump_json(by_alias=True, exclude_none=True).encode("utf-8")
+ A2A_DATA_PART_END_TAG
),
mime_type=A2A_DATA_PART_TEXT_MIME_TYPE,
)
)
def _http_json(method: str, url: str, payload: dict | None = None, timeout: float = 10.0) -> dict:
data = None if payload is None else json.dumps(payload).encode("utf-8")
request = urllib.request.Request(
url,
data=data,
headers={"Content-Type": "application/json", "Accept": "application/json"},
method=method,
)
with urllib.request.urlopen(request, timeout=timeout) as response:
return json.loads(response.read().decode("utf-8"))
async def invoke_go_compliance_service(
case_id: str,
details: dict,
policy: dict | None,
timeout: float = 10.0,
) -> dict:
"""Calls the Go service through ADK RemoteA2aAgent and the A2A SDK."""
agent_card = await asyncio.to_thread(_http_json, "GET", GO_AGENT_CARD_URL, None, timeout)
payload = build_go_message_payload(case_id, details, policy)
request_data = build_contract_handoff_data(case_id, details, policy)
async def add_task_metadata(_ctx, a2a_message, parameters):
parameters.request_metadata = {"task_id": case_id}
return a2a_message, parameters
remote_agent = RemoteA2aAgent(
name="compliance_agent",
agent_card=GO_AGENT_CARD_URL,
description="Validates extracted contract fields with the Go A2A compliance service.",
timeout=timeout,
config=A2aRemoteAgentConfig(
request_interceptors=[
RequestInterceptor(before_request=add_task_metadata)
]
),
)
session_service = InMemorySessionService()
adk_app_name = "contract_compliance_remote_handoff"
await session_service.create_session(
app_name=adk_app_name,
user_id="live-cockpit",
session_id=case_id,
state={},
)
runner = Runner(
app=ADKApp(name=adk_app_name, root_agent=remote_agent),
session_service=session_service,
)
response_task = None
request_message = None
error_message = None
try:
async for event in runner.run_async(
user_id="live-cockpit",
session_id=case_id,
new_message=genai_types.Content(
role="user",
parts=[_a2a_data_part_as_genai_part(request_data)],
),
):
metadata = event.custom_metadata or {}
request_message = metadata.get("a2a:request", request_message)
response_task = metadata.get("a2a:response", response_task)
if event.error_message:
error_message = event.error_message
finally:
await remote_agent.cleanup()
if error_message:
raise RuntimeError(error_message)
if response_task is None:
raise RuntimeError("ADK RemoteA2aAgent returned no A2A task response")
response = {
"jsonrpc": "2.0",
"id": payload["id"],
"result": response_task,
}
return {
"agent_card": agent_card,
"request": payload,
"response": response,
"adk_request_message": request_message,
"adk_response_task": response_task,
"verdict": extract_verdict_from_go_response(response),
}
# --- Session Credentials ---
def get_session_secret() -> str:
"""Generate or load a session secret key."""
if os.getenv("SESSION_SECRET_KEY"):
return os.getenv("SESSION_SECRET_KEY")
secret_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "session_secret.txt")
if os.path.exists(secret_file):
try:
with open(secret_file, "r", encoding="utf-8") as f:
return f.read().strip()
except Exception:
pass
ephemeral = secrets.token_hex(32)
try:
with open(secret_file, "w", encoding="utf-8") as f:
f.write(ephemeral)
except Exception:
pass
return ephemeral
# --- App Configuration ---
AGENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
session_service_uri = "sqlite+aiosqlite:///sessions.db"
app: FastAPI = get_fast_api_app(
agents_dir=AGENT_DIR,
web=True,
session_service_uri=session_service_uri,
allow_origins=[
"http://127.0.0.1:8000",
"http://localhost:8000",
"null",
],
)
app.title = "contract-compliance-pipeline"
app.description = "Multi-agent contract compliance pipeline using ADK and A2A protocol."
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# --- Startup ---
@app.on_event("startup")
def boot_sandbox_bounds():
"""Initialize upload directory and preload sample contracts."""
os.makedirs(UPLOAD_DIR, exist_ok=True)
# Preload sample contracts from project root
sample_dir = str(SAMPLE_CONTRACTS_DIR)
if os.path.exists(sample_dir):
for filename in os.listdir(sample_dir):
src = os.path.join(sample_dir, filename)
dst = os.path.join(UPLOAD_DIR, filename)
if os.path.isfile(src):
try:
shutil.copy2(src, dst)
except Exception as ex:
logger.warning(f"Preloading contract {filename} failed: {ex}")
logger.info(f"Upload directory ready: {UPLOAD_DIR}")
# --- ADK Session Storage ---
db_session_service = DatabaseSessionService(db_url=session_service_uri)
# --- Static Files ---
static_folder = Path(__file__).parent / "static" / "live-compliance"
if static_folder.exists():
app.mount(
"/live-compliance",
StaticFiles(directory=str(static_folder), html=True),
name="live-compliance",
)
@app.get("/demo")
def demo_router() -> RedirectResponse:
return RedirectResponse(url="/live-compliance/")
# --- API Endpoints ---
@app.get("/api/compliance/current")
def get_current_case_status() -> dict:
return latest_case_payload()
@app.get("/api/compliance/cases/{case_id}")
def get_compliance_case(case_id: str) -> dict:
case = get_case(case_id)
return {"active": True, "case": case_payload(case)}
@app.get("/api/compliance/cases/{case_id}/artifacts/{artifact_id}")
def get_compliance_artifact(case_id: str, artifact_id: str):
return artifact_response(case_id, artifact_id)
@app.get("/api/compliance/sample-contracts/{filename}")
def get_sample_contract(filename: str) -> PlainTextResponse:
"""Serves bundled sample contract text for the browser-driven demo."""
safe_filename = os.path.basename(filename)
target = (SAMPLE_CONTRACTS_DIR / safe_filename).resolve()
sample_root = SAMPLE_CONTRACTS_DIR.resolve()
if not str(target).startswith(str(sample_root) + os.path.sep):
raise HTTPException(status_code=403, detail="Invalid sample contract path.")
if not target.exists() or not target.is_file():
raise HTTPException(status_code=404, detail="Sample contract not found.")
if target.suffix.lower() not in {".pdf", ".txt", ".md"}:
raise HTTPException(status_code=400, detail="Unsupported sample contract type.")
return PlainTextResponse(
target.read_text(encoding="utf-8", errors="ignore"),
headers={"Cache-Control": "no-store"},
)
@app.post("/api/compliance/upload")
async def upload_contract_file(
file: UploadFile = File(...),
simulated_latency: float = Form(0.0),
simulated_server_state: str = Form("normal"),
custom_policies: Optional[str] = Form(None),
) -> dict:
"""Upload a text contract fixture and run the compliance pipeline.
Accepts text files and .pdf-named text fixtures, saves them securely, and kicks off the
multi-agent pipeline: extraction → A2A compliance check → reporting.
"""
filename = file.filename or "contract.pdf"
original_filename = os.path.basename(filename)
# Verify extensions
ext = os.path.splitext(filename)[1].lower()
if ext not in [".pdf", ".txt", ".md"]:
raise HTTPException(status_code=400, detail="Only .pdf, .txt, and .md text fixtures are supported.")
# Rename to UUID for security
case_id = str(uuid.uuid4())
safe_filename = f"{case_id}{ext}"
os.makedirs(UPLOAD_DIR, exist_ok=True)
secure_target = os.path.join(UPLOAD_DIR, safe_filename)
# Path traversal check
absolute_root = os.path.abspath(UPLOAD_DIR) + os.path.sep
absolute_target = os.path.abspath(secure_target)
if not absolute_target.startswith(absolute_root):
raise HTTPException(status_code=403, detail="Invalid upload path.")
# Size limit (5MB)
content = await file.read()
if len(content) > 5 * 1024 * 1024:
raise HTTPException(status_code=400, detail="File size exceeds 5MB limit.")
if ext == ".pdf" and content.lstrip().startswith(b"%PDF"):
raise HTTPException(
status_code=400,
detail=(
"Binary PDF parsing is not included in this demo. "
"Use bundled text fixtures or upload .txt/.md contract text."
),
)
# Save file
try:
with open(absolute_target, "wb") as f:
f.write(content)
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Failed to save file: {exc!s}")
# Create compliance case
case = await create_compliance_case(original_filename, db_session_service)
content_text = content.decode("utf-8", errors="ignore")
details = extract_contract_details_from_text(original_filename, content_text)
risk = classify_contract_risk(details)
policy = None
if custom_policies:
try:
policy = json.loads(custom_policies)
except json.JSONDecodeError as exc:
raise HTTPException(status_code=400, detail=f"Invalid custom_policies JSON: {exc!s}")
# Keep simulator values bounded so UI experiments cannot lock the API.
simulated_latency = max(0.0, min(float(simulated_latency), 5.0))
simulated_server_state = simulated_server_state.lower().strip()
trace_logs = [
{
"span": "POST /api/compliance/upload",
"service": "python-extraction-agent",
"duration_ms": 140,
"status": "file_ingested",
},
{
"span": "extract_contract_fields",
"service": "python-extraction-agent",
"duration_ms": 180,
"status": "fields_extracted",
},
{
"span": "classify_risk_level",
"service": "python-extraction-agent",
"duration_ms": 25,
"status": risk["risk_tier"].lower(),
},
]
_event(case, "agent", "Extractor completed", f"{details['contractor_name']} fields extracted and risk classified as {risk['risk_tier']}.")
handoff = {
"status": "prepared",
"source_agent": "python-extraction-agent",
"target_agent": "Security Compliance Validator",
"remote_agent": "google.adk.agents.RemoteA2aAgent(compliance_agent)",
"agent_card_url": GO_AGENT_CARD_URL,
"jsonrpc_url": _go_jsonrpc_url(),
"method": "SendMessage",
"task_id": case.id,
"contract_details": details,
"risk_assessment": risk,
"policy_override": bool(policy),
"request": build_go_message_payload(case.id, details, policy),
}
if simulated_latency:
await asyncio.sleep(simulated_latency)
try:
if simulated_server_state == "crashed":
raise ConnectionError("Simulated Go compliance service 503 failure")
handoff_result = await invoke_go_compliance_service(
case_id=case.id,
details=details,
policy=policy,
timeout=10.0,
)
verdict = handoff_result["verdict"]
handoff.update(
{
"status": "completed",
"agent_card": handoff_result["agent_card"],
"request": handoff_result["request"],
"response": handoff_result["response"],
"adk_request_message": handoff_result["adk_request_message"],
"adk_response_task": handoff_result["adk_response_task"],
"verdict": verdict,
}
)
current_step = (
ComplianceStep.APPROVED
if verdict.get("passed", False)
else ComplianceStep.REVIEW_READY
)
trace_logs.extend(
[
{
"span": "GET /.well-known/agent.json",
"service": "go-compliance-agent",
"duration_ms": 45,
"status": "agent_card_discovered",
},
{
"span": "ADK RemoteA2aAgent SendMessage",
"service": "go-compliance-agent",
"duration_ms": 120,
"status": "completed",
},
{
"span": "go_validate_policy",
"service": "go-compliance-agent",
"duration_ms": 18,
"status": "passed" if verdict.get("passed", False) else "violations_found",
},
]
)
except (ConnectionError, TimeoutError, urllib.error.URLError, RuntimeError) as exc:
logger.warning(f"Go compliance handoff failed for case {case.id}: {exc!s}")
current_step = ComplianceStep.MANUAL_REVIEW
verdict = {
"passed": False,
"violations": [
"SYSTEM TIMEOUT: External compliance service failed to respond within 30-second threshold.",
"FAIL-SAFE ACTION: Document routed for legal manager manual verification.",
],
"verdict_timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}
handoff.update(
{
"status": "failed",
"error": str(exc),
"verdict": verdict,
}
)
trace_logs.append(
{
"span": "resilience_fallback_gate",
"service": "python-extraction-agent",
"duration_ms": 450,
"status": "manual_review_routed",
}
)
session_state = {
"case_id": case.id,
"current_step": current_step,
"contract_filename": original_filename,
"contract_details": details,
"risk_assessment": risk,
"compliance_verdict": verdict,
"handoff": handoff,
"pending_signals": [],
"trace_logs": trace_logs,
"completion_time": time.strftime("%Y-%m-%d %H:%M:%S %Z"),
}
sync_case_with_session_state(case, session_state)
return {"active": True, "case": case_payload(case)}
# --- Local Dev Server ---
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="127.0.0.1", port=8000)