ref: remove cost_map from config

src/bespokelabs/curator/cost.py

@@ -12,15 +12,20 @@
 class _LitellmCostProcessor:
     def __init__(self, config, batch=False) -> None:
         self.batch = batch
-        self.cost_per_token_map = config.cost_per_token_map
-        if self.cost_per_token_map:
+        if config.in_mtok_cost is not None:
+            cost_per_input_token = config.in_mtok_cost / 1e6
+            if config.out_mtok_cost is not None:
+                cost_per_output_token = config.out_mtok_cost / 1e6
+            else:
+                cost_per_output_token = cost_per_input_token
+
             litellm.register_model(
                 {
                     config.model: {
-                        "max_tokens": self.cost_per_token_map.max_tokens,
-                        "input_cost_per_token": self.cost_per_token_map.input,
-                        "output_cost_per_token": self.cost_per_token_map.output,
-                        "litellm_provider": self.cost_per_token_map.provider,
+                        "max_tokens": 8192,
+                        "input_cost_per_token": cost_per_input_token,
+                        "output_cost_per_token": cost_per_output_token,
+                        "litellm_provider": "openai",
                     }
                 }
             )

src/bespokelabs/curator/request_processor/config.py

@@ -6,15 +6,6 @@
 logger = logging.getLogger(__name__)
 
 
-class _CostMap(BaseModel):
-    """BaseModel for cost map."""
-
-    input: int
-    output: int
-    provider: str = "openai"
-    max_tokens: int = 8192
-
-
 class RequestProcessorConfig(BaseModel):
     """Configuration for request processors.
 
@@ -28,7 +19,8 @@ class RequestProcessorConfig(BaseModel):
         require_all_responses: Whether to require successful responses for all requests
         generation_params: Dictionary of model-specific generation parameters
         api_key: Optional API key for authentication
-        cost_per_token_map: Optional cost map for input and output tokens
+        in_mtok_cost: Optional cost per million input tokens
+        out_mtok_cost: Optional cost per million output tokens
     """
 
     model: str
@@ -39,7 +31,8 @@ class RequestProcessorConfig(BaseModel):
     generation_params: dict = Field(default_factory=dict)
     return_completions_object: bool = False
     api_key: str | None = None
-    cost_per_token_map: _CostMap | None = None
+    in_mtok_cost: int | None = None
+    out_mtok_cost: int | None = None
 
     class Config:
         """BaseModel Setup class."""