File size: 5,756 Bytes
df455fc
 
 
 
 
 
 
 
ba2cf79
df455fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ba2cf79
 
 
df455fc
 
 
 
 
 
 
 
 
ba2cf79
 
 
 
 
df455fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
from __future__ import annotations

import json
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

from .utils.reward_models import SourceDocument, SourcePack, TaskConstraints

DEFAULT_THEME = {
    "bg": "#F8FAFC",
    "surface": "#FFFFFF",
    "accent": "#2563EB",
    "primary": "#0F172A",
    "secondary": "#475569",
    "font": "Aptos",
    "title_size": 28,
    "body_size": 16,
    "caption_size": 10,
}

_DATA_PATH = Path(__file__).with_name("data.json")


@dataclass(frozen=True, slots=True)
class TaskScenario:
    task_id: str
    difficulty: str
    prompt_text: str
    source_pack: SourcePack
    task_constraints: TaskConstraints
    theme: dict[str, Any]


class TaskRegistry:
    def __init__(self, scenarios: Iterable[TaskScenario]):
        ordered_scenarios = tuple(scenarios)
        if not ordered_scenarios:
            raise ValueError("TaskRegistry requires at least one scenario")

        scenarios_by_id = {scenario.task_id: scenario for scenario in ordered_scenarios}
        if len(scenarios_by_id) != len(ordered_scenarios):
            raise ValueError("TaskRegistry task ids must be unique")

        self._ordered_scenarios = ordered_scenarios
        self._scenarios_by_id = scenarios_by_id

    def __len__(self) -> int:
        return len(self._ordered_scenarios)

    def all(self) -> tuple[TaskScenario, ...]:
        return self._ordered_scenarios

    def get(self, task_id: str) -> TaskScenario:
        try:
            return self._scenarios_by_id[task_id]
        except KeyError as error:
            raise KeyError(f"Unknown task id '{task_id}'") from error

    def by_difficulty(self, difficulty: str) -> tuple[TaskScenario, ...]:
        normalized = difficulty.strip().lower()
        matches = tuple(
            scenario
            for scenario in self._ordered_scenarios
            if scenario.difficulty.lower() == normalized
        )
        if not matches:
            raise KeyError(f"Unknown difficulty '{difficulty}'")
        return matches

    def sample(
        self,
        rng: random.Random,
        *,
        difficulty: str | None = None,
    ) -> TaskScenario:
        candidates = (
            self.by_difficulty(difficulty)
            if difficulty is not None
            else self._ordered_scenarios
        )
        return rng.choice(candidates)


def _load_raw_scenarios() -> list[dict[str, Any]]:
    payload = json.loads(_DATA_PATH.read_text(encoding="utf-8"))
    scenarios = payload.get("scenarios")
    if not isinstance(scenarios, list) or not scenarios:
        raise ValueError("server/data.json must contain a non-empty 'scenarios' list")
    return scenarios


def _build_source_document(payload: dict[str, Any]) -> SourceDocument:
    pages_payload = payload.get("pages")
    pages = None
    if isinstance(pages_payload, list):
        pages = [str(item).strip() for item in pages_payload if str(item).strip()]
        if not pages:
            pages = None

    return SourceDocument(
        doc_id=str(payload["doc_id"]),
        title=str(payload["title"]),
        path=payload.get("path"),
        mime_type=str(payload.get("mime_type", "text/plain")),
        text=str(payload["text"]).strip() if payload.get("text") is not None else None,
        pages=pages,
        images=payload.get("images"),
        metadata=(
            payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {}
        ),
    )


def _build_task_constraints(payload: dict[str, Any]) -> TaskConstraints:
    return TaskConstraints(
        min_slides=payload.get("min_slides"),
        max_slides=payload.get("max_slides"),
        target_audience=payload.get("target_audience"),
        tone=payload.get("tone"),
        extra_constraints=(
            payload.get("extra_constraints")
            if isinstance(payload.get("extra_constraints"), dict)
            else {}
        ),
    )


def _build_scenario(payload: dict[str, Any]) -> TaskScenario:
    task_id = str(payload["task_id"])
    difficulty = str(payload["difficulty"]).lower()
    source_pack_payload = payload.get("source_pack")
    source_documents = None
    source_pack_brief = None
    if isinstance(source_pack_payload, dict):
        source_documents = source_pack_payload.get("documents")
        brief = source_pack_payload.get("brief")
        if isinstance(brief, str) and brief.strip():
            source_pack_brief = brief.strip()
    else:
        source_documents = payload.get("source_documents")

    if not isinstance(source_documents, list) or not source_documents:
        raise ValueError(
            f"Scenario '{task_id}' must define source_pack.documents or source_documents"
        )

    source_pack = SourcePack(
        task_id=task_id,
        documents=[_build_source_document(item) for item in source_documents],
        brief=source_pack_brief,
        metadata={
            **(
                payload.get("metadata")
                if isinstance(payload.get("metadata"), dict)
                else {}
            ),
            "difficulty": difficulty,
        },
    )

    return TaskScenario(
        task_id=task_id,
        difficulty=difficulty,
        prompt_text=str(payload["prompt_text"]),
        source_pack=source_pack,
        task_constraints=_build_task_constraints(payload.get("task_constraints") or {}),
        theme={**DEFAULT_THEME, **dict(payload.get("theme") or {})},
    )


def _load_default_task_registry() -> TaskRegistry:
    return TaskRegistry(_build_scenario(item) for item in _load_raw_scenarios())


DEFAULT_TASK_REGISTRY = _load_default_task_registry()


__all__ = [
    "DEFAULT_TASK_REGISTRY",
    "DEFAULT_THEME",
    "TaskRegistry",
    "TaskScenario",
]