refactor rubric hook helpers (#892)

Author: burtenshaw

src/openenv/core/rubrics/base.py

@@ -81,25 +81,35 @@ def __call__(self, action: Any, observation: Any):
 
     def _call_sync(self, action: Any, observation: Any, result: float) -> float:
         """Synchronous call path."""
-        self.last_score = result
+        return self._finish_forward(action, observation, result)
 
-        # Post-forward hooks
-        for hook in self._forward_hooks:
-            hook(self, action, observation, result)
-
-        return result
+    def _run_forward_pre_hooks(self, action: Any, observation: Any) -> None:
+        """Run pre-forward hooks synchronously."""
+        for hook in self._forward_pre_hooks:
+            hook(self, action, observation)
 
-    async def _call_async(self, action: Any, observation: Any, result_coro) -> float:
-        """Asynchronous call path."""
-        # Pre-forward hooks
+    async def _run_forward_pre_hooks_async(self, action: Any, observation: Any) -> None:
+        """Run pre-forward hooks from an async call path."""
         for hook in self._forward_pre_hooks:
             if inspect.iscoroutinefunction(hook):
                 await hook(self, action, observation)
             else:
                 hook(self, action, observation)
 
-        # Await the forward result
-        result = await result_coro
+    def _finish_forward(self, action: Any, observation: Any, result: float) -> float:
+        """Store the result and run post-forward hooks synchronously."""
+        self.last_score = result
+
+        # Post-forward hooks
+        for hook in self._forward_hooks:
+            hook(self, action, observation, result)
+
+        return result
+
+    async def _finish_forward_async(
+        self, action: Any, observation: Any, result: float
+    ) -> float:
+        """Store the result and run post-forward hooks from an async call path."""
         self.last_score = result
 
         # Post-forward hooks
@@ -111,6 +121,15 @@ async def _call_async(self, action: Any, observation: Any, result_coro) -> float
 
         return result
 
+    async def _call_async(self, action: Any, observation: Any, result_coro) -> float:
+        """Asynchronous call path."""
+        # Pre-forward hooks
+        await self._run_forward_pre_hooks_async(action, observation)
+
+        # Await the forward result
+        result = await result_coro
+        return await self._finish_forward_async(action, observation, result)
+
     @abstractmethod
     def forward(self, action: Any, observation: Any) -> float:
         """Compute the reward. Implement this in subclasses.

src/openenv/core/rubrics/containers.py

@@ -75,14 +75,8 @@ def __call__(self, action: Any, observation: Any):
             if _in_async_context():
                 return self._empty_async(action, observation)
             else:
-                # Pre-hooks
-                for hook in self._forward_pre_hooks:
-                    hook(self, action, observation)
-                result = 1.0
-                self.last_score = result
-                for hook in self._forward_hooks:
-                    hook(self, action, observation, result)
-                return result
+                self._run_forward_pre_hooks(action, observation)
+                return self._finish_forward(action, observation, 1.0)
 
         # Call first rubric to see if it's async
         first_result = self._rubric_list[0](action, observation)
@@ -92,13 +86,8 @@ def __call__(self, action: Any, observation: Any):
         else:
             # Continue with sync path
             if first_result == 0.0:
-                # Pre-hooks
-                for hook in self._forward_pre_hooks:
-                    hook(self, action, observation)
-                self.last_score = 0.0
-                for hook in self._forward_hooks:
-                    hook(self, action, observation, 0.0)
-                return 0.0
+                self._run_forward_pre_hooks(action, observation)
+                return self._finish_forward(action, observation, 0.0)
 
             final_result = first_result
             for i, rubric in enumerate(self._rubric_list[1:], start=1):
@@ -109,150 +98,68 @@ def __call__(self, action: Any, observation: Any):
                     return self._call_async_mid(
                         action,
                         observation,
-                        final_result,
                         score,
                         self._rubric_list[i + 1 :],
                     )
                 if score == 0.0:
-                    # Pre-hooks
-                    for hook in self._forward_pre_hooks:
-                        hook(self, action, observation)
-                    self.last_score = 0.0
-                    for hook in self._forward_hooks:
-                        hook(self, action, observation, 0.0)
-                    return 0.0
+                    self._run_forward_pre_hooks(action, observation)
+                    return self._finish_forward(action, observation, 0.0)
                 final_result = score
 
             # All sync - check if in async context
             if _in_async_context():
                 return self._wrap_sync_result(action, observation, final_result)
             else:
-                # Pre-hooks
-                for hook in self._forward_pre_hooks:
-                    hook(self, action, observation)
-                self.last_score = final_result
-                for hook in self._forward_hooks:
-                    hook(self, action, observation, final_result)
-                return final_result
+                self._run_forward_pre_hooks(action, observation)
+                return self._finish_forward(action, observation, final_result)
 
     async def _empty_async(self, action, observation):
         """Async path for empty sequential."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
-
-        result = 1.0
-        self.last_score = result
-
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, result)
-            else:
-                hook(self, action, observation, result)
-        return result
+        await self._run_forward_pre_hooks_async(action, observation)
+        return await self._finish_forward_async(action, observation, 1.0)
 
     async def _wrap_sync_result(self, action, observation, result):
         """Wrap sync result for async context."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
-
-        self.last_score = result
-
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, result)
-            else:
-                hook(self, action, observation, result)
-        return result
+        await self._run_forward_pre_hooks_async(action, observation)
+        return await self._finish_forward_async(action, observation, result)
 
     async def _call_async_detected(self, action, observation, first_coro):
         """Async path when first child is async."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
