from __future__ import annotations import logging import random from dataclasses import dataclass from pathlib import Path from typing import Any from uuid import uuid4 from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State from ..models import PptAgentAction, PptAgentObservation from .debug_logging import debug_context, write_debug_event from .llm_client import LLMClient from .task_registry import DEFAULT_TASK_REGISTRY, TaskRegistry, TaskScenario from .tools.pptx_tools import ( create_slide, delete_slide, register_theme, save_presentation, update_slide, update_theme, ) from .utils.pptx_extraction import PptxExtractionService from .utils.pptx_functions import PptxEditor from .utils.rendering.pptx_render_service import PptxRenderService from .utils.reward_kernel import ( build_eval_spec, evaluate_presentation, evaluate_slide, ) from .utils.reward_metrics import clamp_reward from .utils.reward_models import EvalSpec, ExtractedPresentation, to_serializable from .utils.slidesgenbench import ( QuantitativeQuizJudgeService, QuizBankGenerationService, SlidesGenQuantitativeJudgeService, SlidesGenQuizBankService, ) _DEFAULT_MAX_STEPS = 30 _INVALID_ACTION_PENALTY = 0.0 logger = logging.getLogger(__name__) @dataclass(frozen=True, slots=True) class _ActionExecutionResult: action_type: str tool_result: dict[str, Any] affected_slide_index: int | None = None reward: float = 0.0 class PptAgentEnvironment(Environment): """Prompt-to-PPT environment backed by ``PptxEditor`` and reward kernel.""" SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__( self, max_steps: int = _DEFAULT_MAX_STEPS, *, task_registry: TaskRegistry | None = None, quiz_bank_service: QuizBankGenerationService | None = None, quantitative_quiz_judge_service: QuantitativeQuizJudgeService | None = None, render_service: Any | None = None, ): super().__init__() self._max_steps = max_steps self._rng = random.Random() self._inspection_service = PptxExtractionService() self._task_registry = task_registry or DEFAULT_TASK_REGISTRY self._scenario_queue: list[TaskScenario] = [] self._eval_spec_cache: dict[str, EvalSpec] = {} llm_client: LLMClient | None = None if quiz_bank_service is None or quantitative_quiz_judge_service is None: llm_client = LLMClient() self._quiz_bank_service = quiz_bank_service or SlidesGenQuizBankService( llm_client ) self._quantitative_quiz_judge_service = ( quantitative_quiz_judge_service or SlidesGenQuantitativeJudgeService(llm_client) ) self._render_service = render_service or PptxRenderService() self._state = State(episode_id=str(uuid4()), step_count=0) self._done = False self._termination_reason: str | None = None self._scenario: TaskScenario | None = None self._eval_spec: EvalSpec | None = None self._editor: PptxEditor | None = None self._last_score = 0.0 self._last_action_result: dict[str, Any] | None = None self._last_reward_details: dict[str, Any] | None = None self._known_named_shapes_by_slide: dict[int, dict[str, int]] = {} def reset( self, seed: int | None = None, episode_id: str | None = None, **kwargs: Any, ) -> PptAgentObservation: requested_task_id = kwargs.get("task_id") requested_difficulty = kwargs.get("difficulty") next_episode_id = episode_id or str(uuid4()) with debug_context(episode_id=next_episode_id): logger.info("env.reset start episode_id=%s seed=%s", next_episode_id, seed) write_debug_event( "env.reset.start", { "seed": seed, "requested_task_id": requested_task_id, "requested_difficulty": requested_difficulty, }, ) self._reset_rubric() if seed is not None: self._rng.seed(seed) self._scenario_queue = [] self._done = False self._termination_reason = None self._last_score = 0.0 self._last_action_result = None self._last_reward_details = None self._known_named_shapes_by_slide = {} self._scenario = self._resolve_scenario( task_id=requested_task_id, difficulty=requested_difficulty, ) with debug_context( task_id=self._scenario.task_id, difficulty=self._scenario.difficulty, ): logger.info("env.reset sampled task task_id=%s", self._scenario.task_id) write_debug_event( "env.reset.scenario", { "task_id": self._scenario.task_id, "difficulty": self._scenario.difficulty, }, ) self._eval_spec = self._build_eval_spec(self._scenario) self._editor = self._initialize_editor(self._scenario.theme) self._state = State(episode_id=next_episode_id, step_count=0) logger.info( "env.reset ready episode_id=%s task_id=%s max_steps=%s", self._state.episode_id, self._scenario.task_id, self._max_steps, ) write_debug_event( "env.reset.ready", { "max_steps": self._max_steps, "spec_hash": self._eval_spec.spec_hash, }, ) return self._build_observation( reward=0.0, done=False, last_action_error=None, ) def step( self, action: PptAgentAction, timeout_s: float | None = None, **kwargs: Any, ) -> PptAgentObservation: del timeout_s, kwargs if self._editor is None or self._scenario is None or self._eval_spec is None: raise RuntimeError("Environment must be reset before step()") if self._done: return self._build_observation( reward=0.0, done=True, last_action_error="Episode already terminated", ) with debug_context( episode_id=self._state.episode_id, task_id=self._scenario.task_id, difficulty=self._scenario.difficulty, ): self._state.step_count += 1 logger.info( "env.step start episode_id=%s step=%s action_type=%s slide_index=%s", self._state.episode_id, self._state.step_count, action.action_type, action.slide_index, ) write_debug_event( "env.step.start", { "step": self._state.step_count, "action": action.model_dump(mode="json"), }, ) reward = 0.0 invalid_reason: str | None = None inspection: ExtractedPresentation | None = None try: execution_result = self._execute_action(action) self._update_known_named_shapes(action, execution_result) self._last_action_result = { "action_type": execution_result.action_type, "slide_index": self._result_slide_index(action, execution_result), "tool_result": execution_result.tool_result, } write_debug_event( "action.executed", { "step": self._state.step_count, "action_type": execution_result.action_type, "tool_result": execution_result.tool_result, "affected_slide_index": execution_result.affected_slide_index, }, ) if execution_result.affected_slide_index is not None: inspection = self._inspection_service.inspect_presentation( self._editor ) intermediate_result = self._score_intermediate_step( slide_index=execution_result.affected_slide_index, inspection=inspection, ) reward = clamp_reward(intermediate_result.reward_total) self._last_reward_details = { "kind": "intermediate_slide", "result": { **to_serializable(intermediate_result), "reward_total": reward, }, } write_debug_event( "reward.intermediate", { "step": self._state.step_count, "reward_result": to_serializable(intermediate_result), }, ) else: reward = clamp_reward(execution_result.reward) self._last_reward_details = { "kind": "action", "result": {"reward_total": reward}, } write_debug_event( "reward.action", { "step": self._state.step_count, "reward_total": reward, }, ) logger.info( "env.step action complete episode_id=%s step=%s action_type=%s reward=%s tool_result=%s", self._state.episode_id, self._state.step_count, action.action_type, reward, execution_result.tool_result, ) except (IndexError, KeyError, ValueError) as error: invalid_reason = str(error) reward = clamp_reward(_INVALID_ACTION_PENALTY) self._last_action_result = None self._last_reward_details = { "kind": "invalid_action", "result": {"reward_total": reward, "error": invalid_reason}, } write_debug_event( "env.step.invalid_action", { "step": self._state.step_count, "action": action.model_dump(mode="json"), "error": invalid_reason, "reward_total": reward, }, ) logger.warning( "env.step invalid action episode_id=%s step=%s action_type=%s error=%s", self._state.episode_id, self._state.step_count, action.action_type, invalid_reason, ) termination_reason = self._should_terminate() reward = clamp_reward(float(reward)) if termination_reason is not None: reward = self._finalize_episode(termination_reason) logger.info( "env.step end episode_id=%s step=%s done=%s reward=%s score=%s error=%s", self._state.episode_id, self._state.step_count, self._done, reward, self._last_score, invalid_reason, ) write_debug_event( "env.step.end", { "step": self._state.step_count, "done": self._done, "reward": reward, "score": self._last_score, "error": invalid_reason, "termination_reason": self._termination_reason, }, ) return self._build_observation( reward=reward, done=self._done, last_action_error=invalid_reason, inspection=inspection, ) @property def state(self) -> State: slide_count = len(self._editor.prs.slides) if self._editor is not None else 0 return State( episode_id=self._state.episode_id, step_count=self._state.step_count, done=self._done, slide_count=slide_count, max_steps=self._max_steps, task_name=self._current_task_name(), score=self._last_score, ) def _sample_scenario(self) -> TaskScenario: if not self._scenario_queue: self._scenario_queue = list(self._task_registry.all()) self._rng.shuffle(self._scenario_queue) return self._scenario_queue.pop(0) def _resolve_scenario( self, *, task_id: Any | None, difficulty: Any | None, ) -> TaskScenario: if task_id is not None: return self._task_registry.get(str(task_id)) if difficulty is not None: return self._task_registry.sample(self._rng, difficulty=str(difficulty)) return self._sample_scenario() def _build_eval_spec(self, scenario: TaskScenario) -> EvalSpec: cached_eval_spec = self._eval_spec_cache.get(scenario.task_id) if cached_eval_spec is not None: logger.info("env.eval_spec cache hit task_id=%s", scenario.task_id) return cached_eval_spec logger.info("env.eval_spec build start task_id=%s", scenario.task_id) eval_spec = build_eval_spec( scenario.prompt_text, scenario.source_pack, scenario.task_constraints, quiz_bank_service=self._quiz_bank_service, ) self._eval_spec_cache[scenario.task_id] = eval_spec logger.info( "env.eval_spec build complete task_id=%s question_count=%s spec_hash=%s", scenario.task_id, len(eval_spec.slidesgenbench.quiz_bank), eval_spec.spec_hash, ) return eval_spec def _initialize_editor(self, theme: dict[str, Any]) -> PptxEditor: editor = PptxEditor() register_theme(editor, theme) return editor def _build_observation( self, *, reward: float, done: bool, last_action_error: str | None, inspection: ExtractedPresentation | None = None, ) -> PptAgentObservation: if self._editor is None: raise RuntimeError("Environment must be reset before building observations") inspection = inspection or self._inspection_service.inspect_presentation( self._editor ) last_action_result = ( { **self._last_action_result, "reward_details": self._last_reward_details or {}, } if self._last_action_result is not None else None ) return PptAgentObservation( task_name=self._current_task_name(), difficulty=self._difficulty(), slide_count=inspection.slide_count, task_prompt=self._task_prompt(), source_context=self._source_context(), last_action_error=last_action_error, score=self._last_score, prompt_summary=self._prompt_summary(), last_action_result=last_action_result, termination_reason=self._termination_reason, done=done, reward=reward, metadata={ "reward_details": self._last_reward_details or {}, "current_theme": self._editor.get_theme(), "known_named_shapes_by_slide": { slide_index: dict(named_shapes) for slide_index, named_shapes in self._known_named_shapes_by_slide.items() }, "slide_constraints": self._slide_constraints(), "default_save_path": self._default_output_path(), "max_steps": self._max_steps, "step_count": self._state.step_count, }, ) def _current_task_name(self) -> str: return self._scenario.task_id if self._scenario is not None else "" def _prompt_summary(self) -> str: if self._scenario is None: return "" prompt = self._scenario.prompt_text.strip() return prompt if len(prompt) <= 220 else f"{prompt[:217]}..." def _task_prompt(self) -> str: if self._scenario is None: return "" return self._scenario.prompt_text def _source_context(self) -> str: if self._scenario is None: return "" sections: list[str] = [] if self._scenario.source_pack.brief: sections.append(f"[Source Pack Brief]\n{self._scenario.source_pack.brief}") for document in self._scenario.source_pack.documents: title = document.title.strip() if document.title else document.doc_id if document.pages: body = "\n\n".join( f"Page {page_index}:\n{page_text.strip()}" for page_index, page_text in enumerate(document.pages, start=1) if page_text and page_text.strip() ) else: body_parts = [document.text] if document.text else [] body = "\n".join( part.strip() for part in body_parts if part and part.strip() ) if body: sections.append(f"[{title} | {document.doc_id}]\n{body}") return "\n\n".join(sections) def _difficulty(self) -> str: if self._scenario is None: return "" return self._scenario.difficulty def _execute_action(self, action: PptAgentAction) -> _ActionExecutionResult: handler = getattr(self, f"_handle_{action.action_type.lower()}", None) if handler is None: raise ValueError(f"Unsupported action type '{action.action_type}'") return handler(action) def _resolve_slide_id(self, action: PptAgentAction) -> int: if self._editor is None: raise RuntimeError("Environment must be reset before resolving slide ids") if action.slide_index is None: raise ValueError(f"{action.action_type} requires slide_index") return self._editor.get_slide_id(action.slide_index - 1) def _handle_create_slide(self, action: PptAgentAction) -> _ActionExecutionResult: if self._editor is None: raise RuntimeError("Environment must be reset before create_slide") payload = dict(action.payload) tool_result = create_slide( self._editor, background_color=payload.get("background_color", ""), shapes=payload.get("shapes") or [], ) return _ActionExecutionResult( action_type=action.action_type, tool_result=tool_result, affected_slide_index=len(self._editor.prs.slides), ) def _handle_update_slide(self, action: PptAgentAction) -> _ActionExecutionResult: if action.slide_index is None: raise ValueError("update_slide requires slide_index") slide_id = self._resolve_slide_id(action) payload = dict(action.payload) tool_result = update_slide( self._editor, slide_id, background_color=payload.get("background_color"), delete_shape_ids=payload.get("delete_shape_ids") or [], add_shapes=payload.get("add_shapes") or [], update_shapes=payload.get("update_shapes") or [], ) return _ActionExecutionResult( action_type=action.action_type, tool_result=tool_result, affected_slide_index=action.slide_index, ) def _handle_delete_slide(self, action: PptAgentAction) -> _ActionExecutionResult: if self._editor is None: raise RuntimeError("Environment must be reset before delete_slide") slide_id = self._resolve_slide_id(action) tool_result = delete_slide(self._editor, slide_id) return _ActionExecutionResult( action_type=action.action_type, tool_result=tool_result, affected_slide_index=None, reward=0.0, ) def _handle_save_presentation( self, action: PptAgentAction ) -> _ActionExecutionResult: if self._editor is None: raise RuntimeError("Environment must be reset before save_presentation") payload = dict(action.payload) tool_result = save_presentation( self._editor, payload.get("path") or self._default_output_path(), ) return _ActionExecutionResult( action_type=action.action_type, tool_result=tool_result, affected_slide_index=None, reward=0.0, ) def _handle_set_theme(self, action: PptAgentAction) -> _ActionExecutionResult: if self._editor is None: raise RuntimeError("Environment must be reset before set_theme") payload = dict(action.payload) update_theme(self._editor, payload) return _ActionExecutionResult( action_type=action.action_type, tool_result={"theme": self._editor.get_theme()}, affected_slide_index=None, reward=0.0, ) def _default_output_path(self) -> str: task_name = self._current_task_name() or "presentation" episode_id = self._state.episode_id or str(uuid4()) return str(Path("outputs") / f"{task_name}_{episode_id}.pptx") def _slide_constraints(self) -> dict[str, int]: if self._eval_spec is None: return {} task_spec = self._eval_spec.task_spec constraints: dict[str, int] = {} if isinstance(task_spec.min_slides, int): constraints["min_slides"] = task_spec.min_slides if isinstance(task_spec.max_slides, int): constraints["max_slides"] = task_spec.max_slides return constraints def _result_slide_index( self, action: PptAgentAction, execution_result: _ActionExecutionResult ) -> int | None: if execution_result.action_type == "create_slide": return execution_result.affected_slide_index return action.slide_index def _update_known_named_shapes( self, action: PptAgentAction, execution_result: _ActionExecutionResult ) -> None: tool_result = execution_result.tool_result named_shapes = tool_result.get("named_shapes") normalized_named_shapes = { name: shape_id for name, shape_id in (named_shapes or {}).items() if isinstance(name, str) and isinstance(shape_id, int) } if execution_result.action_type == "create_slide": slide_index = execution_result.affected_slide_index if isinstance(slide_index, int): self._known_named_shapes_by_slide[slide_index] = dict( normalized_named_shapes ) return if execution_result.action_type == "update_slide": slide_index = action.slide_index if not isinstance(slide_index, int): return current = dict(self._known_named_shapes_by_slide.get(slide_index, {})) deleted_shape_ids = { shape_id for shape_id in tool_result.get("deleted_shape_ids") or [] if isinstance(shape_id, int) } if deleted_shape_ids: current = { name: shape_id for name, shape_id in current.items() if shape_id not in deleted_shape_ids } current.update(normalized_named_shapes) self._known_named_shapes_by_slide[slide_index] = current return if execution_result.action_type == "delete_slide": slide_index = action.slide_index if not isinstance(slide_index, int): return reindexed: dict[int, dict[str, int]] = {} for ( existing_index, existing_shapes, ) in self._known_named_shapes_by_slide.items(): if existing_index < slide_index: reindexed[existing_index] = existing_shapes elif existing_index > slide_index: reindexed[existing_index - 1] = existing_shapes self._known_named_shapes_by_slide = reindexed def _score_intermediate_step( self, *, slide_index: int, inspection: ExtractedPresentation, ) -> Any: if self._eval_spec is None: raise RuntimeError("Environment must be reset before scoring") slide_extraction = inspection.slides[slide_index - 1] previous_slide_extractions = inspection.slides[: slide_index - 1] or None return evaluate_slide( self._eval_spec, slide_index, slide_extraction=slide_extraction, previous_slide_extractions=previous_slide_extractions, ) def _should_terminate(self) -> str | None: if ( self._last_action_result is not None and self._last_action_result.get("action_type") == "save_presentation" ): return "presentation_saved" if self._state.step_count >= self._max_steps: return "step_budget_exhausted" return None def _finalize_episode(self, termination_reason: str) -> float: if self._eval_spec is None or self._editor is None: raise RuntimeError("Environment must be reset before finalization") self._done = True self._termination_reason = termination_reason logger.info( "env.finalize start episode_id=%s reason=%s", self._state.episode_id, termination_reason, ) result = evaluate_presentation( self._eval_spec, self._editor, render_service=self._render_service, inspection_service=self._inspection_service, quantitative_quiz_judge_service=self._quantitative_quiz_judge_service, ) self._last_score = result.reward_total self._last_reward_details = { "kind": "terminal_presentation", "result": to_serializable(result), } write_debug_event( "reward.terminal", { "termination_reason": termination_reason, "reward_result": to_serializable(result), }, ) logger.info( "env.finalize complete episode_id=%s reward=%s metadata=%s", self._state.episode_id, self._last_score, result.metadata, ) return self._last_score if __name__ == "__main__": # pragma: no cover env = PptAgentEnvironment() observation = env.reset(seed=7) print("Reset prompt:", observation.prompt_summary) observation = env.step(PptAgentAction(action_type="create_slide")) print("Slides after create_slide:", observation.slide_count) observation = env.step( PptAgentAction( action_type="update_slide", slide_index=1, payload={ "add_shapes": [ { "type": "text", "text": "Hello from the PPT skeleton", "x": 0.8, "y": 1.2, "w": 8.4, "h": 1.0, "style": { "font_name": "", "font_size_pt": "", "color_hex": "", }, } ] }, ) ) while not observation.done: observation = env.step(PptAgentAction(action_type="create_slide")) print("Terminal reward:", observation.reward)