symbolicai 0.20.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

symai/strategy.py

@@ -2,7 +2,8 @@ import inspect
 import logging
 import time
 from collections import defaultdict
-from typing import Callable
+from collections.abc import Callable
+from typing import Any, ClassVar
 
 import numpy as np
 from beartype import beartype
@@ -14,9 +15,9 @@ from rich.panel import Panel
 from rich.table import Table
 
 from .components import Function
-from .models import (LLMDataModel, TypeValidationError,
-                     build_dynamic_llm_datamodel)
+from .models import LLMDataModel, TypeValidationError, build_dynamic_llm_datamodel
 from .symbol import Expression
+from .utils import UserMessage
 
 
 class ValidationFunction(Function):
@@ -29,18 +30,18 @@ class ValidationFunction(Function):
       • Error simplification
     """
     # Have some default retry params that don't add overhead
-    _default_retry_params = dict(
-        tries=8,
-        delay=0.015,
-        backoff=1.25,
-        jitter=0.0,
-        max_delay=0.25,
-        graceful=False
-    )
+    _default_retry_params: ClassVar[dict[str, int | float | bool]] = {
+        "tries": 8,
+        "delay": 0.015,
+        "backoff": 1.25,
+        "jitter": 0.0,
+        "max_delay": 0.25,
+        "graceful": False,
+    }
 
     def __init__(
         self,
-        retry_params: dict[str, int | float | bool] = None,
+        retry_params: dict[str, int | float | bool] | None = None,
         verbose: bool = False,
         *args,
         **kwargs,
@@ -92,10 +93,9 @@ class ValidationFunction(Function):
             seed = 42
 
         rnd = np.random.RandomState(seed=seed)
-        seeds = rnd.randint(
+        return rnd.randint(
             0, np.iinfo(np.int16).max, size=num_seeds, dtype=np.int16
         ).tolist()
-        return seeds
 
     def simplify_validation_errors(self, error: ValidationError) -> str:
         """
@@ -130,12 +130,14 @@ class ValidationFunction(Function):
         _delay = min(base + jit, self.retry_params['max_delay'])
         time.sleep(_delay)
 
-    def remedy_prompt(self, *args, **kwargs):
+    def remedy_prompt(self, *_args, **_kwargs):
         """
         Abstract or base remedy prompt method.
         Child classes typically override this to include additional context needed for correction.
         """
-        raise NotImplementedError("Each child class needs its own remedy_prompt implementation.")
+        msg = "Each child class needs its own remedy_prompt implementation."
+        UserMessage(msg)
+        raise NotImplementedError(msg)
 
     def display_panel(self, content, title, border_style="cyan", style="#f0eee6", padding=(1,2)):
         """
@@ -171,7 +173,7 @@ class TypeValidationFunction(ValidationFunction):
         *args,
         **kwargs,
     ):
-        super().__init__(retry_params=retry_params, verbose=verbose, *args, **kwargs)
+        super().__init__(*args, retry_params=retry_params, verbose=verbose, **kwargs)
         self.input_data_model = None
         self.output_data_model = None
         self.accumulate_errors = accumulate_errors
@@ -181,11 +183,15 @@ class TypeValidationFunction(ValidationFunction):
         assert attach_to in ["input", "output"], f"Invalid attach_to value: {attach_to}; must be either 'input' or 'output'"
         if attach_to == "input":
             if self.input_data_model is not None and not override:
-                raise ValueError("There is already a data model attached to the input. If you want to override it, set `override=True`.")
+                msg = "There is already a data model attached to the input. If you want to override it, set `override=True`."
+                UserMessage(msg)
+                raise ValueError(msg)
             self.input_data_model = data_model
         elif attach_to == "output":
             if self.output_data_model is not None and not override:
-                raise ValueError("There is already a data model attached to the output. If you want to override it, set `override=True`.")
+                msg = "There is already a data model attached to the output. If you want to override it, set `override=True`."
+                UserMessage(msg)
+                raise ValueError(msg)
             self.output_data_model = data_model
 
     def remedy_prompt(self, prompt: str, output: str, errors: str) -> str:
@@ -268,118 +274,174 @@ Important guidelines:
 </guidelines>
 """
 
-    def forward(self, prompt: str, f_semantic_conditions: list[Callable] | None = None, *args, **kwargs):
+    def _ensure_output_model(self):
         if self.output_data_model is None:
