agent.py

import base64
import logging
import mimetypes
import re
import time
import tokenize
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, TypeVar
from urllib.parse import urlparse

ResultT = TypeVar("ResultT")
logger = logging.getLogger(__name__)
URL_PATTERN = re.compile(r"https?://[^\s<>\"']+", re.IGNORECASE)
IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp")
AUDIO_SUFFIXES = (".mp3", ".wav", ".m4a", ".ogg", ".flac", ".aac", ".webm")
MAX_PYTHON_SOURCE_BYTES = 1024 * 1024


def extract_urls(text: str) -> list[str]:
    return [match.rstrip(".,;:!?)]") for match in URL_PATTERN.findall(text)]


def is_youtube_url(url: str) -> bool:
    host = (urlparse(url).hostname or "").lower()
    return host in {"youtube.com", "www.youtube.com", "m.youtube.com", "youtu.be"}


@dataclass
class TaskContext:
    task_id: str
    question: str
    file_name: str | None = None
    file_url: str | None = None
    local_path: str | None = None
    diagnostics: list[str] = field(default_factory=list)


@dataclass
class TaskResult:
    task_id: str
    question: str
    submitted_answer: str
    route: str
    status: str
    diagnostics: str = ""


@dataclass
class SolverServices:
    text_search: Callable[[str], str]
    synthesize: Callable[[str, str], str]
    vision: Callable[[str, str], str] | None = None
    transcribe: Callable[[str], str] | None = None
    youtube: Callable[[str, str], str] | None = None


@dataclass
class EvaluationBatch:
    payload: dict[str, Any]
    results: list[TaskResult]


def route_task(context: TaskContext) -> str:
    """Select the deterministic solver route for a task."""
    urls = extract_urls(context.question)
    if any(is_youtube_url(url) for url in urls):
        return "youtube"
    file_name = (context.file_name or "").lower()
    if file_name.endswith(".py"):
        return "python"
    if file_name.endswith((".xlsx", ".xls", ".xlsm")):
        return "workbook"
    if file_name.endswith(IMAGE_SUFFIXES):
        return "vision"
    if file_name.endswith(AUDIO_SUFFIXES):
        return "audio"
    if any(urlparse(url).path.lower().endswith(IMAGE_SUFFIXES) for url in urls):
        return "vision"
    if any(urlparse(url).path.lower().endswith(AUDIO_SUFFIXES) for url in urls):
        return "audio"
    return "text_research"


def retry_call(
    operation: Callable[[], ResultT],
    *,
    attempts: int = 3,
    delay_seconds: float = 1,
    on_retry: Callable[[int, Exception], None] | None = None,
) -> ResultT:
    """Run an operation with a bounded number of attempts."""
    if attempts < 1:
        raise ValueError("attempts must be at least 1")

    for attempt in range(1, attempts + 1):
        try:
            return operation()
        except Exception as exc:
            if attempt == attempts:
                raise
            if on_retry:
                on_retry(attempt, exc)
            if delay_seconds:
                time.sleep(delay_seconds)

    raise RuntimeError("retry loop ended unexpectedly")


