review comments

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

Author: akihikokuroda

mellea/core/base.py

@@ -410,50 +410,63 @@ class ComputedModelOutputThunk(ModelOutputThunk[S]):
     and those that are already computed. It should be returned from synchronous functions
     and sampling strategies to indicate that no awaiting is needed.
 
+    Rather than creating from scratch, this wraps an existing ModelOutputThunk and ensures
+    it's fully computed. All attribute access is delegated to the wrapped thunk.
+
     Key differences from ModelOutputThunk:
-    - Always initialized with a value (cannot be None)
+    - Always initialized from a computed ModelOutputThunk
     - _computed is always True
     - Cannot be used for streaming (generation fields are not set)
     - Provides type safety to indicate "already computed"
     """
 
-    def __init__(
-        self,
-        value: str,
-        meta: dict[str, Any] | None = None,
-        parsed_repr: S | None = None,
-        tool_calls: dict[str, ModelToolCall] | None = None,
-    ):
-        """Initializes a computed ModelOutputThunk with a required value.
+    def __init__(self, thunk: ModelOutputThunk[S]):
+        """Wraps an existing ModelOutputThunk, ensuring it's computed.
 
         Args:
-            value: The computed string value (required, cannot be None)
-            meta: Optional metadata dictionary
-            parsed_repr: Optional parsed representation
-            tool_calls: Optional tool calls dictionary
+            thunk: A ModelOutputThunk that must be fully computed (value cannot be None)
+
+        Raises:
+            ValueError: If the thunk is not computed or has a None value
         """
-        if value is None:
+        if not thunk.is_computed():
+            raise ValueError(
+                "ComputedModelOutputThunk requires a computed ModelOutputThunk"
+            )
+        if thunk.value is None:
             raise ValueError("ComputedModelOutputThunk requires a non-None value")
 
-        super().__init__(value, meta, parsed_repr, tool_calls)
-
-        # Ensure computed flag is set
-        assert self._computed, "ComputedModelOutputThunk must be computed"
+        # Store the wrapped thunk and ensure it's marked as computed
+        self._wrapped_thunk = thunk
+        self._wrapped_thunk._computed = True
 
         # Clear generation-related fields since this is already computed
-        self._generate = None
-        self._generate_type = GenerateType.NONE
-        self._generate_extra = None
-        self._process = None
-        self._post_process = None
+        self._wrapped_thunk._generate = None
+        self._wrapped_thunk._generate_type = GenerateType.NONE
+        self._wrapped_thunk._generate_extra = None
+        self._wrapped_thunk._process = None
+        self._wrapped_thunk._post_process = None
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate all attribute access to the wrapped thunk."""
+        return getattr(self._wrapped_thunk, name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Delegate all attribute setting to the wrapped thunk, except for _wrapped_thunk itself."""
+        if name == "_wrapped_thunk":
+            object.__setattr__(self, name, value)
+        else:
+            setattr(self._wrapped_thunk, name, value)
 
     async def avalue(self) -> str:
         """Returns the value immediately since it's already computed.
 
         Overrides the parent method to avoid unnecessary async operations.
         """
-        assert self.value is not None, "ComputedModelOutputThunk value cannot be None"
-        return self.value
+        assert self._wrapped_thunk.value is not None, (
+            "ComputedModelOutputThunk value cannot be None"
+        )
+        return self._wrapped_thunk.value
 
     async def astream(self) -> str:
         """Returns the value immediately since streaming is not applicable.

mellea/stdlib/functional.py

@@ -44,7 +44,7 @@ def act(
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
     tool_calls: bool = False,
-) -> tuple[ModelOutputThunk[S], Context]: ...
+) -> tuple[ComputedModelOutputThunk[S], Context]: ...
 
 
 @overload
@@ -73,7 +73,7 @@ def act(
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
     tool_calls: bool = False,
-) -> tuple[ModelOutputThunk[S], Context] | SamplingResult[S]:
+) -> tuple[ComputedModelOutputThunk[S], Context] | SamplingResult[S]:
     """Runs a generic action, and adds both the action and the result to the context.
 
     Args:
@@ -129,7 +129,7 @@ def instruct(
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
     tool_calls: bool = False,
-) -> tuple[ModelOutputThunk[str], Context]: ...
+) -> tuple[ComputedModelOutputThunk[str], Context]: ...
 
 
 @overload
@@ -170,7 +170,7 @@ def instruct(
     format: type[BaseModelSubclass] | None = None,
     model_options: dict | None = None,
     tool_calls: bool = False,
-) -> tuple[ModelOutputThunk[str], Context] | SamplingResult[str]:
+) -> tuple[ComputedModelOutputThunk[str], Context] | SamplingResult[str]:
     """Generates from an instruction.
 
     Args:
@@ -555,20 +555,9 @@ async def aact(
             # ._generate_log should never be None after generation.
             assert result._generate_log is not None
             result._generate_log.is_final_result = True
-            generate_logs.append(result._generate_log)
 
             # Wrap in ComputedModelOutputThunk to indicate it's fully computed
-            computed_result = ComputedModelOutputThunk(
-                value=result.value,  # type: ignore
-                meta=result._meta,
-                parsed_repr=result.parsed_repr,
-                tool_calls=result.tool_calls,
-            )
-            # Copy over important fields
-            computed_result._thinking = result._thinking
-            computed_result._context = result._context
-            computed_result._action = result._action
-            computed_result._model_options = result._model_options
+            computed_result = ComputedModelOutputThunk(result)
             computed_result._generate_log = result._generate_log
 
             # Update context to point to the wrapped result instead of original
@@ -611,18 +600,7 @@ async def aact(
         )
 
         # Wrap sampling result in ComputedModelOutputThunk since it's always computed
-        computed_result = ComputedModelOutputThunk(
-            value=result.value,  # type: ignore
-            meta=result._meta,
-            parsed_repr=result.parsed_repr,
-            tool_calls=result.tool_calls,
-        )
-        # Copy over important fields
-        computed_result._thinking = result._thinking
-        computed_result._context = result._context
-        computed_result._action = result._action
-        computed_result._model_options = result._model_options
-        computed_result._generate_log = result._generate_log
+        computed_result = ComputedModelOutputThunk(result)
 
         # Update the sampling result to use the computed thunk
         sampling_result.sample_generations[sampling_result.result_index] = (