-            raise ValueError("While the input data model is optional, the output data model must be provided. Please register it before calling the `forward` method.")
+            msg = (
+                "While the input data model is optional, the output data model must be provided. "
+                "Please register it before calling the `forward` method."
+            )
+            UserMessage(msg)
+            raise ValueError(msg)
+
+    def _display_verbose_panels(self, prompt: str):
+        if not self.verbose:
+            return
+        for label, body in [
+            ("Prompt", prompt),
+            ("Input data model", self.input_data_model.simplify_json_schema() if self.input_data_model else 'N/A'),
+            ("Output data model", self.output_data_model.simplify_json_schema()),
+        ]:
+            self.display_panel(body, title=label)
+
+    def _check_semantic_conditions(
+        self,
+        result,
+        f_semantic_conditions: list[Callable] | None,
+    ) -> str | None:
+        if f_semantic_conditions is None:
+            return None
+        try:
+            assert all(
+                f(result if not getattr(self.output_data_model, '_is_dynamic_model', False) else result.value)
+                for f in f_semantic_conditions
+            )
+        except Exception as err:
+            return f"Semantic validation failed with:\n{err!s}"
+        return None
+
+    def _format_validation_error(self, error: Exception) -> str:
+        if isinstance(error, ValidationError):
+            return self.simplify_validation_errors(error)
+        return str(error)
+
+    def _handle_failed_validation_attempt(
+        self,
+        attempt_index: int,
+        prompt: str,
+        json_str: str,
+        errors: list[str],
+        error: Exception,
+        remedy_seeds: list[Any],
+        kwargs: dict,
+    ) -> str:
+        logger.info(f"Validation attempt {attempt_index + 1} failed, pausing before retry…")
+        self._pause(attempt_index)
+        error_str = self._format_validation_error(error)
+        errors.append(error_str)
+
+        logger.error("Validation errors identified!")
+        if self.verbose:
+            errors_report = "\n".join(errors) if self.accumulate_errors else error_str
+            title = f"Validation Errors ({'accumulated errors' if self.accumulate_errors else 'last error'})"
+            self.display_panel(errors_report, title=title, border_style="red")
+
+        logger.info("Updating remedy function context…")
+        context = self.remedy_prompt(
+            prompt=prompt,
+            output=json_str,
+            errors="\n".join(errors) if self.accumulate_errors else error_str,
+        )
+        self.remedy_function.clear()
+        self.remedy_function.adapt(context)
+        if self.verbose:
+            self.display_panel(self.remedy_function.dynamic_context, title="New Context")
+
+        json_str = self.remedy_function(seed=remedy_seeds[attempt_index], **kwargs).value
+        logger.info("Applied remedy function with updated context!")
+        return json_str
+
+    def _run_validation_attempts(
+        self,
+        prompt: str,
+        f_semantic_conditions: list[Callable] | None,
+        validation_context: dict,
+        remedy_seeds: list[Any],
+        json_str: str,
+        kwargs: dict,
+    ) -> tuple[Any | None, str, list[str]]:
+        errors: list[str] = []
+        result = None
+        total_attempts = self.retry_params["tries"] + 1
+        for attempt in range(total_attempts):
+            if attempt != self.retry_params["tries"]:
+                logger.info(f"Attempt {attempt + 1}/{self.retry_params['tries']}: Attempting validation…")
+            try:
+                result = self.output_data_model.model_validate_json(
+                    json_str,
+                    strict=False,
+                    context=validation_context,
+                )
+                semantic_error = self._check_semantic_conditions(result, f_semantic_conditions)
+                if semantic_error is not None:
+                    if attempt == total_attempts - 1:
+                        result = None
+                        errors.append(semantic_error)
+                        break
+                    raise AssertionError(semantic_error)
+                break
+            except Exception as error:
+                json_str = self._handle_failed_validation_attempt(
+                    attempt,
+                    prompt,
+                    json_str,
+                    errors,
+                    error,
+                    remedy_seeds,
+                    kwargs,
+                )
+        return result, json_str, errors
+
+    def _handle_validation_failure(self, prompt: str, json_str: str, errors: list[str]):
+        logger.error("All validation attempts failed!")
+        if self.retry_params['graceful']:
+            return
+        raise TypeValidationError(
+            prompt=prompt,
+            result=json_str,
+            violations=errors,
+        )
+
+    def forward(self, prompt: str, f_semantic_conditions: list[Callable] | None = None, *args, **kwargs):
+        self._ensure_output_model()
         validation_context = kwargs.pop('validation_context', {})
-        # Force JSON mode
         kwargs["response_format"] = {"type": "json_object"}
         logger.info("Initializing validation…")
-        if self.verbose:
-            for label, body in [
-                ("Prompt", prompt),
-                ("Input data model", self.input_data_model.simplify_json_schema() if self.input_data_model else 'N/A'),
-                ("Output data model", self.output_data_model.simplify_json_schema()),
-            ]:
-                self.display_panel(body, title=label)
-
-        # Zero shot the task
+        self._display_verbose_panels(prompt)
+
         context = self.zero_shot_prompt(prompt=prompt)
         json_str = super().forward(context, *args, **kwargs).value
 
         remedy_seeds = self.prepare_seeds(self.retry_params["tries"] + 1, **kwargs)
         logger.info(f"Prepared {len(remedy_seeds)} remedy seeds for validation attempts…")
 
