quote-judge-studio-v1/agent.py

"""QuoteJudge full-stack A2A agent.

Deterministic quote comparison, per-user persistence in managed Postgres,
browser-safe upload parsing, and production execution receipt storage.
"""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import re
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
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,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    State,
    WorkspaceAccess,
    WorkspaceMode,
)
from a2a_pack.workspace import FileUpload, UploadedFile

try:  # psycopg is installed in deployment and sandbox via requirements.txt.
    import psycopg
    from psycopg.rows import dict_row
except Exception:  # pragma: no cover - import guard for card-only tooling
    psycopg = None
    dict_row = None


MAX_QUOTES = 20
MAX_VENDOR_LENGTH = 120
MAX_UPLOADS = 5
MAX_UPLOAD_BYTES = 128_000
ALLOWED_UPLOAD_MEDIA_TYPES = {"text/csv", "text/plain", "application/json"}
_DATABASE_NAME = "quote-judge-studio-v1-data"
_LOCAL_COMPARISONS: dict[tuple[str, str], dict[str, Any]] = {}
_LOCAL_RECEIPTS: list[dict[str, Any]] = []


class QuoteJudgeStudioV1Config(BaseModel):
    """No caller-configured settings are required."""


class QuoteInput(BaseModel):
    vendor: str = Field(..., min_length=1, max_length=MAX_VENDOR_LENGTH)
    quantity: int = Field(..., gt=0, le=10_000_000)
    unit_price: float = Field(..., ge=0, le=1_000_000_000)
    delivery_days: int = Field(..., ge=0, le=3650)
    warranty_months: int = Field(..., ge=0, le=600)

    @field_validator("vendor")
    @classmethod
    def clean_vendor(cls, value: str) -> str:
        cleaned = " ".join(value.strip().split())
        if not cleaned:
            raise ValueError("vendor is required")
        return cleaned


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

    @field_validator("warranty")
    @classmethod
    def weights_must_be_positive(cls, value: float, info: Any) -> float:
        values = info.data
        total = float(values.get("price", 0)) + float(values.get("delivery", 0)) + float(value)
        if total <= 0:
            raise ValueError("at least one weight must be positive")
        return value


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

    @field_validator("filename")
    @classmethod
    def safe_filename(cls, value: str) -> str:
        cleaned = value.strip().replace("\\", "/").split("/")[-1]
        if not cleaned or cleaned in {".", ".."}:
            raise ValueError("filename is required")
        return cleaned[:160]


class QuoteJudgeState(BaseModel):
    """Durable state is stored in managed Postgres; this schema advertises that fact."""

    persistence: str = "managed_postgres"


