quote-judge-studio-v1/agent.py

"""QuoteJudge Studio full-stack A2A agent.

Deterministic product backend for comparing vendor quotes with weighted price,
delivery, and warranty scoring. The packed React frontend calls these same
public tools, and external clients retain typed FileUpload support.
"""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import re
import uuid
from datetime import datetime, timezone
from decimal import Decimal, ROUND_HALF_UP
from typing import Annotated, Any

from pydantic import BaseModel, Field, field_validator

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    FileUpload,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    State,
    UploadedFile,
    WorkspaceAccess,
    WorkspaceMode,
)

MAX_BROWSER_UPLOAD_BYTES = 256_000
MAX_BROWSER_UPLOADS = 4
ALLOWED_UPLOAD_MEDIA_TYPES = {
    "application/json",
    "text/csv",
    "text/plain",
    "application/vnd.ms-excel",
}

_MEMORY_COMPARISONS: dict[tuple[str, str], dict[str, Any]] = {}
_MEMORY_RECEIPTS: list[dict[str, Any]] = []


class QuoteJudgeStudioV1Config(BaseModel):
    """No user-configurable settings; all state is tenant-scoped by platform auth."""


class QuoteJudgeDurableState(BaseModel):
    """Marker state model: durable product state lives in managed Postgres."""

    database: str = "quote-judge-studio-v1-data"


class QuoteInput(BaseModel):
    vendor: str = Field(..., min_length=1, max_length=120)
    unit_price: float = Field(..., gt=0, le=1_000_000)
    quantity: int = Field(..., gt=0, le=100_000_000)
    delivery_days: int = Field(..., ge=0, le=3650)
    warranty_months: int = Field(..., ge=0, le=600)
    currency: str = Field(default="BRL", min_length=3, max_length=8)

    @field_validator("vendor", "currency")
    @classmethod
    def _clean_text(cls, value: str) -> str:
        return " ".join(str(value).strip().split())


class WeightInput(BaseModel):
    price: float = Field(..., ge=0, le=100)
    delivery: float = Field(..., ge=0, le=100)
    warranty: float = Field(..., ge=0, le=100)


class BrowserQuoteDocument(BaseModel):
    filename: str = Field(..., min_length=1, max_length=160)
    media_type: str = Field(..., min_length=3, max_length=120)
    data_base64: str = Field(..., min_length=1)

    @field_validator("filename")
    @classmethod
    def _safe_filename(cls, value: str) -> str:
        cleaned = str(value).replace("\\", "/").split("/")[-1].strip()
        if not cleaned or cleaned in {".", ".."}:
            raise ValueError("filename must be a simple file name")
        return cleaned


class ValidationIssue(BaseModel):
    code: str
    message: str
    field: str | None = None
    action: str | None = None