-        result = None
-        errors = []
-        for i in range(self.retry_params["tries"] + 1):
-            if i != self.retry_params["tries"]:
-                logger.info(f"Attempt {i+1}/{self.retry_params['tries']}: Attempting validation…")
-            try:
-                #@NOTE: We use strict=False (default) to allow Pydantic's type coercion.
-                # This handles common LLM output issues like:
-                # - String keys "123" -> integer keys 123 for dict[int, ...] types
-                # - String numbers "42" -> int 42
-                # - Float 1.0 -> int 1
-                # These coercions are safe and helpful when dealing with JSON from LLMs,
-                # while still catching actual type errors (e.g., "not_a_number" -> int fails).
-                result = self.output_data_model.model_validate_json(json_str, strict=False, context=validation_context)
-                if f_semantic_conditions is not None:
-                    try:
-                        assert all(
-                            f(result if not getattr(self.output_data_model, '_is_dynamic_model', False) else result.value)
-                            for f in f_semantic_conditions
-                        )
-                    except Exception as e:
-                        # If we are in the last attempt and semantic validation fails, result will be None and we propagate the error
-                        if i == self.retry_params["tries"]:
-                            result = None
-                            errors.append(f"Semantic validation failed with:\n{str(e)}")
-                            break # We break to avoid going into the remedy loop
-                        raise e
-                break
-            except Exception as e:
-                logger.info(f"Validation attempt {i+1} failed, pausing before retry…")
-
-                self._pause(i)
-
-                if isinstance(e, ValidationError):
-                    error_str = self.simplify_validation_errors(e)
-                else:
-                    error_str = str(e)
-
-                errors.append(error_str)
-
-                logger.error(f"Validation errors identified!")
-                if self.verbose:
-                    self.display_panel(
-                        "\n".join(errors) if self.accumulate_errors else error_str,
-                        title=f"Validation Errors ({'accumulated errors' if self.accumulate_errors else 'last error'})",
-                        border_style="red"
-                    )
-
-                # Update remedy function context
-                logger.info("Updating remedy function context…")
-                context = self.remedy_prompt(prompt=prompt, output=json_str, errors="\n".join(errors) if self.accumulate_errors else error_str)
-                self.remedy_function.clear()
-                self.remedy_function.adapt(context)
-                if self.verbose:
-                    self.display_panel(
-                        self.remedy_function.dynamic_context,
-                        title="New Context"
-                    )
-
-                # Apply the remedy function
-                json_str = self.remedy_function(seed=remedy_seeds[i], **kwargs).value
-                logger.info("Applied remedy function with updated context!")
+        result, json_str, errors = self._run_validation_attempts(
+            prompt,
+            f_semantic_conditions,
+            validation_context,
+            remedy_seeds,
+            json_str,
+            kwargs,
+        )
 
         if result is None:
-            logger.error(f"All validation attempts failed!")
-            if self.retry_params['graceful']:
-                return
-            raise TypeValidationError(
-                prompt=prompt,
-                result=json_str,
-                violations=errors,
-            )
+            return self._handle_validation_failure(prompt, json_str, errors)
 
         logger.success("Validation completed successfully!")
-        # Clear artifacts from the remedy function
         self.remedy_function.clear()
-
         return result
 
 
 @beartype
 class contract:
-    _default_remedy_retry_params = dict(
-        tries=8,
-        delay=0.015,
-        backoff=1.25,
-        jitter=0.0,
-        max_delay=0.25,
-        graceful=False
-    )
+    _default_remedy_retry_params: ClassVar[dict[str, int | float | bool]] = {
+        "tries": 8,
+        "delay": 0.015,
+        "backoff": 1.25,
+        "jitter": 0.0,
+        "max_delay": 0.25,
+        "graceful": False,
+    }
+    _internal_forward_kwargs: ClassVar[set[str]] = {"validation_context"}
 
     def __init__(
         self,
@@ -404,22 +466,30 @@ class contract:
         else:
             logger.enable(__name__)
 
-    def _is_valid_input(self, input):
-        if input is None:
+    def _is_valid_input(self, input_value):
+        if input_value is None:
             logger.error("No `input` argument provided!")
-            raise ValueError("Please provide an `input` argument.")
-        if not isinstance(input, LLMDataModel):
-            logger.error(f"Invalid input type: {type(input)}")
-            raise TypeError(f"Expected input to be of type `LLMDataModel`, got {type(input)}")
+            msg = "Please provide an `input` argument."
+            UserMessage(msg)
+            raise ValueError(msg)
+        if not isinstance(input_value, LLMDataModel):
+            logger.error(f"Invalid input type: {type(input_value)}")
+            msg = f"Expected input to be of type `LLMDataModel`, got {type(input_value)}"
+            UserMessage(msg)
+            raise TypeError(msg)
         return True
 
     def _is_valid_output(self, output_type):
         if output_type == inspect._empty:
             logger.error("Missing return type annotation!")
-            raise ValueError("The contract requires a return type annotation.")
+            msg = "The contract requires a return type annotation."
+            UserMessage(msg)
+            raise ValueError(msg)
         if not issubclass(output_type, LLMDataModel):
             logger.error(f"Invalid return type: {output_type}")
-            raise TypeError("The return type annotation must be a subclass of `LLMDataModel`.")
+            msg = "The return type annotation must be a subclass of `LLMDataModel`."
+            UserMessage(msg)
+            raise TypeError(msg)
         return True
 
     def _try_dynamic_type_annotation(self, original_forward, *, context: str = "input"):
@@ -428,23 +498,42 @@ class contract:
         )
         sig = inspect.signature(original_forward)
         try:
