From 7841caaedd1186e5da98722e49410d10b4bdafcd Mon Sep 17 00:00:00 2001
From: Nicolas van Kempen <nvankempen@meta.com>
Date: Wed, 4 Oct 2023 17:54:46 +0100
Subject: [PATCH] Print cost estimation on text completion

---
 pyproject.toml   |  3 +--
 src/cwhy/cwhy.py | 10 ++++++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7312aa2..8b53425 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [
   { name="Nicolas van Kempen", email="nvankemp@gmail.com" },
   { name="Bryce Adelstein Lelbach", email="brycelelbach@gmail.com" }
 ]
-dependencies = ["openai==0.28.1", "llm_utils==0.1.4"]
+dependencies = ["openai==0.28.1", "llm_utils==0.2"]
 description = "Explains and proposes fixes for compile-time errors for many programming languages."
 readme = "README.md"
 requires-python = ">=3.7"
@@ -23,7 +23,6 @@ classifiers = [
 [project.scripts]
 cwhy = "cwhy.__main__:main"
 
-
 [project.urls]
 "Homepage" = "https://github.com/plasma-umass/cwhy"
 "Bug Tracker" = "https://github.com/plasma-umass/cwhy/issues"
diff --git a/src/cwhy/cwhy.py b/src/cwhy/cwhy.py
index c60cebb..79a494e 100755
--- a/src/cwhy/cwhy.py
+++ b/src/cwhy/cwhy.py
@@ -103,8 +103,18 @@ def evaluate(args, stdin):
 def evaluate_text_prompt(args, prompt, wrap=True, **kwargs):
     completion = complete(args, prompt, **kwargs)
     text = completion.choices[0].message.content
+
     if wrap:
         text = llm_utils.word_wrap_except_code_blocks(text)
+
+    cost = llm_utils.calculate_cost(
+        completion.usage.prompt_tokens,
+        completion.usage.completion_tokens,
+        args["llm"]
+    )
+    text += "\n\n"
+    text += f"This request cost ~ {cost:.2f}$."
+
     return text