def normalize_answer(answer: object) -> str:
    """Return grader-safe answer text while preserving the answer's structure."""
    text = str(answer or "").strip()
    text = re.sub(
        r"^(?:\*\*|__)?(?:final\s+answer|answer|response)\s*:\s*(?:\*\*|__)?\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )
    kept_lines: list[str] = []
    for line in text.splitlines():
        stripped = line.strip()
        if re.match(r"^(?:explanation|reasoning|sources?|citations?)\s*:", stripped, re.I):
            break
        if re.fullmatch(r"\[\d+\]", stripped):
            continue
        kept_lines.append(re.sub(r"\s*\[\d+(?:\s*,\s*\d+)*\]", "", stripped))
    return "\n".join(kept_lines).strip()


def acquire_attachment(
    context: TaskContext,
    *,
    http_get: Callable[[str], Any],
    directory: str | Path,
    retry_attempts: int = 3,
    retry_delay_seconds: float = 1,
) -> TaskContext:
    """Download a task attachment and enrich its context before routing."""
    if not context.file_url or not context.file_name:
        return context
    file_url = context.file_url

    def record_retry(attempt: int, exc: Exception) -> None:
        context.diagnostics.append(
            f"download attempt {attempt} failed: {type(exc).__name__}: {exc}"
        )

    def fetch_attachment():
        response = http_get(file_url)
        response.raise_for_status()
        return response

    response = retry_call(
        fetch_attachment,
        attempts=retry_attempts,
        delay_seconds=retry_delay_seconds,
        on_retry=record_retry,
    )
    safe_task_id = re.sub(r"[^A-Za-z0-9_.-]+", "_", context.task_id)
    safe_file_name = Path(context.file_name).name
    destination = Path(directory) / f"{safe_task_id}-{safe_file_name}"
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_bytes(response.content)
    context.local_path = str(destination)
    return context


def inspect_workbook(local_path: str) -> str:
    """Render workbook sheets as bounded tabular evidence for synthesis."""
    import pandas as pd

    workbook = pd.ExcelFile(local_path)
    sections: list[str] = []
    for sheet_name in workbook.sheet_names:
        frame = pd.read_excel(workbook, sheet_name=sheet_name)
        sections.append(
            f"Sheet: {sheet_name}\nRows: {len(frame)}\n"
            f"Columns: {', '.join(map(str, frame.columns))}\n"
            f"Data:\n{frame.head(200).to_csv(index=False)}"
        )
    return "\n\n".join(sections)


def _inspect_python_source(local_path: str) -> str:
    """Read bounded Python source according to its declared encoding without executing it."""
    path = Path(local_path)
    if path.stat().st_size > MAX_PYTHON_SOURCE_BYTES:
        raise ValueError(f"Python source exceeds the {MAX_PYTHON_SOURCE_BYTES}-byte safety limit")
    with tokenize.open(path) as source_file:
        return source_file.read()


def solve_task(
    context: TaskContext,
    services: SolverServices,
    *,
    retry_attempts: int = 3,
    retry_delay_seconds: float = 1,
) -> TaskResult:
    """Solve one task through its deterministic route."""
    route = route_task(context)
    diagnostics = list(context.diagnostics)

    def record_retry(attempt: int, exc: Exception) -> None:
        diagnostics.append(f"attempt {attempt} failed: {type(exc).__name__}: {exc}")

    if route == "youtube":
        if not services.youtube:
            raise RuntimeError("YouTube provider is not configured")
        youtube_provider = services.youtube
        youtube_urls = [url for url in extract_urls(context.question) if is_youtube_url(url)]
        if not youtube_urls:
            raise ValueError("YouTube URL is unavailable")
        raw_answer = retry_call(
            lambda: youtube_provider(context.question, youtube_urls[0]),
            attempts=retry_attempts,
            delay_seconds=retry_delay_seconds,
            on_retry=record_retry,
        )
    elif route == "vision":
        if not services.vision:
            raise RuntimeError("vision provider is not configured")
        vision_provider = services.vision
        image_urls = [
            url
            for url in extract_urls(context.question)
            if urlparse(url).path.lower().endswith(IMAGE_SUFFIXES)
        ]
        if context.local_path:
            content_type = mimetypes.guess_type(context.file_name or "")[0] or "image/jpeg"
            encoded = base64.b64encode(Path(context.local_path).read_bytes()).decode("ascii")
            image_input = f"data:{content_type};base64,{encoded}"
        elif image_urls:
            image_input = image_urls[0]
        else:
            raise FileNotFoundError("image input is unavailable")
        raw_answer = retry_call(
            lambda: vision_provider(context.question, image_input),
            attempts=retry_attempts,
            delay_seconds=retry_delay_seconds,
            on_retry=record_retry,
        )
    elif route == "audio":
        if not services.transcribe:
            raise RuntimeError("audio transcription provider is not configured")
        transcribe_provider = services.transcribe
        audio_urls = [
            url
            for url in extract_urls(context.question)
            if urlparse(url).path.lower().endswith(AUDIO_SUFFIXES)
        ]
        if context.local_path:
            audio_input = context.local_path
        elif audio_urls:
            audio_input = audio_urls[0]
        else:
            raise FileNotFoundError("audio input is unavailable")
        evidence = retry_call(
            lambda: transcribe_provider(audio_input),
            attempts=retry_attempts,
            delay_seconds=retry_delay_seconds,
            on_retry=record_retry,
        )
    elif route == "workbook":
        if not context.local_path:
            raise FileNotFoundError("workbook attachment is unavailable")
        evidence = inspect_workbook(context.local_path)
    elif route == "python":
        if not context.local_path:
            raise FileNotFoundError("Python attachment is unavailable")
        evidence = _inspect_python_source(context.local_path)
    else:
        evidence = retry_call(
            lambda: services.text_search(context.question),
            attempts=retry_attempts,
            delay_seconds=retry_delay_seconds,
            on_retry=record_retry,
        )
    if route not in {"vision", "youtube"}:
        raw_answer = retry_call(
            lambda: services.synthesize(context.question, evidence),
            attempts=retry_attempts,
            delay_seconds=retry_delay_seconds,
            on_retry=record_retry,
        )
    return TaskResult(
        task_id=context.task_id,
        question=context.question,
        submitted_answer=normalize_answer(raw_answer),
        route=route,
        status="ok" if not diagnostics else "recovered",
        diagnostics="; ".join(diagnostics),
    )


def evaluate_items(
    items: Iterable[dict[str, Any]],
    *,
    username: str,
    agent_code: str,
    solve: Callable[[TaskContext], TaskResult],
    fallback: Callable[[TaskContext], object],
    prepare: Callable[[TaskContext], TaskContext] | None = None,
) -> EvaluationBatch:
    """Evaluate every valid task and build the course-compatible payload."""
    results: list[TaskResult] = []
    answers: list[dict[str, str]] = []
    for item in items:
        task_id = item.get("task_id")
        question = item.get("question")
        if not task_id or question is None:
            logger.warning("Skipping invalid task item: %r", item)
            continue
        context = TaskContext(
            task_id=str(task_id),
            question=str(question),
            file_name=item.get("file_name"),
            file_url=item.get("file_url"),
        )
        if prepare:
            try:
                context = prepare(context)
            except Exception as exc:
                logger.exception("Attachment preparation failed for task %s", task_id)
                context.diagnostics.append(
                    f"attachment preparation failed: {type(exc).__name__}: {exc}"
                )
        try:
            result = solve(context)
        except Exception as exc:
            logger.exception("Task %s failed; using fallback", task_id)
            try:
                fallback_answer = normalize_answer(fallback(context))
            except Exception as fallback_exc:
                fallback_answer = ""
                exc = RuntimeError(f"{exc}; fallback failed: {fallback_exc}")
            result = TaskResult(
                task_id=context.task_id,
                question=context.question,
                submitted_answer=fallback_answer,
                route=route_task(context),
                status="error",
                diagnostics=f"{type(exc).__name__}: {exc}",
            )
        if context.diagnostics:
            preparation_diagnostics = "; ".join(context.diagnostics)
            if preparation_diagnostics not in result.diagnostics:
                result.diagnostics = "; ".join(
                    part for part in [preparation_diagnostics, result.diagnostics] if part
                )
            if result.status == "ok":
                result.status = "degraded"
        result.submitted_answer = normalize_answer(result.submitted_answer)
        results.append(result)
        answers.append({"task_id": context.task_id, "submitted_answer": result.submitted_answer})
    return EvaluationBatch(
        payload={
            "username": username.strip(),
            "agent_code": agent_code,
            "answers": answers,
        },
        results=results,
    )


def build_default_services() -> SolverServices:
    """Create production provider adapters without exposing them to core logic."""
    import requests
    from langchain_community.tools import DuckDuckGoSearchResults
    from openai import OpenAI

    search = DuckDuckGoSearchResults(output_format="string", num_results=6)
    client = OpenAI(max_retries=0)

    def text_search(question: str) -> str:
        return str(search.invoke(question))

    def synthesize(question: str, evidence: str) -> str:
        response = client.responses.create(
            model="gpt-5.4-mini",
            instructions=(
                "Answer the question using the supplied evidence. Return only the "
                "exact final answer requested by the user, without labels, reasoning, "
                "citations, or surrounding prose. Preserve requested list formatting."
            ),
            input=f"Question:\n{question}\n\nEvidence:\n{evidence}",
        )
        return response.output_text

    def vision(question: str, image_input: str) -> str:
        vision_input: Any = [
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": question},
                    {"type": "input_image", "image_url": image_input},
                ],
            }
        ]
        response = client.responses.create(
            model="gpt-5.4-mini",
            instructions=(
                "Answer only the exact question about the image. Return the concise "
                "final answer without labels, reasoning, or citations."
            ),
            input=vision_input,
        )
        return response.output_text

    def transcribe(audio_input: str) -> str:
        if audio_input.startswith(("http://", "https://")):
            response = requests.get(audio_input, timeout=30)
            response.raise_for_status()
            filename = Path(urlparse(audio_input).path).name or "audio.mp3"
            content_type = response.headers.get("Content-Type", "audio/mpeg").split(";", 1)[0]
            file_input: Any = (filename, response.content, content_type)
            transcription = client.audio.transcriptions.create(
                model="gpt-4o-transcribe",
                file=file_input,
                response_format="text",
            )
        else:
            with open(audio_input, "rb") as audio_file:
                transcription = client.audio.transcriptions.create(
                    model="gpt-4o-transcribe",
                    file=audio_file,
                    response_format="text",
                )
        return getattr(transcription, "text", transcription)

    def youtube(question: str, video_url: str) -> str:
        from google import genai
        from google.genai import types

        gemini = genai.Client()
        response = gemini.models.generate_content(
            model="gemini-3.5-flash",
            contents=types.Content(
                parts=[
                    types.Part(file_data=types.FileData(file_uri=video_url)),
                    types.Part(text=question),
                ]
            ),
        )
        return response.text or ""

    return SolverServices(
        text_search=text_search,
        synthesize=synthesize,
        vision=vision,
        transcribe=transcribe,
        youtube=youtube,
    )