+            resolved_param = None
             # Fallback: look at the relevant part of the function signature
             # depending on whether we deal with an *input* or *output*
             if context == "input":
                 param = sig.parameters.get("input")
+                if param is None:
+                    for candidate in sig.parameters.values():
+                        if candidate.name == "self":
+                            continue
+                        if candidate.kind in (
+                            inspect.Parameter.POSITIONAL_ONLY,
+                            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                            inspect.Parameter.KEYWORD_ONLY,
+                        ):
+                            param = candidate
+                            break
+                resolved_param = param
                 if param is None or param.annotation == inspect._empty:
-                    raise TypeError("Failed to infer type from input parameter annotation")
+                    msg = "Failed to infer type from input parameter annotation"
+                    UserMessage(msg)
+                    raise TypeError(msg)
                 dynamic_model = build_dynamic_llm_datamodel(param.annotation)
             else:  # context == "output"
                 if sig.return_annotation == inspect._empty:
-                    raise TypeError("Failed to infer type from return annotation")
+                    msg = "Failed to infer type from return annotation"
+                    UserMessage(msg)
+                    raise TypeError(msg)
                 dynamic_model = build_dynamic_llm_datamodel(sig.return_annotation)
-        except Exception:
-            logger.exception(f"Failed to build dynamic LLMDataModel from {sig.parameters.get('input')}!")
-            raise TypeError(
+        except Exception as err:
+            logger.exception(f"Failed to build dynamic LLMDataModel from {resolved_param}!")
+            msg = (
                 "The type annotation must be a subclass of `LLMDataModel` or a "
                 "valid Python typing object supported by Pydantic."
             )
+            UserMessage(msg)
+            raise TypeError(msg) from err
 
         dynamic_model._is_dynamic_model = True
         return dynamic_model
@@ -453,49 +542,54 @@ class contract:
         try:
             data_model = self.f_type_validation_remedy(prompt, f_semantic_conditions=f_semantic_conditions, **remedy_kwargs)
         except Exception as e:
-            logger.error(f"Type validation failed with exception!")
+            logger.error("Type validation failed with exception!")
             raise e
         return data_model
 
-    def _validate_input(self, wrapped_self, input, it, **remedy_kwargs):
+    def _validate_input(self, wrapped_self, input_value, it, **remedy_kwargs):
         logger.info("Starting input validation...")
         if self.pre_remedy:
             logger.info("Validating pre-conditions with remedy...")
             if not hasattr(wrapped_self, 'pre'):
                 logger.error("Pre-condition function not defined!")
-                raise Exception("Pre-condition function not defined. Please define a `pre` method if you want to enforce pre-conditions through a remedy.")
+                msg = "Pre-condition function not defined. Please define a `pre` method if you want to enforce pre-conditions through a remedy."
+                UserMessage(msg)
+                raise Exception(msg)
 
             op_start = time.perf_counter()
             try:
-                assert wrapped_self.pre(input)
+                assert wrapped_self.pre(input_value)
                 logger.success("Pre-condition validation successful!")
-                return input
-            except Exception as e:
+                return input_value
+            except Exception:
                 logger.exception("Pre-condition validation failed!")
-                self.f_type_validation_remedy.register_expected_data_model(input, attach_to="output", override=True)
-                input = self._try_remedy_with_exception(prompt=wrapped_self.prompt, f_semantic_conditions=[wrapped_self.pre], **remedy_kwargs)
+                self.f_type_validation_remedy.register_expected_data_model(input_value, attach_to="output", override=True)
+                input_value = self._try_remedy_with_exception(
+                    prompt=wrapped_self.prompt,
+                    f_semantic_conditions=[wrapped_self.pre],
+                    **remedy_kwargs,
+                )
             finally:
                 wrapped_self._contract_timing[it]["input_validation"] = time.perf_counter() - op_start
-            return input
-        else:
-            if hasattr(wrapped_self, 'pre'):
-                logger.info("Validating pre-conditions without remedy...")
-                op_start = time.perf_counter()
-                try:
-                    assert wrapped_self.pre(input)
-                except Exception as e:
-                    logger.exception(f"Pre-condition validation failed")
-                    raise e
-                finally:
-                    wrapped_self._contract_timing[it]["input_validation"] = time.perf_counter() - op_start
-                logger.success("Pre-condition validation successful!")
-                return input
-            logger.info("Skip; no pre-condition validation was required!")
-            return input
+            return input_value
+        if hasattr(wrapped_self, 'pre'):
+            logger.info("Validating pre-conditions without remedy...")
+            op_start = time.perf_counter()
+            try:
+                assert wrapped_self.pre(input_value)
+            except Exception as e:
+                logger.exception("Pre-condition validation failed")
+                raise e
+            finally:
+                wrapped_self._contract_timing[it]["input_validation"] = time.perf_counter() - op_start
+            logger.success("Pre-condition validation successful!")
+            return input_value
+        logger.info("Skip; no pre-condition validation was required!")
+        return input_value
 
-    def _validate_output(self, wrapped_self, input, output, it, **remedy_kwargs):
+    def _validate_output(self, wrapped_self, input_value, output, it, **remedy_kwargs):
         logger.info("Starting output validation...")
-        self.f_type_validation_remedy.register_expected_data_model(input, attach_to="input", override=True)
+        self.f_type_validation_remedy.register_expected_data_model(input_value, attach_to="input", override=True)
         self.f_type_validation_remedy.register_expected_data_model(output, attach_to="output", override=True)
 
         op_start = time.perf_counter()
@@ -505,7 +599,7 @@ class contract:
             if output is None: # output is None when graceful mode is enabled
                 return output
         except Exception as e:
-            logger.exception(f"Type creation failed!")
+            logger.exception("Type creation failed!")
             raise e
         finally:
             wrapped_self._contract_timing[it]["output_validation"] = time.perf_counter() - op_start
@@ -515,88 +609,111 @@ class contract:
             logger.info("Validating post-conditions with remedy...")
             if not hasattr(wrapped_self, "post"):
                 logger.error("Post-condition function not defined!")
-                raise Exception("Post-condition function not defined. Please define a `post` method if you want to enforce post-conditions through a remedy.")
+                msg = "Post-condition function not defined. Please define a `post` method if you want to enforce post-conditions through a remedy."
+                UserMessage(msg)
+                raise Exception(msg)
 
             op_start = time.perf_counter()
             try:
                 assert wrapped_self.post(output)
                 logger.success("Post-condition validation successful!")
                 return output
-            except Exception as e:
+            except Exception:
                 logger.exception("Post-condition validation failed!")
                 output = self._try_remedy_with_exception(prompt=wrapped_self.prompt, f_semantic_conditions=[wrapped_self.post], **remedy_kwargs)
             finally:
                 wrapped_self._contract_timing[it]["output_validation"] += (time.perf_counter() - op_start)
             logger.success("Post-condition validation successful!")
             return output
-        else:
-            if hasattr(wrapped_self, "post"):
-                logger.info("Validating post-conditions without remedy...")
-                op_start = time.perf_counter()
-                try:
-                    assert wrapped_self.post(output)
-                except Exception as e:
-                    logger.exception("Post-condition validation failed!")
-                    raise e
-                finally:
-                    wrapped_self._contract_timing[it]["output_validation"] = time.perf_counter() - op_start
-                logger.success("Post-condition validation successful!")
-                return output
+        if hasattr(wrapped_self, "post"):
+            logger.info("Validating post-conditions without remedy...")
+            op_start = time.perf_counter()
+            try:
+                assert wrapped_self.post(output)
+            except Exception as e:
+                logger.exception("Post-condition validation failed!")
+                raise e
+            finally:
+                wrapped_self._contract_timing[it]["output_validation"] = time.perf_counter() - op_start
+            logger.success("Post-condition validation successful!")
+            return output
         logger.info("Skip; no post-condition validation was required!")
         return output
 
     def _validate_act_method(self, act_method):
         act_sig = inspect.signature(act_method)
         params = list(act_sig.parameters.values())
-        if not params or params[0].name != 'input':
-            raise TypeError("'act' method first parameter must be named 'input'")
-        if params[0].annotation == inspect._empty:
-            raise TypeError("'act' method parameter 'input' must have a type annotation'")
+
+        first_param = None
+        for param in params:
+            if param.name == "self":
+                continue
+            if param.kind in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                first_param = param
+                break
+
+        if first_param is None:
+            msg = "'act' method must accept at least one positional parameter after `self`."
+            UserMessage(msg)
+            raise TypeError(msg)
+        if first_param.annotation == inspect._empty:
+            msg = f"'act' method parameter '{first_param.name}' must have a type annotation."
+            UserMessage(msg)
+            raise TypeError(msg)
         if act_sig.return_annotation == inspect._empty:
-            raise TypeError("'act' method must have a return type annotation'")
+            msg = "'act' method must have a return type annotation'"
+            UserMessage(msg)
+            raise TypeError(msg)
         return True
 
-    def _act(self, wrapped_self, input, it, **act_kwargs):
+    def _act(self, wrapped_self, input_value, it, **act_kwargs):
         act_method = getattr(wrapped_self, 'act', None)
         if not callable(act_method):
             # Propagate the input if no act method is defined
-            return input
+            return input_value
 
         assert self._validate_act_method(act_method)
 
-        is_dynamic_model = getattr(input, '_is_dynamic_model', False)
-        input = input if not is_dynamic_model else input.value
+        is_dynamic_model = getattr(input_value, '_is_dynamic_model', False)
+        input_value = input_value if not is_dynamic_model else input_value.value
 
         logger.info(f"Executing 'act' method on {wrapped_self.__class__.__name__}…")
 
         op_start = time.perf_counter()
         try:
-            act_output = act_method(input, **act_kwargs)
+            act_output = act_method(input_value, **act_kwargs)
         except Exception as e:
-            logger.exception(f"'act' method execution failed")
+            logger.exception("'act' method execution failed")
             raise e
         finally:
             wrapped_self._contract_timing[it]["act_execution"] = time.perf_counter() - op_start
 
         act_sig = inspect.signature(act_method)
-        if act_sig.return_annotation != inspect.Signature.empty and inspect.isclass(act_sig.return_annotation):
-            if not isinstance(act_output, act_sig.return_annotation):
-                raise TypeError(f"'act' method returned {type(act_output).__name__}, expected {act_sig.return_annotation.__name__}.")
+        if (
+            act_sig.return_annotation != inspect.Signature.empty
+            and inspect.isclass(act_sig.return_annotation)
+            and not isinstance(act_output, act_sig.return_annotation)
+        ):
+            msg = f"'act' method returned {type(act_output).__name__}, expected {act_sig.return_annotation.__name__}."
+            UserMessage(msg)
+            raise TypeError(msg)
 
         logger.success("'act' method executed successfully!")
         return act_output
 
-    def __call__(self, cls):
-        original_init = cls.__init__
-        original_forward = cls.forward
-
+    def _build_wrapped_init(self, original_init):
         def __init__(wrapped_self, *args, **kwargs):
             logger.info("Initializing contract...")
             original_init(wrapped_self, *args, **kwargs)
 
             if not hasattr(wrapped_self, "prompt"):
                 logger.error("Prompt attribute not defined!")
-                raise Exception("Please define a static `prompt` attribute that describes what the contract must do.")
+                msg = "Please define a static `prompt` attribute that describes what the contract must do."
+                UserMessage(msg)
+                raise Exception(msg)
 
             wrapped_self.contract_successful = False
             wrapped_self.contract_result = None
@@ -604,105 +721,207 @@ class contract:
             wrapped_self._contract_timing = defaultdict(dict)
             logger.info("Contract initialization complete!")
 
-        def wrapped_forward(wrapped_self, **kwargs):
-            logger.info("Starting contract execution...")
-            it = len(wrapped_self._contract_timing) # the len is the __call__ op_start
-            contract_start = time.perf_counter()
+        return __init__
+
+    def _start_contract_execution(self, wrapped_self):
+        logger.info("Starting contract execution...")
+        it = len(wrapped_self._contract_timing)  # the len is the __call__ op_start
+        return it, time.perf_counter()
+
+    def _find_input_param_name(self, sig: inspect.Signature) -> str | None:
+        for param in sig.parameters.values():
+            if param.name == "self":
+                continue
+            if param.kind in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            ):
+                return param.name
+        return None
 
