Use result for why

src/chatdbg/chatdbg_utils.py

@@ -1,7 +1,7 @@
 import argparse
 import textwrap
 import time
-from typing import Any, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import llm_utils
 import openai
@@ -74,7 +74,13 @@ def parse_known_args(argv: List[str]) -> Tuple[argparse.Namespace, List[str]]:
 
 
 def explain(
-    source_code: str, traceback: str, exception: str, args: argparse.Namespace
+    source_code: str,
+    traceback: str,
+    exception: str,
+    args: argparse.Namespace,
+    append_message: Callable[[str], None] = print,
+    append_warning: Callable[[str], None] = print,
+    set_error: Callable[[str], None] = print,
 ) -> None:
     user_prompt = f"""
 Explain what the root cause of this error is, given the following source code
@@ -96,32 +102,31 @@ def explain(
     input_tokens = llm_utils.count_tokens(args.llm, user_prompt)
 
     if args.debug:
-        print(user_prompt)
-        print(f"Total input tokens: {input_tokens}")
+        append_message(f"{user_prompt}\n\nTotal input tokens: {input_tokens}.")
         return
 
     start = time.time()
 
     try:
         client = openai.OpenAI(timeout=args.timeout)
     except openai.OpenAIError:
-        print("You need an OpenAI key to use this tool.")
-        print("You can get a key here: https://platform.openai.com/api-keys")
-        print("Set the environment variable OPENAI_API_KEY to your key value.")
+        append_warning("you need an OpenAI key to use this tool.")
+        append_warning("You can get a key here: https://platform.openai.com/api-keys.")
+        append_warning("set the environment variable OPENAI_API_KEY to your key value.")
         return
 
     try:
         completion = client.chat.completions.create(
             model=args.llm, messages=[{"role": "user", "content": user_prompt}]
         )
     except openai.NotFoundError:
-        print(f"'{args.llm}' either does not exist or you do not have access to it.")
+        set_error(f"'{args.llm}' does not exist or you do not have access to it.")
         return
     except openai.RateLimitError:
-        print("You have exceeded a rate limit or have no remaining funds.")
+        set_error("you have exceeded a rate limit or have no remaining funds.")
         return
     except openai.APITimeoutError:
-        print("The OpenAI API timed out.")
+        set_error("the OpenAI API timed out.")
         return
 
     text = completion.choices[0].message.content
@@ -130,7 +135,10 @@ def explain(
     output_tokens = completion.usage.completion_tokens
     cost = llm_utils.calculate_cost(input_tokens, output_tokens, args.llm)
 
-    print(llm_utils.word_wrap_except_code_blocks(text))
-    print()
-    print(f"Elapsed time: {elapsed:.2f} seconds")
-    print(f"Total cost: {cost:.2f}$")
+    append_message(
+        llm_utils.word_wrap_except_code_blocks(text)
+        + "\n\n"
+        + f"Elapsed time: {elapsed:.2f} seconds"
+        + "\n"
+        + f"Total cost: {cost:.2f}$"
+    )