Align local policy prompts and canonical action targets

environment.py

@@ -12,6 +12,14 @@ from models import ActionTypeEnum, FraudCheckAction, ResolutionEnum
 from reward import RewardBreakdown, build_reward_breakdown
 from utils import approximate_token_count, extract_json_object
 
+CANONICAL_INVESTIGATION_ALIASES = [
+    "merchant_profile",
+    "customer_profile",
+    "network_graph",
+    "payment_trace",
+    "policy_review",
+]
+
 INVESTIGATION_ALIAS_TO_ACTION = {
     "merchant_profile": ActionTypeEnum.FETCH_MERCHANT_PROFILE,
     "fetch_merchant_profile": ActionTypeEnum.FETCH_MERCHANT_PROFILE,
@@ -27,6 +35,42 @@ INVESTIGATION_ALIAS_TO_ACTION = {
     "check_policy": ActionTypeEnum.CHECK_POLICY,
 }
 
+ACTION_TYPE_TO_CANONICAL_ALIAS = {
+    ActionTypeEnum.FETCH_MERCHANT_PROFILE: "merchant_profile",
+    ActionTypeEnum.FETCH_CUSTOMER_PROFILE: "customer_profile",
+    ActionTypeEnum.FETCH_NETWORK_GRAPH: "network_graph",
+    ActionTypeEnum.REVIEW_TRANSACTION: "payment_trace",
+    ActionTypeEnum.CHECK_POLICY: "policy_review",
+}
+
+
+def build_fraudshield_prompt(observation) -> str:
+    """Build the canonical prompt used for both training and inference."""
+
+    payload = {
+        "case_id": observation.case_id,
+        "task_name": observation.task_name.value,
+        "visible_panels": observation.visible_panels,
+        "revealed_evidence": observation.revealed_evidence,
+        "linked_case_ids": observation.linked_case_ids,
+        "remaining_steps": observation.remaining_steps,
+        "remaining_sla": observation.remaining_sla,
+        "note_required": observation.note_required,
+        "allowed_actions": [action.value for action in observation.allowed_actions],
+        "case_summary": observation.case_summary.model_dump(mode="json"),
+        "app_context": observation.app_context,
+    }
+    available = observation.app_context.get("available_investigations", CANONICAL_INVESTIGATION_ALIASES)
+    return (
+        "You are a fraud analyst in a multi-step training environment. "
+        "Return JSON only. Use visible evidence, investigation budget, and prior evidence carefully.\n\n"
+        f"Visible observation:\n{json.dumps(payload, sort_keys=True)}\n\n"
+        f"Valid investigation aliases: {available}\n"
+        "JSON schema: "
+        '{"action_type":"investigate|decide","investigation_target":"alias_or_null",'
+        '"decision":"fraud|legitimate|null","confidence":0.0,"reasoning":"one sentence"}'
+    )
+
 
 @dataclass
 class TextStepResult:
@@ -70,33 +114,7 @@ class FraudShieldTextEnvironment:
 
     def build_prompt(self, observation) -> str:
         """Build the prompt shown to an LLM policy."""
-
-        payload = {
-            "case_id": observation.case_id,
-            "task_name": observation.task_name.value,
-            "visible_panels": observation.visible_panels,
-            "revealed_evidence": observation.revealed_evidence,
-            "linked_case_ids": observation.linked_case_ids,
-            "remaining_steps": observation.remaining_steps,
-            "remaining_sla": observation.remaining_sla,
-            "note_required": observation.note_required,
-            "allowed_actions": [action.value for action in observation.allowed_actions],
-            "case_summary": observation.case_summary.model_dump(mode="json"),
-            "app_context": observation.app_context,
-        }
-        available = observation.app_context.get(
-            "available_investigations",
-            ["merchant_profile", "customer_profile", "network_graph", "payment_trace", "policy_review"],
-        )
-        return (
-            "You are a fraud analyst in a multi-step training environment. "
-            "Return JSON only. Use visible evidence, investigation budget, and prior evidence carefully.\n\n"
-            f"Visible observation:\n{json.dumps(payload, sort_keys=True)}\n\n"
-            f"Valid investigation aliases: {available}\n"
-            "JSON schema: "
-            '{"action_type":"investigate|decide","investigation_target":"alias_or_null",'
-            '"decision":"fraud|legitimate|null","confidence":0.0,"reasoning":"one sentence"}'
-        )
+        return build_fraudshield_prompt(observation)
 
     def parse_response(self, response_text: str) -> tuple[FraudCheckAction, dict[str, Any], bool, bool]:
         """Convert model output into a typed environment action."""

llm_agent_openai.py

@@ -8,6 +8,7 @@ import re
 from pathlib import Path
 from typing import Any, Dict, Optional
 
+from environment import ACTION_TYPE_TO_CANONICAL_ALIAS, build_fraudshield_prompt
 from models import ActionTypeEnum, FraudCheckAction, ResolutionEnum, TaskDifficulty
 
 try:  # pragma: no cover - optional in local smoke tests
@@ -95,21 +96,7 @@ class LLMFraudDetectionAgent:
 
     def _build_messages(self, observation) -> list[Dict[str, str]]:
         available_aliases = self._available_investigation_aliases(observation)
-        observation_payload = {
-            "case_id": observation.case_id,
-            "task_name": observation.task_name.value,
-            "current_screen": observation.current_screen.value,
-            "visible_panels": observation.visible_panels,
-            "case_summary": observation.case_summary.model_dump(mode="json"),
-            "revealed_evidence": observation.revealed_evidence,
-            "linked_case_ids": observation.linked_case_ids,
-            "remaining_steps": observation.remaining_steps,
-            "remaining_sla": observation.remaining_sla,
-            "note_required": observation.note_required,
-            "allowed_public_actions": [action.value for action in observation.allowed_actions],
-            "available_investigation_aliases": available_aliases,
-            "app_context": observation.app_context,
-        }
+        prompt = build_fraudshield_prompt(observation)
         system_prompt = (
             "You are a fraud analyst operating inside a simulated investigation workflow. "
             "Only use the visible evidence shown to you. Choose either one investigation alias or one final "
@@ -121,7 +108,7 @@ class LLMFraudDetectionAgent:
         )
         return [
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": json.dumps(observation_payload, separators=(",", ":"))},
+            {"role": "user", "content": prompt},
         ]
 
     def _payload_to_action(self, payload: Dict[str, Any], observation) -> FraudCheckAction:
@@ -316,10 +303,7 @@ class LocalModelFraudDetectionAgent(LLMFraudDetectionAgent):
             return self._fallback(observation, exc)
 
     def _build_local_prompt(self, observation) -> str:
-        messages = self._build_messages(observation)
-        if hasattr(self.tokenizer, "apply_chat_template"):
-            return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        return "\n".join(f"{message['role'].upper()}: {message['content']}" for message in messages)
+        return build_fraudshield_prompt(observation) + "\n"
 
     def _load_model(self) -> None:
         try:

train.py

@@ -13,7 +13,7 @@ import pandas as pd
 from datasets import Dataset, load_dataset
 
 from config import ExperimentConfig
-from environment import FraudShieldTextEnvironment
+from environment import ACTION_TYPE_TO_CANONICAL_ALIAS, FraudShieldTextEnvironment
 from models import ActionTypeEnum, FraudCheckAction
 from utils import ensure_dir, save_json, seed_everything
 
@@ -171,7 +171,7 @@ def build_rollout_dataset(config: ExperimentConfig) -> Dataset:
                 action = agent.decide(text_env)
                 payload = {
                     "action_type": "decide" if action.action_type.value == "resolve_case" else "investigate",
-                    "investigation_target": action.action_type.value,
+                    "investigation_target": ACTION_TYPE_TO_CANONICAL_ALIAS.get(action.action_type),
                     "decision": "fraud" if getattr(action, "resolution", None) and action.resolution.value in {"block", "hold", "escalate"} else "legitimate",
                     "confidence": 0.8,
                     "reasoning": action.reasoning or "Training rollout step.",