From f38efe7b80043f4e28457e6f602de70e3394d0e2 Mon Sep 17 00:00:00 2001
From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com>
Date: Tue, 20 Feb 2024 22:00:28 +0530
Subject: [PATCH] Move num_outputs pricing logic to base.py

---
 daras_ai_v2/base.py       | 2 +-
 recipes/DocSummary.py     | 3 ---
 recipes/GoogleGPT.py      | 3 ---
 recipes/GoogleImageGen.py | 3 ---
 recipes/SmartGPT.py       | 3 ---
 recipes/Text2Audio.py     | 3 ---
 6 files changed, 1 insertion(+), 16 deletions(-)

diff --git a/daras_ai_v2/base.py b/daras_ai_v2/base.py
index 0fc63ea33..f5c5a182c 100644
--- a/daras_ai_v2/base.py
+++ b/daras_ai_v2/base.py
@@ -1838,7 +1838,7 @@ def get_price_roundoff(self, state: dict) -> int:
         return max(1, math.ceil(self.get_raw_price(state)))
 
     def get_raw_price(self, state: dict) -> float:
-        return self.price
+        return self.price * state.get("num_outputs", 1)
 
     def get_example_response_body(
         self,
diff --git a/recipes/DocSummary.py b/recipes/DocSummary.py
index fa4dbd427..476bd5da6 100644
--- a/recipes/DocSummary.py
+++ b/recipes/DocSummary.py
@@ -181,9 +181,6 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
             case _:
                 raise NotImplementedError(f"{chain_type} not implemented")
 
-    def get_raw_price(self, state: dict) -> float:
-        return self.price * state.get("num_outputs", 1)
-
 
 MAP_REDUCE_PROMPT = """
 {documents}
diff --git a/recipes/GoogleGPT.py b/recipes/GoogleGPT.py
index 48c512560..35d650acb 100644
--- a/recipes/GoogleGPT.py
+++ b/recipes/GoogleGPT.py
@@ -279,6 +279,3 @@ def run_v2(
             max_tokens=request.max_tokens,
             avoid_repetition=request.avoid_repetition,
         )
-
-    def get_raw_price(self, state: dict) -> float:
-        return self.price * state.get("num_outputs", 1)
diff --git a/recipes/GoogleImageGen.py b/recipes/GoogleImageGen.py
index a935e54cb..278128a37 100644
--- a/recipes/GoogleImageGen.py
+++ b/recipes/GoogleImageGen.py
@@ -238,6 +238,3 @@ def render_example(self, state: dict):
 
     def preview_description(self, state: dict) -> str:
         return "Enter a Google Image Search query + your Img2Img text prompt describing how to alter the result to create a unique, relevant ai generated images for any search query."
-
-    def get_raw_price(self, state: dict) -> float:
-        return super().get_raw_price(state) * state.get("num_outputs", 1)
diff --git a/recipes/SmartGPT.py b/recipes/SmartGPT.py
index 10c4a3fe7..0aabcfefb 100644
--- a/recipes/SmartGPT.py
+++ b/recipes/SmartGPT.py
@@ -204,9 +204,6 @@ def render_steps(self):
     def preview_description(self, state: dict) -> str:
         return "SmartGPT is a cutting-edge AI technology that can be used to generate natural language responses to any given input. We have combined the power of [CoT](https://arxiv.org/abs/2305.02897), [Reflexion](https://arxiv.org/abs/2303.11366) & [DERA](https://arxiv.org/abs/2303.17071) into one pipeline so that you can use ChatGPT to its full potential! Input your prompt + a reflection/research prompt + a resolver prompt to use SmartGPT for enhanced text generation, natural language and incredible question-answer results."
 
-    def get_raw_price(self, state: dict) -> float:
-        return self.price * state.get("num_outputs", 1)
-
 
 def answers_as_prompt(texts: list[str], sep="\n\n") -> str:
     return sep.join(
diff --git a/recipes/Text2Audio.py b/recipes/Text2Audio.py
index 5140fd2c2..f7ddf5aab 100644
--- a/recipes/Text2Audio.py
+++ b/recipes/Text2Audio.py
@@ -140,9 +140,6 @@ def render_example(self, state: dict):
     def preview_description(self, state: dict) -> str:
         return "Generate AI Music with text instruction prompts. AudiLDM is capable of generating realistic audio samples by process any text input. Learn more [here](https://huggingface.co/cvssp/audioldm-m-full)."
 
-    def get_raw_price(self, state: dict) -> float:
-        return super().get_raw_price(state) * state.get("num_outputs", 1)
-
 
 def _render_output(state):
     selected_models = state.get("selected_models", [])