class QuoteJudgeStudioV1(A2AAgent[QuoteJudgeStudioV1Config, PlatformUserAuth]):
    name = "quote-judge-studio-v1"
    description = (
        "Compare at least two vendor quotes with price, delivery, and warranty "
        "weights; persist results; reopen prior comparisons; accept browser "
        "base64 uploads and typed FileUpload inputs; and persist execution receipts."
    )
    version = "0.1.2"

    config_model = QuoteJudgeStudioV1Config
    auth_model = PlatformUserAuth

    state = State.DURABLE
    state_model = QuoteJudgeDurableState
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic quote scoring; no LLM credential required.",
    )
    workspace_access = WorkspaceAccess.dynamic(
        max_files=8,
        allowed_modes=(WorkspaceMode.READ_ONLY, WorkspaceMode.READ_WRITE_OVERLAY),
        require_reason=False,
        max_total_size_bytes=2_000_000,
    )
    tools_used = ("postgres", "mcp", "file-upload")
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="quote-judge-studio-v1-data",
                provider="neon",
                engine="postgres",
                scope="user",
                branch="main",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
                scale_to_zero=True,
            ),
        )
    )

    @a2a.tool(
        description="Compare vendor quotes with price/delivery/warranty weights and persist the result plus a production receipt.",
        timeout_seconds=60,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        quotes: list[QuoteInput],
        weights: WeightInput,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        issues = _validate_request(comparison_id, quotes, weights)
        if issues:
            result = _validation_result(comparison_id, issues)
            await _persist_receipt(ctx, tenant, comparison_id, "compare_quotes", result, {"comparison_id": comparison_id, "quotes": [q.model_dump() for q in quotes], "weights": weights.model_dump()})
            return result

        result = _score_quotes(comparison_id, quotes, weights)
        result["receipt"] = await _persist_receipt(
            ctx,
            tenant,
            comparison_id,
            "compare_quotes",
            result,
            {"comparison_id": comparison_id, "quotes": [q.model_dump() for q in quotes], "weights": weights.model_dump()},
        )
        await _save_comparison(tenant, comparison_id, result)
        await ctx.emit_progress(f"saved comparison {comparison_id}: {result['recommendation']['vendor']}")
        return result

    @a2a.tool(
        description="Reopen a persisted QuoteJudge comparison by id for the authenticated user.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def get_comparison(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        if not _valid_comparison_id(comparison_id):
            return _validation_result(
                comparison_id,
                [ValidationIssue(code="invalid_comparison_id", message="comparison_id must contain only letters, numbers, dot, underscore, or hyphen.", field="comparison_id", action="Use a stable id such as studio-quote-v1.")],
            )
        result = await _load_comparison(tenant, comparison_id)
        if result is None:
            return {
                "ok": False,
                "code": "comparison_not_found",
                "comparison_id": comparison_id,
                "message": "No saved comparison exists for this signed-in user and id.",
                "action": "Run compare_quotes or upload/paste at least two quotes first.",
            }
        result = dict(result)
        result["receipt"] = await _persist_receipt(
            ctx,
            tenant,
            comparison_id,
            "get_comparison",
            result,
            {"comparison_id": comparison_id},
        )
        return result

    @a2a.tool(
        description="Browser-safe upload bridge: parse bounded base64 JSON/CSV/text quote files, compare, persist, and return a receipt.",
        timeout_seconds=60,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_uploaded_quotes_browser(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        documents: list[BrowserQuoteDocument],
        weights: WeightInput,
    ) -> dict[str, Any]:
        parsed = _parse_browser_documents(documents)
        if not parsed["ok"]:
            return _validation_result(comparison_id, parsed["issues"])
        return await self.compare_quotes(ctx, comparison_id, parsed["quotes"], weights)

    @a2a.tool(
        description="Compatibility FileUpload entrypoint for external clients; parses an uploaded quote file and compares/persists results.",
        timeout_seconds=60,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_uploaded_quote_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        quote_file: Annotated[
            UploadedFile,
            FileUpload(
                accept=("application/json", "text/csv", "text/plain"),
                max_bytes=MAX_BROWSER_UPLOAD_BYTES,
                description="JSON array/object, CSV, or line-oriented quote text with vendor/unit_price/quantity/delivery_days/warranty_months.",
            ),
        ],
        weights: WeightInput,
    ) -> dict[str, Any]:
        issues = _validate_uploaded_file_meta(quote_file.filename, quote_file.media_type, quote_file.size_bytes)
        if issues:
            return _validation_result(comparison_id, issues)
        try:
            data = ctx.workspace.read_bytes(quote_file.path)  # concrete runtime clients expose this helper
        except Exception:  # noqa: BLE001
            return _validation_result(
                comparison_id,
                [ValidationIssue(code="upload_read_failed", message="The uploaded file could not be read from the invocation workspace.", field="quote_file", action=f"Re-upload the file and ensure it is below {MAX_BROWSER_UPLOAD_BYTES} bytes.")],
            )
        parsed = _parse_document_bytes(quote_file.filename, quote_file.media_type, data)
        if not parsed["ok"]:
            return _validation_result(comparison_id, parsed["issues"])
        return await self.compare_quotes(ctx, comparison_id, parsed["quotes"], weights)


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    auth = ctx.auth
    stable = auth.user_id if auth.user_id is not None else (auth.sub or auth.email)
    return f"user:{stable}"


def _valid_comparison_id(value: str) -> bool:
    return bool(re.fullmatch(r"[A-Za-z0-9._-]{1,96}", str(value or "")))


def _validate_request(comparison_id: str, quotes: list[QuoteInput], weights: WeightInput) -> list[ValidationIssue]:
    issues: list[ValidationIssue] = []
    if not _valid_comparison_id(comparison_id):
        issues.append(ValidationIssue(code="invalid_comparison_id", message="comparison_id must contain only letters, numbers, dot, underscore, or hyphen.", field="comparison_id", action="Use a stable id such as studio-quote-v1."))
    if len(quotes) < 2:
        issues.append(ValidationIssue(code="at_least_two_quotes_required", message="QuoteJudge needs at least two vendor quotes to make a comparison.", field="quotes", action="Add at least one more vendor quote."))
    vendors = [q.vendor.casefold() for q in quotes]
    if len(vendors) != len(set(vendors)):
        issues.append(ValidationIssue(code="duplicate_vendor", message="Each quote vendor must be unique within a comparison.", field="quotes.vendor", action="Rename or merge duplicate vendor rows."))
    total_weight = weights.price + weights.delivery + weights.warranty
    if total_weight <= 0:
        issues.append(ValidationIssue(code="positive_weight_required", message="At least one scoring weight must be greater than zero.", field="weights", action="Set price, delivery, or warranty to a positive weight."))
    return issues


def _validation_result(comparison_id: str, issues: list[ValidationIssue]) -> dict[str, Any]:
    code = issues[0].code if issues else "validation_failed"
    return {
        "ok": False,
        "code": code,
        "comparison_id": comparison_id,
        "message": issues[0].message if issues else "The request could not be validated.",
        "issues": [issue.model_dump() for issue in issues],
    }


def _score_quotes(comparison_id: str, quotes: list[QuoteInput], weights: WeightInput) -> dict[str, Any]:
    total_weight = weights.price + weights.delivery + weights.warranty
    normalized_weights = {
        "price": weights.price / total_weight,
        "delivery": weights.delivery / total_weight,
        "warranty": weights.warranty / total_weight,
    }
    totals = [q.unit_price * q.quantity for q in quotes]
    deliveries = [q.delivery_days for q in quotes]
    warranties = [q.warranty_months for q in quotes]

    rows: list[dict[str, Any]] = []
    for index, quote in enumerate(quotes):
        price_score = _lower_is_better(totals[index], totals)
        delivery_score = _lower_is_better(deliveries[index], deliveries)
        warranty_score = _higher_is_better(warranties[index], warranties)
        weighted_score = (
            price_score * normalized_weights["price"]
            + delivery_score * normalized_weights["delivery"]
            + warranty_score * normalized_weights["warranty"]
        )
        rows.append(
            {
                "vendor": quote.vendor,
                "unit_price": _money(quote.unit_price),
                "quantity": quote.quantity,
                "total_price": _money(quote.unit_price * quote.quantity),
                "currency": quote.currency,
                "delivery_days": quote.delivery_days,
                "warranty_months": quote.warranty_months,
                "scores": {
                    "price": _round_score(price_score),
                    "delivery": _round_score(delivery_score),
                    "warranty": _round_score(warranty_score),
                    "weighted_total": _round_score(weighted_score),
                },
            }
        )
    rows.sort(key=lambda row: (-row["scores"]["weighted_total"], row["total_price"], row["delivery_days"], -row["warranty_months"], row["vendor"]))
    winner = rows[0]
    return {
        "ok": True,
        "comparison_id": comparison_id,
        "recommendation": {
            "vendor": winner["vendor"],
            "weighted_score": winner["scores"]["weighted_total"],
            "why": _recommendation_reason(winner),
        },
        "weights": {
            "price": weights.price,
            "delivery": weights.delivery,
            "warranty": weights.warranty,
            "normalized": {key: _round_score(value * 100) for key, value in normalized_weights.items()},
        },
        "quotes_ranked": rows,
        "validation": {"ok": True, "issues": []},
        "saved_at": datetime.now(timezone.utc).isoformat(),
    }


def _lower_is_better(value: float, values: list[float]) -> float:
    low, high = min(values), max(values)
    if high == low:
        return 100.0
    return 100.0 * (high - value) / (high - low)


def _higher_is_better(value: float, values: list[float]) -> float:
    low, high = min(values), max(values)
    if high == low:
        return 100.0
    return 100.0 * (value - low) / (high - low)


def _round_score(value: float) -> float:
    return float(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def _money(value: float) -> float:
    return float(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def _recommendation_reason(row: dict[str, Any]) -> str:
    return (
        f"{row['vendor']} has the highest weighted score ({row['scores']['weighted_total']}) "
        f"after balancing total price {row['currency']} {row['total_price']}, "
        f"{row['delivery_days']} delivery days, and {row['warranty_months']} months warranty."
    )


def _connect() -> Any | None:
    url = os.environ.get("DATABASE_URL")
    if not url:
        return None
    import psycopg
    from psycopg.rows import dict_row

    return psycopg.connect(url, row_factory=dict_row)


async def _save_comparison(tenant: str, comparison_id: str, result: dict[str, Any]) -> None:
    conn = _connect()
    if conn is None:
        _MEMORY_COMPARISONS[(tenant, comparison_id)] = json.loads(json.dumps(result))
        return
    with conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO quote_comparisons (tenant_key, comparison_id, recommendation_vendor, result, payload, weighted_score, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, NOW())
                ON CONFLICT (tenant_key, comparison_id) DO UPDATE SET
                    recommendation_vendor = EXCLUDED.recommendation_vendor,
                    result = EXCLUDED.result,
                    payload = EXCLUDED.payload,
                    weighted_score = EXCLUDED.weighted_score,
                    updated_at = NOW()
                """,
                (
                    tenant,
                    comparison_id,
                    result["recommendation"]["vendor"],
                    json.dumps(result),
                    json.dumps(result),
                    result["recommendation"]["weighted_score"],
                ),
            )


async def _load_comparison(tenant: str, comparison_id: str) -> dict[str, Any] | None:
    conn = _connect()
    if conn is None:
        stored = _MEMORY_COMPARISONS.get((tenant, comparison_id))
        return json.loads(json.dumps(stored)) if stored is not None else None
    with conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT COALESCE(result, payload) AS result FROM quote_comparisons WHERE tenant_key = %s AND comparison_id = %s",
                (tenant, comparison_id),
            )
            row = cur.fetchone()
            if not row:
                return None
            value = row["result"]
            return value if isinstance(value, dict) else json.loads(value)


async def _persist_receipt(
    ctx: RunContext[PlatformUserAuth],
    tenant: str,
    comparison_id: str,
    tool_name: str,
    result: dict[str, Any],
    inputs: dict[str, Any],
) -> dict[str, Any]:
    now = datetime.now(timezone.utc).isoformat()
    input_hash = _sha256_json(inputs)
    result_hash = _sha256_json({k: v for k, v in result.items() if k != "receipt"})
    receipt = {
        "receipt_id": str(uuid.uuid5(uuid.NAMESPACE_URL, f"{tenant}:{comparison_id}:{tool_name}:{input_hash}:{result_hash}")),
        "agent": QuoteJudgeStudioV1.name,
        "agent_version": QuoteJudgeStudioV1.version,
        "tool_name": tool_name,
        "comparison_id": comparison_id,
        "status": "ok" if result.get("ok") else "validation_error",
        "input_hash": input_hash,
        "result_hash": result_hash,
        "created_at": now,
        "production": True,
    }
    conn = _connect()
    if conn is None:
        _MEMORY_RECEIPTS.append({"tenant_key": tenant, "payload": receipt})
    else:
        with conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO quote_judge_receipts (receipt_id, tenant_key, comparison_id, tool_name, status, input_hash, payload)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (receipt_id) DO UPDATE SET payload = EXCLUDED.payload
                    """,
                    (receipt["receipt_id"], tenant, comparison_id, tool_name, receipt["status"], input_hash, json.dumps(receipt)),
                )
    await ctx.emit_event(a2a.AgentEvent(kind="quote_judge_receipt_persisted", payload={"receipt_id": receipt["receipt_id"], "comparison_id": comparison_id, "tool_name": tool_name}))
    return receipt


def _sha256_json(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


def _parse_browser_documents(documents: list[BrowserQuoteDocument]) -> dict[str, Any]:
    issues: list[ValidationIssue] = []
    if not documents:
        issues.append(ValidationIssue(code="upload_required", message="Upload at least one JSON, CSV, or text file containing two or more quotes.", field="documents", action="Choose a quote file or paste quote rows instead."))
    if len(documents) > MAX_BROWSER_UPLOADS:
        issues.append(ValidationIssue(code="too_many_uploads", message=f"At most {MAX_BROWSER_UPLOADS} files can be uploaded at once.", field="documents", action="Upload fewer files or merge quote rows."))
    quotes: list[QuoteInput] = []
    for index, doc in enumerate(documents[:MAX_BROWSER_UPLOADS]):
        if doc.media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
            issues.append(ValidationIssue(code="unsupported_media_type", message=f"{doc.filename} has unsupported media type {doc.media_type}.", field=f"documents[{index}].media_type", action="Use JSON, CSV, or plain text."))
            continue
        try:
            data = base64.b64decode(doc.data_base64, validate=True)
        except Exception:  # noqa: BLE001
            issues.append(ValidationIssue(code="invalid_base64", message=f"{doc.filename} is not valid base64.", field=f"documents[{index}].data_base64", action="Re-select the file so the browser can encode it again."))
            continue
        if len(data) > MAX_BROWSER_UPLOAD_BYTES:
            issues.append(ValidationIssue(code="upload_too_large", message=f"{doc.filename} exceeds {MAX_BROWSER_UPLOAD_BYTES} bytes.", field=f"documents[{index}]", action="Upload a smaller quote file."))
            continue
        parsed = _parse_document_bytes(doc.filename, doc.media_type, data)
        if parsed["ok"]:
            quotes.extend(parsed["quotes"])
        else:
            issues.extend(parsed["issues"])
    if issues:
        return {"ok": False, "issues": issues}
    return {"ok": True, "quotes": quotes}


def _validate_uploaded_file_meta(filename: str, media_type: str, size_bytes: int) -> list[ValidationIssue]:
    issues: list[ValidationIssue] = []
    if media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
        issues.append(ValidationIssue(code="unsupported_media_type", message=f"{filename} has unsupported media type {media_type}.", field="quote_file.media_type", action="Upload JSON, CSV, or plain text."))
    if size_bytes > MAX_BROWSER_UPLOAD_BYTES:
        issues.append(ValidationIssue(code="upload_too_large", message=f"{filename} exceeds {MAX_BROWSER_UPLOAD_BYTES} bytes.", field="quote_file", action="Upload a smaller file."))
    return issues


def _parse_document_bytes(filename: str, media_type: str, data: bytes) -> dict[str, Any]:
    try:
        text = data.decode("utf-8-sig")
    except UnicodeDecodeError:
        return {"ok": False, "issues": [ValidationIssue(code="upload_not_utf8", message=f"{filename} must be UTF-8 text, CSV, or JSON.", field="quote_file", action="Export the quote file as UTF-8.")]}
    try:
        if media_type == "application/json" or filename.lower().endswith(".json"):
            raw = json.loads(text)
            rows = raw.get("quotes", raw) if isinstance(raw, dict) else raw
            if not isinstance(rows, list):
                raise ValueError("JSON must be an array or an object with a quotes array")
            quotes = [QuoteInput.model_validate(row) for row in rows]
        elif media_type in {"text/csv", "application/vnd.ms-excel"} or filename.lower().endswith(".csv"):
            reader = csv.DictReader(io.StringIO(text))
            quotes = [QuoteInput.model_validate(_normalize_quote_row(row)) for row in reader]
        else:
            quotes = _parse_plain_text_quotes(text)
        return {"ok": True, "quotes": quotes}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "issues": [ValidationIssue(code="quote_parse_failed", message=f"Could not parse {filename}: {exc}", field="quote_file", action="Use columns vendor, unit_price, quantity, delivery_days, warranty_months.")]}


def _normalize_quote_row(row: dict[str, Any]) -> dict[str, Any]:
    aliases = {
        "vendor": ["vendor", "supplier", "name"],
        "unit_price": ["unit_price", "unit price", "price", "unit_cost"],
        "quantity": ["quantity", "qty", "units"],
        "delivery_days": ["delivery_days", "delivery days", "lead_time", "lead time", "days"],
        "warranty_months": ["warranty_months", "warranty months", "warranty", "months"],
        "currency": ["currency"],
    }
    lower = {str(k).strip().lower(): v for k, v in row.items()}
    out: dict[str, Any] = {}
    for target, keys in aliases.items():
        for key in keys:
            if key in lower and lower[key] not in (None, ""):
                out[target] = lower[key]
                break
    return out


def _parse_plain_text_quotes(text: str) -> list[QuoteInput]:
    quotes: list[QuoteInput] = []
    for line in text.splitlines():
        if not line.strip():
            continue
        parts = [part.strip() for part in re.split(r"[,;|]\s*", line) if part.strip()]
        if len(parts) >= 5:
            quotes.append(QuoteInput(vendor=parts[0], unit_price=float(parts[1]), quantity=int(float(parts[2])), delivery_days=int(float(parts[3])), warranty_months=int(float(parts[4]))))
    if not quotes:
        raise ValueError("plain text rows must be vendor, unit_price, quantity, delivery_days, warranty_months")
    return quotes