+        await self._run_forward_pre_hooks_async(action, observation)
 
         result = await first_coro
         if result == 0.0:
-            self.last_score = 0.0
-            for hook in self._forward_hooks:
-                if inspect.iscoroutinefunction(hook):
-                    await hook(self, action, observation, result)
-                else:
-                    hook(self, action, observation, result)
-            return 0.0
+            return await self._finish_forward_async(action, observation, result)
 
         for rubric in self._rubric_list[1:]:
             score = rubric(action, observation)
             if inspect.iscoroutine(score):
                 score = await score
             if score == 0.0:
-                self.last_score = 0.0
-                for hook in self._forward_hooks:
-                    if inspect.iscoroutinefunction(hook):
-                        await hook(self, action, observation, 0.0)
-                    else:
-                        hook(self, action, observation, 0.0)
-                return 0.0
+                return await self._finish_forward_async(action, observation, 0.0)
             result = score
 
-        self.last_score = result
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, result)
-            else:
-                hook(self, action, observation, result)
-        return result
+        return await self._finish_forward_async(action, observation, result)
 
-    async def _call_async_mid(
-        self, action, observation, current_result, first_async_coro, remaining
-    ):
+    async def _call_async_mid(self, action, observation, first_async_coro, remaining):
         """Async path when async detected mid-execution."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
+        await self._run_forward_pre_hooks_async(action, observation)
 
         # Await the first async rubric (already called)
         result = await first_async_coro
         if result == 0.0:
-            self.last_score = 0.0
-            for hook in self._forward_hooks:
-                if inspect.iscoroutinefunction(hook):
-                    await hook(self, action, observation, 0.0)
-                else:
-                    hook(self, action, observation, 0.0)
-            return 0.0
+            return await self._finish_forward_async(action, observation, 0.0)
 
         # Continue with remaining rubrics
         for rubric in remaining:
             score = rubric(action, observation)
             if inspect.iscoroutine(score):
                 score = await score
             if score == 0.0:
-                self.last_score = 0.0
-                for hook in self._forward_hooks:
-                    if inspect.iscoroutinefunction(hook):
-                        await hook(self, action, observation, 0.0)
-                    else:
-                        hook(self, action, observation, 0.0)
-                return 0.0
+                return await self._finish_forward_async(action, observation, 0.0)
             result = score
 
-        self.last_score = result
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, result)
-            else:
-                hook(self, action, observation, result)
-        return result
+        return await self._finish_forward_async(action, observation, result)
 
     def __len__(self) -> int:
         return len(self._rubric_list)
@@ -304,33 +211,17 @@ def __call__(self, action: Any, observation: Any):
             return self._call_async(action, observation, score)
         else:
             # Child is sync
-            # Pre-hooks
-            for hook in self._forward_pre_hooks:
-                hook(self, action, observation)
+            self._run_forward_pre_hooks(action, observation)
             result = 0.0 if score < self.threshold else score
-            self.last_score = result
-            for hook in self._forward_hooks:
-                hook(self, action, observation, result)
-            return result
+            return self._finish_forward(action, observation, result)
 
     async def _call_async(self, action, observation, score_coro):
         """Async path."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
+        await self._run_forward_pre_hooks_async(action, observation)
 
         score = await score_coro
         result = 0.0 if score < self.threshold else score
-        self.last_score = result
-
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, result)
-            else:
-                hook(self, action, observation, result)
-        return result
+        return await self._finish_forward_async(action, observation, result)
 
 
 class WeightedSum(Rubric):
@@ -395,24 +286,15 @@ def __call__(self, action: Any, observation: Any):
             return self._call_async(action, observation, results)
         else:
             # Sync path
-            # Pre-hooks
-            for hook in self._forward_pre_hooks:
-                hook(self, action, observation)
+            self._run_forward_pre_hooks(action, observation)
             total = 0.0
             for score, weight in zip(results, self._weights):
                 total += score * weight
-            self.last_score = total
-            for hook in self._forward_hooks:
-                hook(self, action, observation, total)
-            return total
+            return self._finish_forward(action, observation, total)
 
     async def _call_async(self, action, observation, results):
         """Async path with parallel execution."""
-        for hook in self._forward_pre_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation)
-            else:
-                hook(self, action, observation)
+        await self._run_forward_pre_hooks_async(action, observation)
 
         # Separate sync and async results
         async_tasks = []
@@ -437,14 +319,7 @@ async def _call_async(self, action, observation, results):
         for score, weight in zip(scores, self._weights):
             total += score * weight
 
-        self.last_score = total
-
-        for hook in self._forward_hooks:
-            if inspect.iscoroutinefunction(hook):
-                await hook(self, action, observation, total)
-            else:
-                hook(self, action, observation, total)
-        return total
+        return await self._finish_forward_async(action, observation, total)
 
     @property
     def weights(self) -> List[float]: