proposal-system/lambdas/suggestions/app.py

"""Proposal System - Suggestion Engine Lambda.

Queries Bedrock Knowledge Base for similar proposals and invokes Claude
to generate line item suggestions for new proposals.
"""

import json
import logging
import os
import time

import boto3
import httpx

logger = logging.getLogger(__name__)
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))

KNOWLEDGE_BASE_ID = os.environ.get("KNOWLEDGE_BASE_ID", "")
MODEL_ID = os.environ.get("MODEL_ID", "us.anthropic.claude-sonnet-4-5-20250929-v1:0")
API_BASE_URL = os.environ.get("API_BASE_URL", "")
INTERNAL_API_KEY_SECRET_ARN = os.environ.get("INTERNAL_API_KEY_SECRET_ARN", "")

bedrock_agent = boto3.client("bedrock-agent-runtime")
bedrock_runtime = boto3.client("bedrock-runtime")
secrets_client = boto3.client("secretsmanager")

_cached_api_key: str | None = None


def _get_api_key() -> str:
    global _cached_api_key
    if _cached_api_key is None:
        if INTERNAL_API_KEY_SECRET_ARN:
            resp = secrets_client.get_secret_value(SecretId=INTERNAL_API_KEY_SECRET_ARN)
            _cached_api_key = resp["SecretString"]
        else:
            _cached_api_key = ""
    return _cached_api_key


def handler(event, context):
    batch_item_failures = []
    for record in event.get("Records", []):
        try:
            body = json.loads(record["body"])
            payload = body.get("payload", body)
            proposal_id = payload["proposalId"]
            trigger = payload.get("trigger", "generate")
            process_suggestion(proposal_id, trigger)
        except Exception as e:
            logger.error("Failed to process record %s: %s", record.get("messageId"), e)
            batch_item_failures.append({"itemIdentifier": record["messageId"]})
    return {"batchItemFailures": batch_item_failures}


def process_suggestion(proposal_id: str, trigger: str):
    proposal = fetch_proposal(proposal_id)
    if not proposal:
        logger.warning("Proposal %s not found", proposal_id)
        return

    scope = proposal.get("refinedScope") or proposal.get("scopeOfWork", "")
    category = proposal.get("serviceCategory", "")
    priority = proposal.get("priority", "")

    existing_items = fetch_line_items(proposal_id)

    similar_proposals = retrieve_similar(scope, category)

    suggested_items = generate_line_items(scope, category, priority, similar_proposals)

    post_line_items(proposal_id, suggested_items, existing_items)

    store_similar_references(proposal_id, similar_proposals)

    update_status_to_in_review(proposal_id)


def fetch_line_items(proposal_id: str) -> list[dict]:
    try:
        resp = _retry_request(
            "GET",
            f"{API_BASE_URL}/api/proposals/{proposal_id}/line-items",
            headers=_api_headers(),
        )
        if resp.status_code == 200:
            return resp.json()
    except Exception as e:
        logger.error("Error fetching line items: %s", e)
    return []


def fetch_proposal(proposal_id: str) -> dict | None:
    try:
        resp = _retry_request(
            "GET",
            f"{API_BASE_URL}/api/proposals/{proposal_id}",
            headers=_api_headers(),
        )
        if resp.status_code == 200:
            return resp.json()
    except Exception as e:
        logger.error("Error fetching proposal: %s", e)
    return None


def retrieve_similar(scope: str, category: str) -> list[dict]:
    if not KNOWLEDGE_BASE_ID:
        logger.info("No Knowledge Base configured, skipping retrieval")
        return []

    try:
        filter_config = (
            {"equals": {"key": "service_category", "value": category}}
            if category
            else None
        )

        params = {
            "knowledgeBaseId": KNOWLEDGE_BASE_ID,
            "retrievalQuery": {"text": scope},
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "numberOfResults": 10,
                }
            },
        }

        if filter_config:
            params["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"] = (
                filter_config
            )

        response = bedrock_agent.retrieve(**params)

        results = []
        for result in response.get("retrievalResults", []):
            content = result.get("content", {}).get("text", "")
            score = result.get("score", 0.0)
            metadata = result.get("metadata", {})
            source_uri = result.get("location", {}).get("s3Location", {}).get("uri", "")

            results.append(
                {
                    "content": content,
                    "score": score,
                    "metadata": metadata,
                    "sourceUri": source_uri,
                }
            )

        return results

    except Exception as e:
        logger.error("Error retrieving from KB: %s", e)
        return []


