server/tasks/task2/actions.py

"""
Actions for Task 2: Property Inference.
Defines the logic for each action type that the agent can take in the environment, including
querying function code, NatSpec, related functions, and submitting the inferred property.
ctx: context object containing episode state and data
qkey: a unique key representing the specific query (used for tracking repeated queries)
params: parameters for the action, such as the submitted property text for SUBMIT_PROPERTY
"""

from typing import Any, Dict, Tuple
from data.data_loader import get_function_by_name, get_related_functions
from utils import PropertyRetriever
from env.schemas import ActionType, Reward

PropertyRetrieverInstance = PropertyRetriever()  # Load once at module level

# TODO: Can separate into signature, visiblity

def get_function_code(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_FUNCTION_CODE action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")

    fn = ctx._target_fn
    code = fn.get("code", "// no code available")
    return (
        code,
        Reward(value=-0.06, reason="get_function_code cost"),
    )

# TODO: Can separate comment and output_property(output_comment)

def get_function_natspec(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_FUNCTION_NATSPEC action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")

    fn = ctx._target_fn
    name = fn["name"]
    natspec = fn.get("natspec") or fn.get("comment") or "No NatSpec available."
    out_prop = fn.get("output_property", "")
    result = f"NatSpec for '{name}':\n{natspec}"
    if out_prop:
        result += f"\n\nExpected output: {out_prop}"
    return result, Reward(value=-0.08, reason="get_function_natspec cost")


def get_file_natspec(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_FILE_NATSPEC action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")

    meta = ctx._contract.get("metadata", {})
    natspec = meta.get("natspec") or meta.get("description", "No file NatSpec available.")
    return (
        f"File NatSpec for {ctx._contract['contract_name']}:\n{natspec}",
        Reward(value=-0.03, reason="get_file_natspec cost"),
    )


def get_related_functions_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_RELATED_FUNCTIONS action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")

    name = ctx._target_fn["name"]
    related = get_related_functions(ctx._contract, name)
    if not related:
        text = f"No related functions found for '{name}'."
    else:
        summaries = []
        for rn in related:
            rfn = get_function_by_name(ctx._contract, rn)
            if rfn:
                sig = rfn.get("signature", rn)
                comment = rfn.get("comment", "")
                summaries.append(f"  • {sig} — {comment}")
        text = f"Related functions for '{name}':\n" + "\n".join(summaries)
    return text, Reward(value=-0.06, reason="get_related_functions cost")


def get_signature(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_SIGNATURE action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")

    fn = ctx._target_fn
    sig = fn.get("signature")
    return sig, Reward(value=-0.04, reason="get_signature cost")


def get_similar_rule_action(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle GET_SIMILAR_RULE action."""
    if ctx._is_repeated(qkey):
        return "Repeated query.", Reward(value=-0.40, reason="Repeated query")
    
    PropertyRetrieverInstance.load_model()  # Ensure model is loaded before querying
    similar_rule = PropertyRetrieverInstance.get_similar_property(ctx._target_fn["code"])
    if similar_rule is None:
        return (
            "No similar rule available for this function.",
            Reward(value=-0.20, reason="get_similar_rule cost (not found)"),
        )
    return similar_rule, Reward(value=-0.20, reason="get_similar_rule cost")


def submit_property(ctx: Any, qkey: str, params: Dict) -> Tuple[str, Reward]:
    """Handle SUBMIT_PROPERTY action."""
    if ctx._submitted:
        return (
            "❌ You have already submitted a property for this episode. "
            "Only one submission is allowed.",
            Reward(value=-1.0, reason="Second submit_property attempt", partial=False),
        )

    submitted_text = params.get("property", "").strip()
    if not submitted_text:
        return (
            "Submit requires 'property' key in params with a non-empty string.",
            Reward(value=-0.5, reason="Empty property submission"),
        )

    ctx._submitted = True
    ctx._done = True
    score, confidence = ctx._grader.grade(submitted_text)
    reward = round(score * 5.0, 4)

    msg = f'Score: {score:.2f}/1.00 → Confidence: {confidence}\n'    
    return msg, Reward(
        value=reward,
        reason=f"Property submission score={score:.3f}",
        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}'. Valid: {[a.value for a in ActionType]}",
        Reward(value=-0.10, reason="Unknown action"),
    )