-            input = kwargs.pop("input", None)
-            input_type = None
-            try:
-                assert self._is_valid_input(input)
-            except TypeError:
-                input_type = self._try_dynamic_type_annotation(original_forward, context="input")
-                input = input_type(value=input)
-
-            maybe_payload = getattr(wrapped_self, "payload", None)
-            maybe_template = getattr(wrapped_self, "template")
-            if inspect.ismethod(maybe_template):
-                # `template` is a primitive in symbolicai case in which we actually don't have a template
-                maybe_template = None
-
-            # Create validation kwargs that include all original kwargs plus payload and template
-            validation_kwargs = {
-                **kwargs,
-                "payload": maybe_payload,
-                "template_suffix": maybe_template
-            }
+    def _prepare_forward_args(self, args, kwargs):
+        args_list = list(args)
+        kwargs_without_input = dict(kwargs)
+        original_kwargs = dict(kwargs)
+        return args_list, kwargs_without_input, original_kwargs
 
-            sig = inspect.signature(original_forward)
-            output_type = sig.return_annotation
-            try:
-                assert self._is_valid_output(output_type)
-            except TypeError:
-                output_type = self._try_dynamic_type_annotation(original_forward, context="output")
+    def _extract_input_value(
+        self,
+        args_list,
+        kwargs_without_input,
+        original_kwargs,
+        input_param_name: str | None,
+    ):
+        if args_list:
+            return args_list[0], ("args", 0)
+        if input_param_name and input_param_name in kwargs_without_input:
+            return kwargs_without_input.pop(input_param_name), ("kwargs", input_param_name)
+        if "input" in kwargs_without_input:
+            return kwargs_without_input.pop("input"), ("kwargs", "input")
+        return original_kwargs.get("input"), ("fallback_kw", "input")
+
+    def _coerce_input_value(self, original_forward, input_value):
+        try:
+            assert self._is_valid_input(input_value)
+            return input_value
+        except TypeError:
+            input_type_model = self._try_dynamic_type_annotation(original_forward, context="input")
+            return input_type_model(value=input_value)
+
+    def _collect_validation_kwargs(self, wrapped_self, kwargs_without_input):
+        maybe_payload = getattr(wrapped_self, "payload", None)
+        maybe_template = getattr(wrapped_self, "template", None)
+        if inspect.ismethod(maybe_template):
+            maybe_template = None
+        return {
+            **kwargs_without_input,
+            "payload": maybe_payload,
+            "template_suffix": maybe_template,
+        }
+
+    def _resolve_output_type(self, sig: inspect.Signature, original_forward):
+        output_type = sig.return_annotation
+        try:
+            assert self._is_valid_output(output_type)
+        except TypeError:
+            output_type = self._try_dynamic_type_annotation(original_forward, context="output")
+        return output_type
 