def generate_line_items(
    scope: str,
    category: str,
    priority: str,
    similar_proposals: list[dict],
) -> list[dict]:
    context_block = ""
    if similar_proposals:
        context_block = "Here are similar historical proposals and their line items for reference:\n\n"
        for i, sp in enumerate(similar_proposals[:5], 1):
            context_block += (
                f"--- Similar Proposal {i} (relevance: {sp['score']:.2f}) ---\n"
            )
            context_block += sp["content"] + "\n\n"

    prompt = f"""You are a construction/facilities proposal estimator for Sea Haven Industries.
Based on the scope of work and similar historical proposals, generate a detailed list of line items
with quantities, units, and estimated pricing.

Service Category: {category}
Priority: {priority}

Scope of Work:
{scope}

{context_block}

Generate line items as a JSON array. Each item should have:
- description: clear description of the work/material
- quantity: numeric quantity
- unit: unit of measurement (e.g., "sq ft", "hours", "each", "linear ft")
- unitPrice: price per unit in dollars (or null if lump sum)
- totalPrice: total price for this line item in dollars
- pricingMode: "UnitPrice" if unit price provided, "TotalPrice" if lump sum

Respond ONLY with the JSON array, no additional text."""

    try:
        response = bedrock_runtime.invoke_model(
            modelId=MODEL_ID,
            contentType="application/json",
            accept="application/json",
            body=json.dumps(
                {
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": 4096,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.3,
                }
            ),
        )

        response_body = json.loads(response["body"].read())
        content = response_body["content"][0]["text"]

        content = content.strip()
        if content.startswith("```"):
            content = content.split("\n", 1)[1]
            content = content.rsplit("```", 1)[0]

        line_items = json.loads(content)
        return line_items if isinstance(line_items, list) else []

    except Exception as e:
        logger.error("Error generating line items: %s", e)
        return []


def post_line_items(proposal_id: str, items: list[dict], existing_items: list[dict]):
    if not items and not existing_items:
        return

    line_items_payload = []

    # Preserve non-AI items (Manual, Vendor, Historical)
    preserved = [li for li in existing_items if li.get("source") != "AI"]
    for i, li in enumerate(preserved):
        line_items_payload.append(
            {
                "id": li.get("id"),
                "description": li["description"],
                "quantity": float(li.get("quantity", 1)),
                "unit": li.get("unit", "each"),
                "unitPrice": li.get("unitPrice"),
                "totalPrice": float(li.get("totalPrice", 0)),
                "pricingMode": li.get("pricingMode", "TotalPrice"),
                "sortOrder": i + 1,
                "source": li.get("source", "Manual"),
            }
        )

    # Add new AI-generated items after preserved ones
    offset = len(line_items_payload)
    for i, item in enumerate(items):
        pricing_mode = item.get("pricingMode", "TotalPrice")
        if pricing_mode not in ("UnitPrice", "TotalPrice", "Both"):
            pricing_mode = "UnitPrice" if item.get("unitPrice") else "TotalPrice"

        line_items_payload.append(
            {
                "id": None,
                "description": item["description"],
                "quantity": float(item.get("quantity", 1)),
                "unit": item.get("unit", "each"),
                "unitPrice": item.get("unitPrice"),
                "totalPrice": float(item.get("totalPrice", 0)),
                "pricingMode": pricing_mode,
                "sortOrder": offset + i + 1,
                "source": "AI",
            }
        )

    try:
        resp = _retry_request(
            "PUT",
            f"{API_BASE_URL}/api/proposals/{proposal_id}/line-items",
            json={"lineItems": line_items_payload},
            headers=_api_headers(),
            timeout=15,
        )
        if resp.status_code not in (200, 201):
            logger.error(
                "Failed to post line items: %s %s", resp.status_code, resp.text
            )
    except Exception as e:
        logger.error("Error posting line items: %s", e)


def store_similar_references(proposal_id: str, similar_proposals: list[dict]):
    if not similar_proposals:
        return

    for sp in similar_proposals[:5]:
        source_uri = sp.get("sourceUri", "")
        library_item_id = source_uri.split("/")[-1] if source_uri else ""
        if not library_item_id:
            continue

        try:
            _retry_request(
                "POST",
                f"{API_BASE_URL}/api/proposals/{proposal_id}/similar-references",
                json={
                    "referencedLibraryItemId": library_item_id,
                    "similarityScore": sp["score"],
                },
                headers=_api_headers(),
            )
        except Exception as e:
            logger.error("Error storing similar reference: %s", e)


def update_status_to_in_review(proposal_id: str):
    try:
        _retry_request(
            "PUT",
            f"{API_BASE_URL}/api/proposals/{proposal_id}",
            json={"status": "InReview"},
            headers=_api_headers(),
        )
    except Exception as e:
        logger.error("Error updating status: %s", e)


def _api_headers() -> dict:
    headers = {"Content-Type": "application/json"}
    api_key = _get_api_key()
    if api_key:
        headers["X-Internal-Api-Key"] = api_key
    return headers


def _retry_request(
    method: str, url: str, *, max_retries: int = 3, **kwargs
) -> httpx.Response:
    kwargs.setdefault("timeout", 10)
    for attempt in range(max_retries):
        try:
            resp = httpx.request(method, url, **kwargs)
            if resp.status_code < 500:
                return resp
        except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as exc:
            if attempt == max_retries - 1:
                raise
            logger.warning(
                "Retryable error (attempt %d/%d): %s", attempt + 1, max_retries, exc
            )
        time.sleep(min(2**attempt, 4))
    return resp  # type: ignore[possibly-undefined]