Update app/bandit.py

app/bandit.py

@@ -1,56 +1,58 @@
 from __future__ import annotations
-import math
 import random
-from typing import Dict, Any, List, Tuple
+from typing import Dict, Any, Tuple, Callable, Optional
 from . import storage
 
 class ThompsonBandit:
     """
-    CTR（クリック）とCVR（コンバージョン）の二段ベータモデル。
-    目的関数は E[value] = p_click * p_conv * value_per_conversion。
-    季節性補正は、意思決定時に小さな事前カウント（virtual counts）として加算。
+    CTR（クリック）とCVR（コンバージョン）の二段ベータ。
+    目的関数: E[value] = p_click * p_conv * value_per_conversion
+    季節性は CTR 側の仮想カウントで補正。
     """
-
     def __init__(self, campaign_id: str, seasonality_boost: float = 5.0):
         self.campaign_id = campaign_id
         self.seasonality_boost = seasonality_boost
 
-    def sample_arm(self, context: Dict[str, Any], seasonal_fin) -> Tuple[str, float]:
+    def sample_arm(
+        self,
+        context: Dict[str, Any],
+        seasonal_fn: Optional[Callable[[Dict[str, Any]], float]] = None,  # ← キーワード引数OK
+    ) -> Tuple[Optional[str], float]:
         metrics = storage.get_metrics(self.campaign_id)
+        if not metrics:
+            return None, -1.0
+
         vpc = storage.get_campaign_value_per_conversion(self.campaign_id)
-        best_score = -1.0
-        best_variant = None
 
+        # 季節性スコア s ∈ (0, 1)
         s = 0.5
-        if seasonal_fin:
+        if seasonal_fn is not None:
             try:
-                s = float(seasonal_fin(context))
-                s = max(0.01, min(0.99,s))
+                s = float(seasonal_fn(context))
+                s = max(0.01, min(0.99, s))
             except Exception:
                 s = 0.5
 
+        best_score, best_variant = -1.0, None
         for row in metrics:
-            ac = row["alpha_click"]; bc = row["beta_click"]
-            av = row["alpha_conv"]; bv =row["beta_conv"]
+            ac, bc = float(row["alpha_click"]), float(row["beta_click"])
+            av, bv = float(row["alpha_conv"]), float(row["beta_conv"])
 
+            # 季節性でクリック側の事前分布を微調整
             ac_eff = ac + self.seasonality_boost * s
-            bc_eff = bc + self.seasonality_boost * (1-s)
-            av_eff = av
-            bv_eff = bv
+            bc_eff = bc + self.seasonality_boost * (1.0 - s)
 
-            pc = random.betavariate(ac_eff, bc_eff)
-            pv = random.betavariate(av_eff, bv_eff)
+            pc = random.betavariate(max(1e-6, ac_eff), max(1e-6, bc_eff))
+            pv = random.betavariate(max(1e-6, av),     max(1e-6, bv))
             score = pc * pv * vpc
             if score > best_score:
-                best_score = score
-                best_variant = row["variant_id"]
-
+                best_score, best_variant = score, row["variant_id"]
         return best_variant, best_score
-    
+
     @staticmethod
     def update_with_event(campaign_id: str, variant_id: str, event_type: str):
         if event_type == "impression":
-            storage.update_metric(campaign_id, variant_id, "impressions",1)
+            storage.update_metric(campaign_id, variant_id, "impressions", 1)
             storage.update_metric(campaign_id, variant_id, "beta_click", 1)
         elif event_type == "click":
             storage.update_metric(campaign_id, variant_id, "clicks", 1)
@@ -59,4 +61,4 @@ class ThompsonBandit:
             storage.update_metric(campaign_id, variant_id, "conversions", 1)
             storage.update_metric(campaign_id, variant_id, "alpha_conv", 1)
         else:
-            raise ValueError("unknown event_type")
+            raise ValueError("unknown event_type")