class QuoteJudgeStudioV1(A2AAgent[QuoteJudgeStudioV1Config, PlatformUserAuth]):
    name = "quote-judge-studio-v1"
    description = (
        "QuoteJudge compares vendor quotes with weighted price, delivery, and "
        "warranty scoring, persists per-user comparisons, and serves a one-page product UI."
    )
    version = "0.1.0"

    config_model = QuoteJudgeStudioV1Config
    auth_model = PlatformUserAuth
    state = State.DURABLE
    state_model = QuoteJudgeState
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=300)
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name=_DATABASE_NAME,
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic local quote scoring; no LLM credential is required.",
    )
    workspace_access = WorkspaceAccess.dynamic(
        max_files=MAX_UPLOADS,
        allowed_modes=(WorkspaceMode.READ_ONLY,),
        require_reason=False,
        max_total_size_bytes=MAX_UPLOADS * MAX_UPLOAD_BYTES,
    )
    tools_used = ("postgres", "mcp")

    @a2a.tool(
        description=(
            "Compare at least two vendor quotes using typed price, delivery, and warranty "
            "weights; persist the recommendation and a production receipt."
        ),
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def compare_quotes(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        quotes: list[QuoteInput],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        cleaned_id = _clean_comparison_id(comparison_id)
        started = _now_iso()
        input_payload = {
            "comparison_id": cleaned_id,
            "quotes": [quote.model_dump() for quote in quotes],
            "weights": weights.model_dump(),
        }

        if len(quotes) < 2:
            result = {
                "ok": False,
                "code": "at_least_two_quotes_required",
                "message": "Add at least two vendor quotes before comparing.",
                "comparison_id": cleaned_id,
            }
            await _persist_receipt(
                tenant_key=tenant,
                comparison_id=cleaned_id,
                skill_name="compare_quotes",
                input_payload=input_payload,
                result_payload=result,
                status="validation_error",
                started_at=started,
            )
            return result
        if len(quotes) > MAX_QUOTES:
            result = {
                "ok": False,
                "code": "too_many_quotes",
                "message": f"Compare at most {MAX_QUOTES} quotes at a time.",
                "comparison_id": cleaned_id,
            }
            await _persist_receipt(
                tenant_key=tenant,
                comparison_id=cleaned_id,
                skill_name="compare_quotes",
                input_payload=input_payload,
                result_payload=result,
                status="validation_error",
                started_at=started,
            )
            return result

        result = _score_quotes(cleaned_id, quotes, weights)
        await _save_comparison(tenant, result)
        await _persist_receipt(
            tenant_key=tenant,
            comparison_id=cleaned_id,
            skill_name="compare_quotes",
            input_payload=input_payload,
            result_payload=result,
            status="ok",
            started_at=started,
        )
        return result

    @a2a.tool(
        description="Reopen a saved QuoteJudge comparison for the authenticated user.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="cheap",
    )
    async def get_comparison(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        cleaned_id = _clean_comparison_id(comparison_id)
        started = _now_iso()
        record = await _load_comparison(tenant, cleaned_id)
        if record is None:
            result = {
                "ok": False,
                "code": "comparison_not_found",
                "message": "No saved comparison exists for this id and user.",
                "comparison_id": cleaned_id,
            }
        else:
            result = dict(record)
            result["ok"] = True
        await _persist_receipt(
            tenant_key=tenant,
            comparison_id=cleaned_id,
            skill_name="get_comparison",
            input_payload={"comparison_id": cleaned_id},
            result_payload=result,
            status="ok" if result.get("ok") else "not_found",
            started_at=started,
        )
        return result

    @a2a.tool(
        description=(
            "Browser upload bridge: parse bounded base64 JSON/CSV/text quote files, "
            "compare extracted quotes, and persist the result."
        ),
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def compare_uploaded_quotes_browser(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        documents: list[BrowserQuoteDocument],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        parsed = _parse_browser_documents(documents)
        if not parsed["ok"]:
            return parsed
        return await self.compare_quotes(
            ctx,
            comparison_id=comparison_id,
            quotes=[QuoteInput.model_validate(item) for item in parsed["quotes"]],
            weights=weights,
        )

    @a2a.tool(
        description=(
            "External-client upload path: parse an uploaded CSV, JSON, or plain-text quote "
            "file and compare extracted quotes."
        ),
        timeout_seconds=30,
        idempotent=True,
        cost_class="cheap",
    )
    async def compare_uploaded_quote_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=("text/csv", "text/plain", "application/json"),
                max_bytes=MAX_UPLOAD_BYTES,
                description="CSV, JSON, or text file containing vendor quote rows.",
            ),
        ],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        if document.size_bytes > MAX_UPLOAD_BYTES:
            return {
                "ok": False,
                "code": "upload_too_large",
                "message": f"Upload must be {MAX_UPLOAD_BYTES} bytes or smaller.",
                "comparison_id": _clean_comparison_id(comparison_id),
            }
        if document.media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
            return {
                "ok": False,
                "code": "unsupported_media_type",
                "message": "Upload must be CSV, JSON, or plain text.",
                "comparison_id": _clean_comparison_id(comparison_id),
            }
        try:
            reader = getattr(ctx.workspace, "read_bytes")
            raw = reader(document.path)
        except Exception:
            return {
                "ok": False,
                "code": "upload_read_failed",
                "message": "Could not read the uploaded file from the granted workspace.",
                "comparison_id": _clean_comparison_id(comparison_id),
            }
        parsed = _parse_document_bytes(document.filename, document.media_type, raw)
        if not parsed["ok"]:
            return parsed
        return await self.compare_quotes(
            ctx,
            comparison_id=comparison_id,
            quotes=[QuoteInput.model_validate(item) for item in parsed["quotes"]],
            weights=weights,
        )


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


def _clean_comparison_id(value: str) -> str:
    cleaned = str(value or "").strip()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,79}", cleaned):
        raise ValueError("comparison_id must be 1-80 safe characters")
    return cleaned


def _score_sort_key(row: dict[str, Any]) -> tuple[float, float, int, int, str]:
    return (
        -float(row["scores"]["weighted_total"]),
        float(row["subtotal"]),
        int(row["delivery_days"]),
        -int(row["warranty_months"]),
        str(row["vendor"]).lower(),
    )


def _subtotal_key(row: dict[str, Any]) -> float:
    return float(row["subtotal"])


def _score_quotes(comparison_id: str, quotes: list[QuoteInput], weights: QuoteWeights) -> dict[str, Any]:
    quote_rows = [quote.model_dump() for quote in quotes]
    for row in quote_rows:
        row["subtotal"] = round(row["quantity"] * row["unit_price"], 2)

    totals = [row["subtotal"] for row in quote_rows]
    deliveries = [row["delivery_days"] for row in quote_rows]
    warranties = [row["warranty_months"] for row in quote_rows]
    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,
    }

    scored: list[dict[str, Any]] = []
    for row in quote_rows:
        price_score = _lower_is_better(row["subtotal"], totals)
        delivery_score = _lower_is_better(row["delivery_days"], deliveries)
        warranty_score = _higher_is_better(row["warranty_months"], warranties)
        weighted_score = (
            price_score * normalized_weights["price"]
            + delivery_score * normalized_weights["delivery"]
            + warranty_score * normalized_weights["warranty"]
        )
        scored.append(
            {
                **row,
                "scores": {
                    "price": round(price_score, 4),
                    "delivery": round(delivery_score, 4),
                    "warranty": round(warranty_score, 4),
                    "weighted_total": round(weighted_score, 4),
                },
            }
        )

    scored.sort(
        key=lambda item: (
            -item["scores"]["weighted_total"],
            item["subtotal"],
            item["delivery_days"],
            -item["warranty_months"],
            item["vendor"].lower(),
        )
    )
    winner = scored[0]
    baseline = min(scored, key=lambda item: item["subtotal"])
    rationale_bits = [
        f"{winner['vendor']} has the highest weighted score ({winner['scores']['weighted_total']:.2f}).",
        f"Its normalized total is {winner['subtotal']:.2f} for {winner['quantity']} units.",
    ]
    if winner["delivery_days"] == min(deliveries):
        rationale_bits.append("It also has the fastest delivery among the submitted quotes.")
    if winner["warranty_months"] == max(warranties):
        rationale_bits.append("It offers the strongest warranty among the submitted quotes.")
    if winner["vendor"] == baseline["vendor"]:
        rationale_bits.append("It is the lowest-price option under the selected quantity.")

    return {
        "ok": True,
        "comparison_id": comparison_id,
        "recommendation": {
            "vendor": winner["vendor"],
            "score": winner["scores"]["weighted_total"],
            "normalized_total": winner["subtotal"],
            "rationale": " ".join(rationale_bits),
        },
        "weights": {
            "price": weights.price,
            "delivery": weights.delivery,
            "warranty": weights.warranty,
            "normalized": {key: round(value, 4) for key, value in normalized_weights.items()},
        },
        "quotes": scored,
        "created_at": _now_iso(),
    }


def _lower_is_better(value: float, values: list[float]) -> float:
    min_value = min(values)
    max_value = max(values)
    if max_value == min_value:
        return 1.0
    return (max_value - value) / (max_value - min_value)


def _higher_is_better(value: float, values: list[float]) -> float:
    min_value = min(values)
    max_value = max(values)
    if max_value == min_value:
        return 1.0
    return (value - min_value) / (max_value - min_value)


def _parse_browser_documents(documents: list[BrowserQuoteDocument]) -> dict[str, Any]:
    if not documents:
        return {"ok": False, "code": "no_uploads", "message": "Upload at least one quote file."}
    if len(documents) > MAX_UPLOADS:
        return {"ok": False, "code": "too_many_uploads", "message": f"Upload at most {MAX_UPLOADS} files."}
    quotes: list[dict[str, Any]] = []
    for document in documents:
        if document.media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
            return {"ok": False, "code": "unsupported_media_type", "message": "Upload must be CSV, JSON, or plain text."}
        try:
            raw = base64.b64decode(document.data_base64, validate=True)
        except Exception:
            return {"ok": False, "code": "invalid_base64", "message": "Uploaded file data is not valid base64."}
        if len(raw) > MAX_UPLOAD_BYTES:
            return {"ok": False, "code": "upload_too_large", "message": f"Each upload must be {MAX_UPLOAD_BYTES} bytes or smaller."}
        parsed = _parse_document_bytes(document.filename, document.media_type, raw)
        if not parsed["ok"]:
            return parsed
        quotes.extend(parsed["quotes"])
    return {"ok": True, "quotes": quotes[:MAX_QUOTES]}


def _parse_document_bytes(filename: str, media_type: str, raw: bytes) -> dict[str, Any]:
    if len(raw) > MAX_UPLOAD_BYTES:
        return {"ok": False, "code": "upload_too_large", "message": f"Upload must be {MAX_UPLOAD_BYTES} bytes or smaller."}
    text = raw.decode("utf-8-sig", errors="replace")
    try:
        if media_type == "application/json" or filename.lower().endswith(".json"):
            data = json.loads(text)
            rows = data.get("quotes", data) if isinstance(data, dict) else data
            if not isinstance(rows, list):
                raise ValueError("JSON must be a list or contain a quotes list")
            quotes = [QuoteInput.model_validate(row).model_dump() for row in rows]
        elif media_type == "text/csv" or filename.lower().endswith(".csv"):
            reader = csv.DictReader(io.StringIO(text))
            quotes = [QuoteInput.model_validate(_coerce_quote_row(row)).model_dump() for row in reader]
        else:
            quotes = _parse_plain_text_quotes(text)
    except Exception as exc:
        return {"ok": False, "code": "quote_parse_failed", "message": f"Could not parse quote upload: {exc}"}
    if not quotes:
        return {"ok": False, "code": "no_quotes_found", "message": "No quote rows were found in the upload."}
    return {"ok": True, "quotes": quotes}


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


def _parse_plain_text_quotes(text: str) -> list[dict[str, Any]]:
    quotes: list[dict[str, Any]] = []
    for line in text.splitlines():
        if not line.strip():
            continue
        parts = [part.strip() for part in re.split(r"[,;|]", line) if part.strip()]
        if len(parts) >= 5:
            quotes.append(
                QuoteInput.model_validate(
                    {
                        "vendor": parts[0],
                        "quantity": parts[1],
                        "unit_price": parts[2],
                        "delivery_days": parts[3],
                        "warranty_months": parts[4],
                    }
                ).model_dump()
            )
    return quotes


@contextmanager
def _db_conn() -> Any:
    url = os.environ.get("DATABASE_URL")
    if not url or psycopg is None:
        yield None
        return
    conn = psycopg.connect(url, row_factory=dict_row)
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()


async def _save_comparison(tenant_key: str, payload: dict[str, Any]) -> None:
    key = (tenant_key, payload["comparison_id"])
    with _db_conn() as conn:
        if conn is None:
            _LOCAL_COMPARISONS[key] = dict(payload)
            return
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO quote_comparisons
                    (tenant_key, comparison_id, recommended_vendor, score, normalized_total, payload, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, NOW())
                ON CONFLICT (tenant_key, comparison_id) DO UPDATE SET
                    recommended_vendor = EXCLUDED.recommended_vendor,
                    score = EXCLUDED.score,
                    normalized_total = EXCLUDED.normalized_total,
                    payload = EXCLUDED.payload,
                    updated_at = NOW()
                """,
                (
                    tenant_key,
                    payload["comparison_id"],
                    payload["recommendation"]["vendor"],
                    payload["recommendation"]["score"],
                    payload["recommendation"]["normalized_total"],
                    json.dumps(payload),
                ),
            )


async def _load_comparison(tenant_key: str, comparison_id: str) -> dict[str, Any] | None:
    key = (tenant_key, comparison_id)
    with _db_conn() as conn:
        if conn is None:
            return _LOCAL_COMPARISONS.get(key)
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT payload
                FROM quote_comparisons
                WHERE tenant_key = %s AND comparison_id = %s
                """,
                (tenant_key, comparison_id),
            )
            row = cur.fetchone()
    if not row:
        return None
    payload = row["payload"]
    return payload if isinstance(payload, dict) else json.loads(payload)