-            output = None
-            current_input = input
-            try:
-                # 1. Start with original input and apply pre-validation
-                maybe_new_input = self._validate_input(wrapped_self, current_input, it, **validation_kwargs)
-                if maybe_new_input is not None:
-                    current_input = maybe_new_input
+    def _run_contract_pipeline(
+        self,
+        wrapped_self,
+        current_input_value,
+        output_type,
+        it,
+        validation_kwargs,
+    ):
+        output = None
+        try:
+            maybe_new_input = self._validate_input(wrapped_self, current_input_value, it, **validation_kwargs)
+            if maybe_new_input is not None:
+                current_input_value = maybe_new_input
+
+            current_input_value = self._act(wrapped_self, current_input_value, it, **validation_kwargs)
+
+            output = self._validate_output(
+                wrapped_self,
+                current_input_value,
+                output_type,
+                it,
+                **validation_kwargs,
+            )
+            wrapped_self.contract_successful = output is not None
+            wrapped_self.contract_result = output
+            wrapped_self.contract_exception = None
+        except Exception as exc:
+            logger.exception("Contract execution failed in main path!")
+            wrapped_self.contract_successful = False
+            wrapped_self.contract_exception = exc
+        return output, current_input_value
 
-                # 2. Check if 'act' method exists and execute it
-                current_input = self._act(wrapped_self, current_input, it, **validation_kwargs)
+    def _execute_forward_call(
+        self,
+        wrapped_self,
+        original_forward,
+        args_list,
+        original_kwargs,
+        input_param_name,
+        input_source,
+        forward_input_value,
+        it,
+        contract_start,
+    ):
+        forward_kwargs = original_kwargs.copy()
+        for internal_kw in self._internal_forward_kwargs:
+            forward_kwargs.pop(internal_kw, None)
+
+        logger.info("Executing original forward method...")
+
+        if input_param_name:
+            if input_param_name in forward_kwargs or input_source == ("kwargs", input_param_name):
+                forward_kwargs[input_param_name] = forward_input_value
+            elif input_source == ("args", 0) and args_list:
+                args_list[0] = forward_input_value
+            else:
+                forward_kwargs[input_param_name] = forward_input_value
+        else:
+            forward_kwargs['input'] = forward_input_value
 
