From 62258bcd862f3ac830198465bc950bbebc50c0b2 Mon Sep 17 00:00:00 2001
From: signebedi <signe@siftingwinnowing.com>
Date: Thu, 6 Apr 2023 16:20:25 -0500
Subject: [PATCH] Modified: error handling and improved API usage query (#62)

---
 gptty/__main__.py |  2 +-
 gptty/gptty.py    | 23 ++++++++++++++---------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/gptty/__main__.py b/gptty/__main__.py
index 375924d..4438276 100644
--- a/gptty/__main__.py
+++ b/gptty/__main__.py
@@ -110,7 +110,7 @@ async def chat_async_wrapper(config_path:str, verbose:bool):
   """
 
   # Print the text in cyan
-  click.echo(f"{CYAN}{title}\nWelcome to gptty (v.{__version__}), a ChatGPT wrapper in your TTY.\nType :help in the chat interface if you need help getting started.{' Verbose / Debug mode is on.' if verbose else ''}{RESET}\n")
+  click.echo(f"{CYAN}{title}\nWelcome to gptty (v.{__version__}), a ChatGPT wrapper in your TTY. Type :help in the chat interface if you need help getting started.{' Verbose / Debug mode is on. Query prompts will be preceded by your daily API usage in the format (query count, query tokens, response tokens).' if verbose else ''}{RESET}\n")
   
   if not os.path.exists(config_path):
       click.echo(f"{RED}FAILED to access app config file at {config_path}. Are you sure this is a valid config file? Run `gptty chat --help` for more information. You can get a sample config at <https://github.com/signebedi/gptty/blob/master/assets/gptty.ini.example>.")
diff --git a/gptty/gptty.py b/gptty/gptty.py
index 1648421..0bfad0b 100644
--- a/gptty/gptty.py
+++ b/gptty/gptty.py
@@ -49,13 +49,19 @@
 
 def usage_stats_today():
 
-    r = openai.api_requestor.APIRequestor()
-    resp = r.request("GET", f'/usage?date={datetime.now().strftime("%Y-%m-%d")}')
-    resp_object = resp[0].data
+    try:
+        r = openai.api_requestor.APIRequestor()
+        resp = r.request("GET", f'/usage?date={datetime.now().strftime("%Y-%m-%d")}')
+        resp_object = resp[0].data
+    except:
+        return None
+    # requests_today = resp_object['data'][0]['n_requests'] # num requests
+    # query_tokens_today = resp_object['data'][0]['n_context_tokens_total'] # query tokens
+    # response_tokens_today = resp_object['data'][0]['n_generated_tokens_total'] # response tokens
 
-    requests_today = resp_object['data'][0]['n_requests'] # num requests
-    query_tokens_today = resp_object['data'][0]['n_context_tokens_total'] # query tokens
-    response_tokens_today = resp_object['data'][0]['n_generated_tokens_total'] # response tokens
+    requests_today = sum(item["n_requests"] for item in resp_object['data'])
+    query_tokens_today = sum(item["n_context_tokens_total"] for item in resp_object['data'])
+    response_tokens_today = sum(item["n_generated_tokens_total"] for item in resp_object['data'])
 
     return requests_today, query_tokens_today, response_tokens_today
 
@@ -235,11 +241,10 @@ async def create_chat_room(configs=get_config_data(), log_responses:bool=True, c
         # Get user input
         try:
 
-            usage = usage_stats_today() if verbose else ""
 
             with patch_stdout():
-                i = await session.prompt_async(ANSI(f"{CYAN}{usage}> "), style=Style.from_dict({'': 'ansicyan'}))
-            print(f"{ERASE_LINE}{MOVE_CURSOR_UP}{GREY}{usage}> {i}\n", end="")
+                i = await session.prompt_async(ANSI(f"{CYAN}{usage_stats_today() if verbose else ''}> "), style=Style.from_dict({'': 'ansicyan'}))
+            print(f"{ERASE_LINE}{MOVE_CURSOR_UP}{GREY}{usage_stats_today() if verbose else ''}> {i}\n", end="")
 
             # i = await ainput(f"{CYAN}> ")
             tag,question = get_tag_from_text(i)