Task1 actions reviewed

Task1 actions separated

tasks/task1/actions.py

@@ -0,0 +1,161 @@
+"""Actions for Task 1: Targeted Vulnerability Detection.
+    Actions & rewards:
+    list_functions       -0.05  (broad overview of contract)
+    get_function_code    -0.10 (wrong function) / +0.05 (correct function)
+    get_function_summary -0.05 (wrong function) / +0.03 (correct function)
+    get_file_metadata    -0.04  (general contract info)
+"""
+
+from typing import Any, Dict, Tuple
+from env.schemas import Reward
+from data.data_loader import (
+    list_function_names,
+    get_function_by_name,
+    list_state_variable_names,
+    get_state_variable_by_name,
+)
+
+
+def list_functions(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle LIST_FUNCTIONS action."""
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    names = list_function_names(ctx._contract)
+    return (
+        f"Functions in {ctx._contract['contract_name']}: {', '.join(names)}",
+        Reward(value=-0.05, reason="list_functions cost", partial=True),
+    )
+
+
+def get_function_code(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle GET_FUNCTION_CODE action."""
+    fn_name = params.get("function_name", "")
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    fn = get_function_by_name(ctx._contract, fn_name)
+    if fn is None:
+        return (
+            f"Function '{fn_name}' not found. Available: {list_function_names(ctx._contract)}",
+            Reward(value=-0.10, reason="Wrong/unknown function name", partial=True),
+        )
+    is_target = fn["name"].lower() == ctx._target_fn["name"].lower()
+    code = fn.get("code", "// no code available")
+    reward_val = 0.05 if is_target else -0.10
+    reason = "Fetched target function code (+)" if is_target else "Fetched non-target function (-)"
+    return (
+        f"// {fn['name']}\n{code}",
+        Reward(value=reward_val, reason=reason, partial=True),
+    )
+
+
+def get_function_summary(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle GET_FUNCTION_SUMMARY action."""
+    fn_name = params.get("function_name", "")
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    fn = get_function_by_name(ctx._contract, fn_name)
+    if fn is None:
+        return (
+            f"Function '{fn_name}' not found.",
+            Reward(value=-0.05, reason="Wrong function name", partial=True),
+        )
+    is_target = fn["name"].lower() == ctx._target_fn["name"].lower()
+    comment = fn.get("comment", "No summary available.")
+    reward_val = 0.03 if is_target else -0.05
+    reason = "Fetched target function summary (+)" if is_target else "Fetched non-target summary (-)"
+    return (
+        f"Summary of '{fn['name']}': {comment}",
+        Reward(value=reward_val, reason=reason, partial=True),
+    )
+
+
+def get_file_metadata(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle GET_FILE_METADATA action."""
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    meta = ctx._contract.get("metadata", {})
+    result = (
+        f"Contract: {ctx._contract['contract_name']} | "
+        f"Solidity: {meta.get('solidity_version', 'N/A')} | "
+        f"Description: {meta.get('description', 'N/A')}"
+    )
+    return result, Reward(value=-0.04, reason="get_file_metadata cost", partial=True)
+
+
+def get_state_variable(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle GET_STATE_VARIABLE action."""
+    var_name = params.get("variable_name", "")
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    if not var_name:
+        names = list_state_variable_names(ctx._contract)
+        return (
+            f"State variables: {', '.join(names)}",
+            Reward(value=-0.05, reason="Listed state variables", partial=True),
+        )
+    sv = get_state_variable_by_name(ctx._contract, var_name)
+    if sv is None:
+        return (
+            f"Variable '{var_name}' not found.",
+            Reward(value=-0.05, reason="Unknown state variable", partial=True),
+        )
+    return (
+        f"{sv['type']} {sv['visibility']} {sv['name']}: {sv.get('description', '')}",
+        Reward(value=-0.05, reason="get_state_variable cost", partial=True),
+    )
+
+
+def get_call_graph(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle GET_CALL_GRAPH action."""
+    if ctx._is_repeated(qkey):
+        return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
+    cg = ctx._contract.get("call_graph", {})
+    cg_str = "; ".join(f"{fn} → [{', '.join(callees)}]" for fn, callees in cg.items())
+    return (
+        f"Call graph: {cg_str}",
+        Reward(value=-0.08, reason="get_call_graph cost", partial=True),
+    )
+
+
+def submit(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
+    """Handle SUBMIT action."""
+    fn_name = params.get("function_name", "")
+    vuln_type = params.get("vulnerability_type", "")
+    if not fn_name or not vuln_type:
+        return (
+            "Submit requires 'function_name' and 'vulnerability_type' in params.",
+            Reward(value=-0.5, reason="Malformed submission", partial=True),
+        )
+    score = ctx._grader.grade_submission(fn_name, vuln_type)
+    reward_val = ctx._grader.reward_for_score(score)
+    ctx._done = True
+
+    if score == 1.0:
+        msg = (
+            f"✅ CORRECT! '{fn_name}' is the vulnerable function. "
+            f"Vulnerability type '{vuln_type}' matches. Score: 1.0"
+        )
+    elif score == 0.5:
+        msg = (
+            f"⚠️  PARTIAL. '{fn_name}' is the right function, but the vulnerability type "
+            f"'{vuln_type}' was not precise. Score: 0.5"
+        )
+    else:
+        correct = ctx._grader.get_canonical_answer()
+        msg = (
+            f"❌ INCORRECT. '{fn_name}' is not the target vulnerable function. "
+            f"Correct answer: {correct['function']} ({correct['vulnerability']}). Score: 0.0"
+        )
+    return msg, Reward(
+        value=reward_val,
+        reason=f"Submission score={score:.1f}",
+        partial=False,
+    )
+
+
+def unknown_action(ctx: Any, qkey: str, params: Dict, action_type: str) -> Tuple[str, Reward]:
+    """Fallback for unknown actions."""
+    return (
+        f"Unknown action type: {action_type}",
+        Reward(value=-0.10, reason="Unknown action", partial=True),
+    )

tasks/task1/environment.py

@@ -10,16 +10,16 @@ Episode flow:
   4. When the agent submits, the Grader scores the answer and the episode ends.
 
 Reward shaping:
-  list_functions          : -0.05
-  get_function_code       : -0.10 (wrong function) / +0.05 (correct function)
-  get_function_summary    : -0.05 (wrong function) / +0.03 (correct function)
-  get_file_metadata       : -0.04
-  get_state_variable      : -0.05
-  get_call_graph          : -0.08
-  submit (score=1.0)      : +5.0
-  submit (score=0.5)      : +1.0
-  submit (score=0.0)      : -1.5
-  repeated query          : -0.40
+  list_functions                            : -0.05
+  get_function_code                         : -0.10 (wrong function) / +0.05 (correct function)
+  get_function_summary                      : -0.05 (wrong function) / +0.03 (correct function)
+  get_file_metadata                         : -0.04
+  get_state_variable                        : -0.05
+  get_call_graph                            : -0.08
+  correct submit (score=1.0)                : +5.0
+  partially correct submit (score=0.5)      : +1.0
+  wrong submit (score=0.0)                  : -1.5
+  repeated query                            : -0.40
 """
 
 from __future__ import annotations
@@ -46,6 +46,7 @@ from env.schemas import (
     StepResult,
 )
 from tasks.task1.grader import Task1Grader
+from tasks.task1 import actions
 
 TASK_ID = "task1_vuln_detection"
 
@@ -188,142 +189,19 @@ class Task1Environment(BaseEnv):
         params = action.params
         qkey = self._query_key(at, params)
 
-        # ---- list_functions ----------------------------------------
-        if at == ActionType.LIST_FUNCTIONS:
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            names = list_function_names(self._contract)
-            return (
-                f"Functions in {self._contract['contract_name']}: {', '.join(names)}",
-                Reward(value=-0.05, reason="list_functions cost", partial=True),
-            )
-
-        # ---- get_function_code -------------------------------------
-        if at == ActionType.GET_FUNCTION_CODE:
-            fn_name = params.get("function_name", "")
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            fn = get_function_by_name(self._contract, fn_name)
-            if fn is None:
-                return (
-                    f"Function '{fn_name}' not found. Available: {list_function_names(self._contract)}",
-                    Reward(value=-0.10, reason="Wrong/unknown function name", partial=True),
-                )
-            is_target = fn["name"].lower() == self._target_fn["name"].lower()
-            code = fn.get("code", "// no code available")
-            reward_val = 0.05 if is_target else -0.10
-            reason = "Fetched target function code (+)" if is_target else "Fetched non-target function (-)"
-            return (
-                f"// {fn['name']}\n{code}",
-                Reward(value=reward_val, reason=reason, partial=True),
-            )
-
-        # ---- get_function_summary ----------------------------------
-        if at == ActionType.GET_FUNCTION_SUMMARY:
-            fn_name = params.get("function_name", "")
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            fn = get_function_by_name(self._contract, fn_name)
-            if fn is None:
-                return (
-                    f"Function '{fn_name}' not found.",
-                    Reward(value=-0.05, reason="Wrong function name", partial=True),
-                )
-            is_target = fn["name"].lower() == self._target_fn["name"].lower()
-            comment = fn.get("comment", "No summary available.")
-            reward_val = 0.03 if is_target else -0.05
-            reason = "Fetched target function summary (+)" if is_target else "Fetched non-target summary (-)"
-            return (
-                f"Summary of '{fn['name']}': {comment}",
-                Reward(value=reward_val, reason=reason, partial=True),
-            )
-
-        # ---- get_file_metadata -------------------------------------
-        if at == ActionType.GET_FILE_METADATA:
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            meta = self._contract.get("metadata", {})
-            result = (
-                f"Contract: {self._contract['contract_name']} | "
-                f"File: {self._contract.get('file_name', 'N/A')} | "
-                f"Solidity: {meta.get('solidity_version', 'N/A')} | "
-                f"License: {meta.get('license', 'N/A')} | "
-                f"Author: {meta.get('author', 'N/A')} | "
-                f"Description: {meta.get('description', 'N/A')}"
-            )
-            return result, Reward(value=-0.04, reason="get_file_metadata cost", partial=True)
-
-        # ---- get_state_variable ------------------------------------
-        if at == ActionType.GET_STATE_VARIABLE:
-            var_name = params.get("variable_name", "")
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            if not var_name:
-                # Return list of all state variables
-                names = list_state_variable_names(self._contract)
-                return (
-                    f"State variables: {', '.join(names)}",
-                    Reward(value=-0.05, reason="Listed state variables", partial=True),
-                )
-            sv = get_state_variable_by_name(self._contract, var_name)
-            if sv is None:
-                return (
-                    f"Variable '{var_name}' not found.",
-                    Reward(value=-0.05, reason="Unknown state variable", partial=True),
-                )
-            return (
-                f"{sv['type']} {sv['visibility']} {sv['name']}: {sv.get('description', '')}",
-                Reward(value=-0.05, reason="get_state_variable cost", partial=True),
-            )
-
-        # ---- get_call_graph ----------------------------------------
-        if at == ActionType.GET_CALL_GRAPH:
-            if self._is_repeated(qkey):
-                return "Repeated query.", Reward(value=-0.40, reason="Repeated query", partial=True)
-            cg = self._contract.get("call_graph", {})
-            cg_str = "; ".join(f"{fn} → [{', '.join(callees)}]" for fn, callees in cg.items())
-            return (
-                f"Call graph: {cg_str}",
-                Reward(value=-0.08, reason="get_call_graph cost", partial=True),
-            )
-
-        # ---- submit ------------------------------------------------
-        if at == ActionType.SUBMIT:
-            fn_name = params.get("function_name", "")
-            vuln_type = params.get("vulnerability_type", "")
-            if not fn_name or not vuln_type:
-                return (
-                    "Submit requires 'function_name' and 'vulnerability_type' in params.",
-                    Reward(value=-0.5, reason="Malformed submission", partial=True),
-                )
-            score = self._grader.grade_submission(fn_name, vuln_type) # type: ignore
-            reward_val = self._grader.reward_for_score(score) # type: ignore
-            self._done = True
-
-            if score == 1.0:
-                msg = (
-                    f"✅ CORRECT! '{fn_name}' is the vulnerable function. "
-                    f"Vulnerability type '{vuln_type}' matches. Score: 1.0"
-                )
-            elif score == 0.5:
-                msg = (
-                    f"⚠️  PARTIAL. '{fn_name}' is the right function, but the vulnerability type "
-                    f"'{vuln_type}' was not precise. Score: 0.5"
-                )
-            else:
-                correct = self._grader.get_canonical_answer() # type: ignore
-                msg = (
-                    f"❌ INCORRECT. '{fn_name}' is not the target vulnerable function. "
-                    f"Correct answer: {correct['function']} ({correct['vulnerability']}). Score: 0.0"
-                )
-            return msg, Reward(
-                value=reward_val,
-                reason=f"Submission score={score:.1f}",
-                partial=False,
-            )
-
-        # ---- unknown action ----------------------------------------
-        return (
-            f"Unknown action type: {at}",
-            Reward(value=-0.10, reason="Unknown action", partial=True),
-        )
+        # Mapping from ActionType to handler function
+        handlers = {
+            ActionType.LIST_FUNCTIONS:       actions.list_functions,
+            ActionType.GET_FUNCTION_CODE:    actions.get_function_code,
+            ActionType.GET_FUNCTION_SUMMARY: actions.get_function_summary,
+            ActionType.GET_FILE_METADATA:    actions.get_file_metadata,
+            ActionType.GET_STATE_VARIABLE:   actions.get_state_variable,
+            ActionType.GET_CALL_GRAPH:       actions.get_call_graph,
+            ActionType.SUBMIT:               actions.submit,
+        }
+
+        handler = handlers.get(at)
+        if handler is None:
+            return actions.unknown_action(self, qkey, params, at)
+
+        return handler(self, qkey, params)

tasks/task2/environment.py

@@ -27,7 +27,6 @@ from __future__ import annotations
 
 import random
 from typing import Any, Dict, List, Optional, Set
-import actions
 
 from data.data_loader import load_contracts, sample_property_episode
 from env.base_env import BaseEnv
@@ -41,6 +40,7 @@ from env.schemas import (
     StepResult,
 )
 from tasks.task2.grader import Task2Grader
+from tasks.task2 import actions
 
 TASK_ID    = "task2_property_discovery"
 MAX_STEPS  = 15