From 70ffe470aa96440e07c5f411a74b7a947b7f3ee3 Mon Sep 17 00:00:00 2001
From: Tim Asp <707699+timothyasp@users.noreply.github.com>
Date: Thu, 13 Apr 2023 11:28:42 -0700
Subject: [PATCH] Add easy print method to openai callback (#2848)

Found myself constantly copying the snippet outputting all the callback
tracking details. so adding a simple way to output the full context
---
 .../llms/examples/token_usage_tracking.ipynb       | 14 +++++---------
 langchain/callbacks/openai_info.py                 |  9 +++++++++
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/docs/modules/models/llms/examples/token_usage_tracking.ipynb b/docs/modules/models/llms/examples/token_usage_tracking.ipynb
index 1f8de942..cbb0c9bd 100644
--- a/docs/modules/models/llms/examples/token_usage_tracking.ipynb
+++ b/docs/modules/models/llms/examples/token_usage_tracking.ipynb
@@ -43,22 +43,18 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Total Tokens: 39\n",
-      "Prompt Tokens: 4\n",
-      "Completion Tokens: 35\n",
+      "Tokens Used: 42\n",
+      "\tPrompt Tokens: 4\n",
+      "\tCompletion Tokens: 38\n",
       "Successful Requests: 1\n",
-      "Total Cost (USD): $0.0007800000000000001\n"
+      "Total Cost (USD): $0.00084\n"
      ]
     }
    ],
    "source": [
     "with get_openai_callback() as cb:\n",
     "    result = llm(\"Tell me a joke\")\n",
-    "    print(f\"Total Tokens: {cb.total_tokens}\")\n",
-    "    print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
-    "    print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
-    "    print(f\"Successful Requests: {cb.successful_requests}\")\n",
-    "    print(f\"Total Cost (USD): ${cb.total_cost}\")"
+    "    print(cb)"
    ]
   },
   {
diff --git a/langchain/callbacks/openai_info.py b/langchain/callbacks/openai_info.py
index 299a3c8f..42005acb 100644
--- a/langchain/callbacks/openai_info.py
+++ b/langchain/callbacks/openai_info.py
@@ -53,6 +53,15 @@ class OpenAICallbackHandler(BaseCallbackHandler):
     successful_requests: int = 0
     total_cost: float = 0.0
 
+    def __repr__(self) -> str:
+        return (
+            f"Tokens Used: {self.total_tokens}\n"
+            f"\tPrompt Tokens: {self.prompt_tokens}\n"
+            f"\tCompletion Tokens: {self.completion_tokens}\n"
+            f"Successful Requests: {self.successful_requests}\n"
+            f"Total Cost (USD): ${self.total_cost}"
+        )
+
     @property
     def always_verbose(self) -> bool:
         """Whether to call verbose callbacks even if verbose is False."""