-                # 3. Validate output type and prepare for original_forward
-                output = self._validate_output(wrapped_self, current_input, output_type, it, **validation_kwargs)
-                wrapped_self.contract_successful = True if output is not None else False
-                wrapped_self.contract_result = output
-                wrapped_self.contract_exception = None
+        if input_param_name and input_param_name != "input" and "input" in forward_kwargs:
+            forward_kwargs.pop("input")
 
-            except Exception as e:
-                logger.exception(f"Contract execution failed in main path!")
-                wrapped_self.contract_successful = False
-                wrapped_self.contract_exception = e
-                # contract_result remains None or its value before the exception.
-                # final_output remains None or its value before the exception.
-                # The finally block's execution of original_forward will determine the actual returned value.
-            finally:
-                # Execute the original forward method with appropriate input
-                logger.info("Executing original forward method...")
-
-                # If contract was successful, use the processed input (after pre-validation and act, both optional)
-                # `current_input` at this stage is the result of the try block's processing up to the point of exception,
-                # or the full processing if successful.
-                # If contract failed, use original_input (fallback).
-                forward_input = current_input if wrapped_self.contract_successful else input
-
-                # Prepare kwargs for original_forward
-                forward_kwargs = kwargs.copy()
-                forward_kwargs['input'] = forward_input
-
-                try:
-                    op_start = time.perf_counter()
-                    output = original_forward(wrapped_self, **forward_kwargs)
-                finally:
-                    wrapped_self._contract_timing[it]["forward_execution"] = time.perf_counter() - op_start
-                wrapped_self._contract_timing[it]["contract_execution"] = time.perf_counter() - contract_start
-
-                if not isinstance(output, output_type):
-                    logger.error(f"Output type mismatch: {type(output)}")
-                    if self.remedy_retry_params["graceful"]:
-                        # In graceful mode, skip type mismatch error and return raw output
-                        if hasattr(output_type, '_is_dynamic_model') and output_type._is_dynamic_model and hasattr(output, 'value'):
-                            return output.value
-                        return output
-                    raise TypeError(
-                        f"Expected output to be an instance of {output_type}, "
-                        f"but got {type(output)}! Forward method must return an instance of {output_type}!"
-                    )
-                if not wrapped_self.contract_successful:
-                    logger.warning("Contract validation failed!")
-                else:
-                    logger.success("Contract validation successful!")
+        try:
+            op_start = time.perf_counter()
+            output = original_forward(wrapped_self, *args_list, **forward_kwargs)
+        finally:
+            wrapped_self._contract_timing[it]["forward_execution"] = time.perf_counter() - op_start
+        wrapped_self._contract_timing[it]["contract_execution"] = time.perf_counter() - contract_start
+        return output
 