async def _persist_receipt(
    *,
    tenant_key: str,
    comparison_id: str,
    skill_name: str,
    input_payload: dict[str, Any],
    result_payload: dict[str, Any],
    status: str,
    started_at: str,
) -> dict[str, Any]:
    receipt = {
        "receipt_id": str(uuid.uuid4()),
        "agent": QuoteJudgeStudioV1.name,
        "agent_version": QuoteJudgeStudioV1.version,
        "skill": skill_name,
        "comparison_id": comparison_id,
        "status": status,
        "started_at": started_at,
        "completed_at": _now_iso(),
        "input_hash": _hash_json(input_payload),
        "result_hash": _hash_json(result_payload),
    }
    with _db_conn() as conn:
        if conn is None:
            _LOCAL_RECEIPTS.append({"tenant_key": tenant_key, **receipt})
            return receipt
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO quote_execution_receipts
                    (tenant_key, comparison_id, receipt_id, skill_name, status, input_hash,
                     result_hash, payload, created_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, NOW())
                """,
                (
                    tenant_key,
                    comparison_id,
                    receipt["receipt_id"],
                    skill_name,
                    status,
                    receipt["input_hash"],
                    receipt["result_hash"],
                    json.dumps(receipt),
                ),
            )
    return receipt


def _hash_json(payload: dict[str, Any]) -> str:
    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()