feat(budget_checker):refactor budget checker (#276)

* feat(budget_checker):refactor budget checker

* feat(budget_checker): fix pre commit

* feat(budget_checker): refactor budget checker as main class

Author: Zijie-Lei

tiny_scientist/__init__.py

@@ -1,9 +1,9 @@
+from .budget_checker import BudgetChecker
 from .coder import Coder
 from .reviewer import Reviewer
 from .safety_checker import SafetyChecker
 from .scientist import TinyScientist
 from .thinker import Thinker
-from .utils.checker import Checker
 from .writer import Writer
 
 __all__ = [
@@ -13,5 +13,5 @@
     "Thinker",
     "Writer",
     "TinyScientist",
-    "Checker",
+    "BudgetChecker",
 ]

tiny_scientist/utils/checker.py renamed to tiny_scientist/budget_checker.py

@@ -9,7 +9,7 @@ class BudgetExceededError(Exception):
     """Raised when a call would exceed the configured budget."""
 
 
-class Checker:
+class BudgetChecker:
     """Track API usage cost and enforce a spending budget."""
 
     def __init__(self, budget: Optional[float] = None) -> None:
@@ -48,7 +48,7 @@ def add_cost(
             if task_name not in self.per_task_cost:
                 self.per_task_cost[task_name] = 0.0
             self.per_task_cost[task_name] += cost
-        return cost
+        return float(cost)
 
     def report(self) -> None:
         print(f"Total cost: ${self.total_cost:.4f}")

tiny_scientist/coder.py

@@ -12,8 +12,8 @@
 from aider.models import Model
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .configs import Config
-from .utils.checker import Checker
 from .utils.llm import create_client, get_response_from_llm
 
 
@@ -28,7 +28,7 @@ def __init__(
         prompt_template_dir: Optional[str] = None,
         chat_history: Optional[str] = None,
         auto_install: bool = True,
-        cost_tracker: Optional[Checker] = None,
+        cost_tracker: Optional[BudgetChecker] = None,
     ):
         """Initialize the ExperimentCoder with configuration and Aider setup."""
         self.client, self.model = create_client(model)
@@ -38,7 +38,7 @@ def __init__(
         self.max_stderr_output = max_stderr_output
         self.auto_install = auto_install
         self.config = Config()
-        self.cost_tracker = cost_tracker or Checker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
 
         # Load prompts
         self.prompts = self.config.prompt_template.coder_prompt

tiny_scientist/reviewer.py

@@ -3,9 +3,9 @@
 
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .configs import Config
 from .tool import BaseTool, PaperSearchTool
-from .utils.checker import Checker
 from .utils.error_handler import api_calling_error_exponential_backoff
 from .utils.input_formatter import InputFormatter
 from .utils.llm import (
@@ -24,7 +24,7 @@ def __init__(
         num_reflections: int = 2,
         temperature: float = 0.75,
         prompt_template_dir: Optional[str] = None,
-        cost_tracker: Optional[Checker] = None,
+        cost_tracker: Optional[BudgetChecker] = None,
         pre_reflection_threshold: float = 0.5,
         post_reflection_threshold: float = 0.8,
         s2_api_key: Optional[str] = None,
@@ -38,7 +38,7 @@ def __init__(
         self.searcher: BaseTool = PaperSearchTool(s2_api_key=s2_api_key)
         self._query_cache: Dict[str, List[Dict[str, Any]]] = {}
         self.last_related_works_string = ""
-        self.cost_tracker = cost_tracker or Checker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
         self.pre_reflection_threshold = pre_reflection_threshold
         self.post_reflection_threshold = post_reflection_threshold
 

tiny_scientist/safety_checker.py

@@ -5,8 +5,8 @@
 
 import yaml
 
+from .budget_checker import BudgetChecker
 from .configs import Config
-from .utils.cost_tracker import CostTracker
 from .utils.error_handler import api_calling_error_exponential_backoff
 from .utils.llm import create_client, get_response_from_llm
 
@@ -44,10 +44,10 @@ def is_rejection_response(response: str) -> bool:
 
 class PromptAttackDetector:
     def __init__(
-        self, model: str = "gpt-4o", cost_tracker: Optional[CostTracker] = None
+        self, model: str = "gpt-4o", cost_tracker: Optional[BudgetChecker] = None
     ) -> None:
         self.client, self.model = create_client(model)
-        self.cost_tracker = cost_tracker or CostTracker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
         self.config = Config()
         self.prompts = self.config.prompt_template.safety_prompt
 
@@ -148,10 +148,10 @@ class SafetyChecker:
     """
 
     def __init__(
-        self, model: str = "gpt-4o", cost_tracker: Optional[CostTracker] = None
+        self, model: str = "gpt-4o", cost_tracker: Optional[BudgetChecker] = None
     ) -> None:
         self.model = model
-        self.cost_tracker = cost_tracker or CostTracker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
         self.detector = PromptAttackDetector(
             model=model, cost_tracker=self.cost_tracker
         )
@@ -199,14 +199,16 @@ def check_safety(self, intent: str) -> Tuple[bool, Dict[str, Any]]:
         self.cost_tracker.report()
         return is_safe, safety_report
 
-    def _load_ethics_prompts(self) -> Dict[str, str]:
+    def _load_ethics_prompts(self) -> Dict[str, Any]:
         """Load ethics prompts from the YAML file."""
         prompt_path = os.path.join(
             os.path.dirname(__file__), "prompts", "safetychecker_prompt.yaml"
         )
         try:
             with open(prompt_path, "r", encoding="utf-8") as file:
-                return yaml.safe_load(file)
+                import typing
+
+                return typing.cast(Dict[str, Any], yaml.safe_load(file))
         except FileNotFoundError:
             print(f"Warning: Ethics prompts file not found at {prompt_path}")
             return {}
@@ -399,7 +401,9 @@ def comprehensive_safety_check(
         # Overall safety determination
         overall_safe = is_intent_safe
         if idea is not None:
-            overall_safe = overall_safe and result["idea_ethics"]["is_ethically_sound"]
+            overall_safe = overall_safe and bool(
+                result["idea_ethics"]["is_ethically_sound"]
+            )
 
         result["overall_safety"] = {
             "is_safe": overall_safe,

tiny_scientist/scientist.py

@@ -4,12 +4,11 @@
 import toml
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .coder import Coder
 from .reviewer import Reviewer
 from .safety_checker import SafetyChecker
 from .thinker import Thinker
-from .utils.checker import Checker
-from .utils.cost_tracker import CostTracker
 from .utils.input_formatter import InputFormatter
 from .writer import Writer
 
@@ -80,7 +79,7 @@ def __init__(
 
         self.safety_checker = (
             SafetyChecker(
-                model=model, cost_tracker=CostTracker(budget=per_module_budget)
+                model=model, cost_tracker=BudgetChecker(budget=per_module_budget)
             )
             if enable_safety_check
             else None
@@ -94,8 +93,9 @@ def __init__(
             iter_num=3,
             search_papers=True,
             generate_exp_plan=True,
+            enable_ethical_defense=False,
             enable_safety_check=enable_safety_check,
-            cost_tracker=Checker(budget=allocation.get("thinker")),
+            cost_tracker=BudgetChecker(budget=allocation.get("thinker")),
         )
 
         self.coder = Coder(
@@ -104,22 +104,22 @@ def __init__(
             prompt_template_dir=prompt_template_dir,
             max_iters=4,
             max_runs=3,
-            cost_tracker=Checker(budget=allocation.get("coder")),
+            cost_tracker=BudgetChecker(budget=allocation.get("coder")),
         )
 
         self.writer = Writer(
             model=model,
             output_dir=output_dir,
             prompt_template_dir=prompt_template_dir,
             template=template,
-            cost_tracker=Checker(budget=allocation.get("writer")),
+            cost_tracker=BudgetChecker(budget=allocation.get("writer")),
         )
 
         self.reviewer = Reviewer(
             model=model,
             prompt_template_dir=prompt_template_dir,
             tools=[],
-            cost_tracker=Checker(budget=allocation.get("reviewer")),
+            cost_tracker=BudgetChecker(budget=allocation.get("reviewer")),
         )
 
     def think(

tiny_scientist/thinker.py

@@ -4,10 +4,10 @@
 
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .configs import Config
 from .safety_checker import SafetyChecker
 from .tool import PaperSearchTool
-from .utils.checker import Checker
 from .utils.error_handler import api_calling_error_exponential_backoff
 from .utils.llm import (
     create_client,
@@ -27,10 +27,11 @@ def __init__(
         output_dir: str = "",
         temperature: float = 0.75,
         prompt_template_dir: Optional[str] = None,
-        cost_tracker: Optional[Checker] = None,
+        cost_tracker: Optional[BudgetChecker] = None,
         enable_safety_check: bool = False,
         pre_reflection_threshold: float = 0.5,
         post_reflection_threshold: float = 0.8,
+        enable_ethical_defense: bool = False,
     ):
         self.tools = tools
         self.iter_num = iter_num
@@ -70,12 +71,13 @@ def __init__(
         3. Novelty: How original is the idea compared to existing work?
         4. Feasibility: How practical is implementation within reasonable resource constraints?
         5. Impact: What is the potential impact of this research on the field and broader applications?"""
-        self.cost_tracker = cost_tracker or Checker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
         self.pre_reflection_threshold = pre_reflection_threshold
         self.post_reflection_threshold = post_reflection_threshold
 
         self.enable_safety_check = enable_safety_check
-        # Initialize SafetyChecker for comprehensive safety checks
+        self.enable_ethical_defense = enable_ethical_defense
+        self.safety_checker: Optional[SafetyChecker]
         if self.enable_safety_check:
             self.safety_checker = SafetyChecker(
                 model=self.model, cost_tracker=self.cost_tracker

tiny_scientist/tool.py

@@ -11,8 +11,8 @@
 import toml
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .configs import Config
-from .utils.checker import Checker
 from .utils.error_handler import api_calling_error_exponential_backoff
 from .utils.llm import create_client, get_response_from_llm
 
@@ -22,8 +22,8 @@
 
 
 class BaseTool(abc.ABC):
-    def __init__(self, cost_tracker: Optional[Checker] = None) -> None:
-        self.cost_tracker = cost_tracker or Checker()
+    def __init__(self, cost_tracker: Optional[BudgetChecker] = None) -> None:
+        self.cost_tracker = cost_tracker or BudgetChecker()
         self.github_token = config["core"].get("github_token", None)
 
     @abc.abstractmethod

tiny_scientist/utils/llm.py

@@ -9,7 +9,7 @@
 import toml
 from google.generativeai.types import GenerationConfig
 
-from tiny_scientist.utils.checker import Checker
+from tiny_scientist.budget_checker import BudgetChecker
 
 # Load config
 config_path = os.path.join(
@@ -84,7 +84,7 @@ def get_batch_responses_from_llm(
     msg_history: Any = None,
     temperature: float = 0.75,
     n_responses: int = 1,
-    cost_tracker: Optional[Checker] = None,
+    cost_tracker: Optional[BudgetChecker] = None,
     task_name: Optional[str] = None,
 ) -> Tuple[List[str], List[List[Dict[str, str]]]]:
     if msg_history is None:
@@ -214,7 +214,7 @@ def get_response_from_llm(
     print_debug: bool = False,
     msg_history: Any = None,
     temperature: float = 0.75,
-    cost_tracker: Optional[Checker] = None,
+    cost_tracker: Optional[BudgetChecker] = None,
     task_name: Optional[str] = None,
 ) -> Tuple[str, List[Dict[str, Any]]]:
     if msg_history is None:
@@ -442,7 +442,7 @@ def get_batch_responses_from_llm_with_tools(
     msg_history: Optional[List[Dict[str, str]]] = None,
     temperature: float = 0.75,
     n_responses: int = 1,
-    cost_tracker: Optional[Checker] = None,
+    cost_tracker: Optional[BudgetChecker] = None,
     task_name: Optional[str] = None,
 ) -> Tuple[List[Union[str, Dict[str, Any]]], List[List[Dict[str, str]]]]:
     """

tiny_scientist/writer.py

@@ -10,9 +10,9 @@
 import cairosvg
 from rich import print
 
+from .budget_checker import BudgetChecker
 from .configs import Config
 from .tool import BaseTool, DrawerTool, PaperSearchTool
-from .utils.checker import Checker
 from .utils.llm import (
     create_client,
     extract_json_between_markers,
@@ -33,7 +33,7 @@ def __init__(
         template: str,
         temperature: float = 0.75,
         prompt_template_dir: Optional[str] = None,
-        cost_tracker: Optional[Checker] = None,
+        cost_tracker: Optional[BudgetChecker] = None,
         s2_api_key: Optional[str] = None,
     ) -> None:
         self.client, self.model = create_client(model)
@@ -50,7 +50,7 @@ def __init__(
             self.formatter = ICLROutputFormatter(model=self.model, client=self.client)
 
         self.prompts = self.config.prompt_template.writer_prompt
-        self.cost_tracker = cost_tracker or Checker()
+        self.cost_tracker = cost_tracker or BudgetChecker()
 
         with resources.files("tiny_scientist.fewshot_sample").joinpath(
             "automated_relational.txt"