-            if hasattr(output_type, '_is_dynamic_model') and output_type._is_dynamic_model:
-                return output.value
+    def _finalize_contract_output(self, output, output_type, wrapped_self):
+        if not isinstance(output, output_type):
+            logger.error(f"Output type mismatch: {type(output)}")
+            if self.remedy_retry_params["graceful"]:
+                if getattr(output_type, '_is_dynamic_model', False) and hasattr(output, 'value'):
+                    return output.value
+                return output
+            msg = (
+                f"Expected output to be an instance of {output_type}, "
+                f"but got {type(output)}! Forward method must return an instance of {output_type}!"
+            )
+            UserMessage(msg)
+            raise TypeError(msg)
+        if not wrapped_self.contract_successful:
+            logger.warning("Contract validation failed!")
+        else:
+            logger.success("Contract validation successful!")
 
-            return output
+        if getattr(output_type, '_is_dynamic_model', False):
+            return output.value
+        return output
 
+    def _contract_forward_impl(self, wrapped_self, original_forward, *args, **kwargs):
+        it, contract_start = self._start_contract_execution(wrapped_self)
+        sig = inspect.signature(original_forward)
+        input_param_name = self._find_input_param_name(sig)
+        args_list, kwargs_without_input, original_kwargs = self._prepare_forward_args(args, kwargs)
+        input_value, input_source = self._extract_input_value(args_list, kwargs_without_input, original_kwargs, input_param_name)
+        current_input_value = self._coerce_input_value(original_forward, input_value)
+        input_value = current_input_value
+        validation_kwargs = self._collect_validation_kwargs(wrapped_self, kwargs_without_input)
+        output_type = self._resolve_output_type(sig, original_forward)
+
+        output, current_input_value = self._run_contract_pipeline(
+            wrapped_self,
+            current_input_value,
+            output_type,
+            it,
+            validation_kwargs,
+        )
+
+        forward_input_value = current_input_value if wrapped_self.contract_successful else input_value
+        output = self._execute_forward_call(
+            wrapped_self,
+            original_forward,
+            args_list,
+            original_kwargs,
+            input_param_name,
+            input_source,
+            forward_input_value,
+            it,
+            contract_start,
+        )
+
+        return self._finalize_contract_output(output, output_type, wrapped_self)
+
+    def _build_wrapped_forward(self, original_forward):
+        def wrapped_forward(wrapped_self, *args, **kwargs):
+            return self._contract_forward_impl(wrapped_self, original_forward, *args, **kwargs)
+
+        return wrapped_forward
+
+    def _build_contract_perf_stats(self):
         def contract_perf_stats(wrapped_self):
             """Analyzes and prints timing statistics across all forward calls."""
             console = Console()
@@ -828,17 +1047,29 @@ class contract:
 
             return stats
 
-        cls.__init__ = __init__
-        cls.forward = wrapped_forward
-        cls.contract_perf_stats = contract_perf_stats
+        return contract_perf_stats
+
+    def __call__(self, cls):
+        original_init = cls.__init__
+        original_forward = cls.forward
 
+        cls.__init__ = self._build_wrapped_init(original_init)
+        cls.forward = self._build_wrapped_forward(original_forward)
+        cls.contract_perf_stats = self._build_contract_perf_stats()
         return cls
 
 
 class BaseStrategy(TypeValidationFunction):
-    def __init__(self, data_model: BaseModel, *args, **kwargs):
+    def __init__(self, data_model: BaseModel, *_args, **kwargs):
         super().__init__(
-            retry_params=dict(tries=8, delay=0.015, backoff=1.25, jitter=0.0, max_delay=0.25, graceful=False),
+            retry_params={
+                "tries": 8,
+                "delay": 0.015,
+                "backoff": 1.25,
+                "jitter": 0.0,
+                "max_delay": 0.25,
+                "graceful": False,
+            },
             **kwargs,
         )
         super().register_expected_data_model(data_model, attach_to="output")
@@ -851,14 +1082,13 @@ class BaseStrategy(TypeValidationFunction):
         pass
 
     def forward(self, *args, **kwargs):
-        result = super().forward(
+        return super().forward(
             *args,
             payload=self.payload,
             template_suffix=self.template,
             response_format={"type": "json_object"},
             **kwargs,
         )
-        return result
 
     @property
     def payload(self):
@@ -866,7 +1096,7 @@ class BaseStrategy(TypeValidationFunction):
 
     @property
     def static_context(self):
-        raise NotImplementedError()
+        raise NotImplementedError
 
     @property
     def template(self):
@@ -878,7 +1108,7 @@ class Strategy(Expression):
         super().__init__(*args, **kwargs)
         self.logger = logging.getLogger(__name__)
 
-    def __new__(self, module: str, *args, **kwargs):
-        self._module = module
-        self.module_path = 'symai.extended.strategies'
-        return Strategy.load_module_class(self.module)
+    def __new__(cls, module: str, *_args, **_kwargs):
+        cls._module = module
+        cls.module_path = 'symai.extended.strategies'
+        return Strategy.load_module_class(cls.module)