Support resolution agent

import asyncio
import json
import os
from dataclasses import dataclass, field

import flyte

MODEL = "anthropic/claude-haiku-4-5"

env = flyte.TaskEnvironment(
    name="support-resolution",
    secrets=[
        flyte.Secret(key="youdotcom-api-key", as_env_var="YOU_API_KEY"),
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
    image=flyte.Image.from_uv_script(__file__, name="support-resolution", pre=True),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

@dataclass
class Source:
    title: str
    url: str
    snippet: str
    domain: str = ""
    favicon: str = ""


def _domain(url: str) -> str:
    from urllib.parse import urlparse

    try:
        return urlparse(url).netloc.replace("www.", "")
    except Exception:
        return ""


def _favicon_for(url: str) -> str:
    return f"https://ydc-index.io/favicon?domain={_domain(url)}&size=128"


@dataclass
class Ticket:
    ticket_id: str
    question: str
    context: str = ""


@dataclass
class Grounding:
    answer: str
    sources: list[Source] = field(default_factory=list)


@dataclass
class Resolution:
    ticket_id: str
    ticket: str
    grounded_answer: str
    draft_reply: str
    sources: list[Source] = field(default_factory=list)


@dataclass
class ResolutionReport:
    resolutions: list[Resolution] = field(default_factory=list)

YOU_RESEARCH_URL = "https://api.you.com/v1/research"


async def _you_post(url: str, body: dict, timeout: float = 120.0) -> dict:
    """POST with exponential backoff + jitter on 429 rate limits."""
    import random

    import httpx

    headers = {
        "X-API-Key": os.environ["YOU_API_KEY"],
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(timeout=timeout) as client:
        for attempt in range(7):
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 429 and attempt < 6:
                wait = float(resp.headers.get("retry-after") or 0) or min(2**attempt, 30)
                await asyncio.sleep(wait + random.uniform(0, 2))
                continue
            resp.raise_for_status()
            return resp.json()
    resp.raise_for_status()
    return resp.json()


@flyte.trace
async def you_research(question: str, research_effort: str = "lite") -> dict:
    """Fast, citation-backed grounding for a support question."""
    body = {"input": question, "research_effort": research_effort}
    return await _you_post(YOU_RESEARCH_URL, body)

@env.task(retries=3)
async def ground_answer(ticket: str, context: str, research_effort: str) -> Grounding:
    """Ground the ticket in fresh public sources via the Research API."""
    question = ticket if not context else f"{ticket}\n\nContext: {context}"
    result = await you_research(question, research_effort)

    output = result.get("output", {})
    answer = output.get("content", "")
    if not isinstance(answer, str):
        answer = json.dumps(answer)

    sources = []
    for s in output.get("sources", []) or []:
        url = str(s.get("url", ""))
        sources.append(
            Source(
                title=str(s.get("title", "") or url),
                url=url,
                snippet=str((s.get("snippets") or [""])[0]),
                domain=_domain(url),
                favicon=_favicon_for(url),
            )
        )
    return Grounding(answer=answer, sources=sources)

@flyte.trace
async def _draft(ticket: str, answer: str, sources_text: str) -> str:
    from litellm import acompletion

    system = (
        "You are a senior customer-support agent. Using ONLY the grounded "
        "answer and sources provided, draft a concise, friendly, customer-ready "
        "reply. Cite the relevant source URL inline in parentheses after any "
        "factual claim so a human agent can verify before sending. If the "
        "sources do not answer the question, say so plainly."
    )
    user = (
        f"Customer ticket: {ticket}\n\n"
        f"Grounded answer:\n{answer}\n\nSources:\n{sources_text}"
    )
    resp = await acompletion(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0.2,
        max_tokens=1024,
    )
    return resp.choices[0].message.content


@env.task
async def draft_reply(ticket: Ticket, grounding: Grounding) -> Resolution:
    """Turn the grounded answer into a cited, customer-ready reply."""
    sources_text = "\n".join(
        f"- {s.title} ({s.domain}): {s.url}\n  \"{s.snippet}\""
        for s in grounding.sources
    )
    reply = await _draft(ticket.question, grounding.answer, sources_text)

    return Resolution(
        ticket_id=ticket.ticket_id,
        ticket=ticket.question,
        grounded_answer=grounding.answer,
        draft_reply=reply,
        sources=grounding.sources,
    )

async def resolve_ticket(ticket: Ticket, research_effort: str) -> Resolution:
    """Ground one ticket then draft its reply."""
    grounding = await ground_answer(ticket.question, ticket.context, research_effort)
    return await draft_reply(ticket, grounding)

def _default_tickets() -> list[Ticket]:
    return [
        Ticket(
            "tkt-1",
            "Is there a recall on the DeWalt DCD777 cordless drill, and what should "
            "the customer do if there is?",
            "Customer purchased the drill recently and is asking about safety recalls.",
        ),
        Ticket(
            "tkt-2",
            "What is Sony's current return policy for the WH-1000XM5 headphones?",
            "Customer wants to return an opened pair bought 20 days ago.",
        ),
        Ticket(
            "tkt-3",
            "Are there any current weather advisories that could delay flights out of "
            "Denver International Airport today?",
            "Customer is worried about a connecting flight.",
        ),
        Ticket(
            "tkt-4",
            "What are the dimensions and weight capacity of the IKEA BEKANT desk?",
            "Customer is checking if it fits their space before resolving a complaint.",
        ),
        Ticket(
            "tkt-5",
            "Has Samsung issued any recall or safety notice for the Galaxy Z Fold5?",
            "Customer reports overheating and wants to know about known issues.",
        ),
        Ticket(
            "tkt-6",
            "What is the warranty period for a Dyson V15 Detect vacuum in the US?",
            "Customer's vacuum stopped working and asks about coverage.",
        ),
    ]


@env.task(report=True)
async def support_resolution(
    tickets: list[Ticket] | None = None,
    research_effort: str = "lite",
) -> ResolutionReport:
    """Fan out across support tickets, grounding and drafting cited replies."""
    if tickets is None:
        tickets = _default_tickets()

    with flyte.group("resolve-tickets"):
        resolutions = await asyncio.gather(
            *[resolve_ticket(t, research_effort) for t in tickets]
        )

    report = ResolutionReport(resolutions=list(resolutions))
    await flyte.report.replace.aio(_render_report(report), do_flush=True)
    await flyte.report.flush.aio()
    return report