garvitsachdeva/SpindleFlow-RL

1"""2Task Bank — LLM-generated tasks derived from the specialist catalog.3 4Tasks are generated dynamically using GPT-4o-mini based on:5  1. The sector defined in training_config.yaml6  2. The specialist roster in specialist_catalog.yaml7  3. The current curriculum phase (controls complexity)8 9No hardcoded task lists. Any sector works by swapping the catalog + sector config.10"""11 12from __future__ import annotations13import random14import threading15import yaml16import os17from pathlib import Path18from dataclasses import dataclass19from typing import Optional20 21 22def _load_complexity_config(config_path: str) -> tuple[dict, dict]:23    """Load COMPLEXITY_BY_PHASE and COMPLEXITY_DESCRIPTIONS from config files."""24    import os25    base = os.path.dirname(os.path.abspath(config_path))26 27    with open(config_path) as f:28        cfg = yaml.safe_load(f)29    cur = cfg.get("curriculum", {})30    by_phase = {31        1: cur.get("phase1_task_types", ["atomic", "simple"]),32        2: cur.get("phase2_task_types", ["moderate"]),33        3: cur.get("phase3_task_types", ["complex", "enterprise"]),34    }35 36    desc_path = os.path.join(base, "complexity_descriptions.yaml")37    try:38        with open(desc_path) as f:39            descriptions = yaml.safe_load(f)40    except FileNotFoundError:41        descriptions = {42            "atomic": "a very simple, single-step",43            "simple": "a straightforward, well-scoped",44            "moderate": "a multi-component, realistic",45            "complex": "a complex, multi-system",46            "enterprise": "a large-scale, enterprise-grade",47        }48    return by_phase, descriptions49 50 51@dataclass52class Task:53    description: str54    complexity_class: str55    domain: str56 57 58class TaskBank:59    """60    Generates tasks dynamically using GPT-4o-mini.61    Falls back to catalog-derived tasks if OpenAI is unavailable.62 63    Tasks are pre-cached in batches to avoid per-episode API latency.64    """65 66    def __init__(67        self,68        phase: int = 1,69        config_path: str = "configs/training_config.yaml",70        catalog_path: str = "configs/specialist_catalog.yaml",71    ):72        self.phase = phase73        self._cache: list[Task] = []74        self._client = None75        self._cache_lock = threading.Lock()76        self._refill_running = False77 78        # Load complexity config from yaml files (not hardcoded)79        self._complexity_by_phase, self._complexity_descriptions = (80            _load_complexity_config(config_path)81        )82 83        # Load sector config84        with open(config_path) as f:85            cfg = yaml.safe_load(f)86        sector_cfg = cfg.get("sector", {})87        self.sector_name = sector_cfg.get("name", "software_engineering")88        self.sector_description = sector_cfg.get(89            "description",90            "Software product development"91        )92        self.use_llm = sector_cfg.get("use_llm_task_generation", True)93        self.llm_model = sector_cfg.get("llm_task_model", "gpt-4o-mini")94        self.cache_size = sector_cfg.get("task_cache_size", 50)95 96        # Load specialist roles from catalog (for context in prompts)97        with open(catalog_path) as f:98            catalog = yaml.safe_load(f)99        self._specialist_roles = [100            s["role"] for s in catalog.get("specialists", [])101        ]102 103        if self.use_llm:104            self._init_openai()105 106        # Pre-fill cache107        self._refill_cache()108 109    def _init_openai(self):110        try:111            from openai import OpenAI112            self._client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))113        except Exception as e:114            print(f"[TaskBank] OpenAI unavailable: {e}. Using catalog-derived tasks.")115            self._client = None116 117    def _refill_cache(self):118        """119        Synchronously generate a batch of tasks and extend the cache.120        Thread-safe: holds _cache_lock while writing; clears _refill_running on exit.121        Called directly on first fill (init) and from the background thread thereafter.122        """123        complexities = self._complexity_by_phase.get(self.phase, ["simple"])124        n_per_complexity = max(1, self.cache_size // len(complexities))125        new_tasks: list[Task] = []126 127        for complexity in complexities:128            if self._client and self.use_llm:129                batch = self._generate_llm_tasks(complexity, n_per_complexity)130            else:131                batch = self._generate_catalog_tasks(complexity, n_per_complexity)132            new_tasks.extend(batch)133 134        random.shuffle(new_tasks)135        with self._cache_lock:136            self._cache.extend(new_tasks)137            self._refill_running = False138 139    def _refill_cache_background(self):140        """Trigger a non-blocking background refill if one isn't already running."""141        with self._cache_lock:142            if self._refill_running:143                return          # already in flight — don't pile up threads144            self._refill_running = True145 146        t = threading.Thread(target=self._refill_cache, daemon=True)147        t.start()148 149    def _generate_llm_tasks(self, complexity: str, n: int) -> list[Task]:150        """Generate n tasks of the given complexity using GPT-4o-mini.151 152        Batches requests at max 20 tasks per API call to avoid JSON truncation153        from max_tokens limits. Results are concatenated into a single list.154        """155        complexity_desc = self._complexity_descriptions.get(complexity, "a realistic")156        roles_str = ", ".join(self._specialist_roles)157        batch_size = 20  # safe upper bound — 20 tasks × ~40 tokens each ≈ 800 tokens158        all_tasks: list[Task] = []159 160        for batch_start in range(0, n, batch_size):161            batch_n = min(batch_size, n - batch_start)162            prompt = f"""You are generating training tasks for a multi-agent RL environment.163 164Sector: {self.sector_name}165Sector description: {self.sector_description}166Available specialist roles: {roles_str}167 168Generate exactly {batch_n} different {complexity_desc} task descriptions for this sector.169Each task should:170- Be 1-2 sentences long171- Be specific and realistic for the {self.sector_name} sector172- Potentially require one or more of the available specialists to complete173- Vary in subject matter (don't repeat similar tasks)174 175Return ONLY a JSON array of strings, no other text:176["task 1 description", "task 2 description", ...]"""177 178            try:179                import json180                response = self._client.chat.completions.create(181                    model=self.llm_model,182                    max_tokens=1200,183                    messages=[{"role": "user", "content": prompt}],184                )185                raw = response.choices[0].message.content.strip()186                raw = raw.replace("```json", "").replace("```", "").strip()187                task_strings = json.loads(raw)188                all_tasks.extend([189                    Task(190                        description=t,191                        complexity_class=complexity,192                        domain=self.sector_name,193                    )194                    for t in task_strings195                    if isinstance(t, str) and len(t) > 10196                ])197            except Exception as e:198                print(f"[TaskBank] LLM generation failed for {complexity} batch: {e}. Using fallback.")199                all_tasks.extend(self._generate_catalog_tasks(complexity, batch_n))200 201        return all_tasks202 203    def _generate_catalog_tasks(self, complexity: str, n: int) -> list[Task]:204        """205        Fallback: derive tasks from specialist catalog without API calls.206        Produces formulaic but valid tasks for any sector.207        """208        complexity_desc = self._complexity_descriptions.get(complexity, "a realistic")209        tasks = []210        specialists = self._specialist_roles.copy()211        random.shuffle(specialists)212 213        for i in range(n):214            if len(specialists) >= 2:215                s1 = specialists[i % len(specialists)]216                s2 = specialists[(i + 1) % len(specialists)]217                desc = (218                    f"Design {complexity_desc} {self.sector_name} solution "219                    f"involving {s1} and {s2} working together"220                )221            else:222                s1 = specialists[0] if specialists else "specialist"223                desc = (224                    f"Create {complexity_desc} {self.sector_name} deliverable "225                    f"for a {s1}"226                )227            tasks.append(Task(228                description=desc,229                complexity_class=complexity,230                domain=self.sector_name,231            ))232        return tasks233 234    def sample(self) -> str:235        """236        Sample a random task description for a new episode.237 238        Never blocks for a refill.  When the cache drops below a low-water mark239        (10% of cache_size) a background thread is kicked off to replenish it.240        If the cache is completely empty (should only happen at init or after a241        phase switch drains it before the background fill completes) we fall back242        to a catalog-derived task immediately so reset() is never stalled.243        """244        low_water = max(5, self.cache_size // 10)245 246        with self._cache_lock:247            if self._cache:248                task = self._cache.pop()249            else:250                task = None251 252        if task is None:253            # Cache exhausted — generate one catalog task inline (fast, no API)254            fallback = self._generate_catalog_tasks(255                random.choice(self._complexity_by_phase.get(self.phase, ["simple"])), 1256            )257            task_desc = fallback[0].description if fallback else (258                f"Complete a {self.sector_name} task requiring specialist collaboration"259            )260            self._refill_cache_background()261            return task_desc262 263        with self._cache_lock:264            cache_len = len(self._cache)265 266        if cache_len < low_water:267            self._refill_cache_background()268 269        return task.description270 271    def sample_task(self) -> Task:272        """Sample a full Task object."""273        desc = self.sample()274        complexity = random.choice(self._complexity_by_phase.get(self.phase, ["simple"]))275        return Task(description=desc, complexity_class=complexity, domain=self.sector_name)276 277    def set_phase(self, phase: int) -> None:278        self.phase = phase279        with self._cache_lock:280            self._cache.clear()281            self._refill_running = False282        self._refill_cache()   # synchronous — phase switches are rare and intentional283 284    @property285    def pool_size(self) -> int:286        return len(self._cache)287