Add function for running chat query with OpenAI

src/zcmds/cmds/common/askai.py

@@ -111,6 +111,51 @@ def parse_args() -> argparse.Namespace:
     return argparser.parse_args()
 
 
+def run_chat_query(
+    chatbot: ChatBot,
+    prompts: list[str],
+    output_stream: OutStream,
+    args: argparse.Namespace,
+) -> Optional[int]:
+    # allow exit() and exit to exit the app
+    as_json = args.json
+    if not as_json:
+        print("############ OPEN-AI QUERY")
+    try:
+        chat_stream: ChatStream = chatbot.query(prompts, no_stream=args.no_stream)
+    except ChatGPTConnectionError as err:
+        print(err)
+        return 1
+    except ChatGPTAuthenticationError as e:
+        print(
+            "Error authenticating with OpenAI, deleting password from config and exiting."
+        )
+        print(e)
+        save_config({})
+        return 1
+    except ChatGPTRateLimitError:
+        print("Rate limit exceeded, set a new key with --set-key")
+        return 1
+    if as_json:
+        print(chat_stream)
+        return 0
+    if chat_stream is None or not chat_stream.success():
+        print("No error response received from OpenAI, response was:")
+        output_stream.write(str(chat_stream.response()))
+        return 1
+    if not args.output:
+        print("############ OPEN-AI RESPONSE\n")
+    response_text = ""
+    for text in chat_stream:
+        if text is None:
+            break
+        response_text += text
+        output_stream.write(response_text)
+    output_stream.write(response_text + "\n")
+    prompts.append(response_text)
+    return None
+
+
 def cli() -> int:
     args = parse_args()
 
@@ -152,7 +197,6 @@ def cli() -> int:
             "\nInteractive mode - press return three times to submit your code to OpenAI"
         )
     prompt = args.prompt or prompt_input()
-    as_json = args.json
 
     def log(*pargs, **kwargs):
         if not args.verbose:
@@ -175,66 +219,28 @@ def log(*pargs, **kwargs):
         ai_assistant_prompt=ai_assistant_prompt,
     )
 
-    def run_chat_query(prompts: list[str], output_stream: OutStream) -> Optional[int]:
-        # allow exit() and exit to exit the app
-        new_cmd = prompts[-1].strip().replace("()", "")
-        if new_cmd.startswith("!"):
-            prompts = prompts[0:-1]
-            new_cmd = new_cmd[1:]
-            rtn = os.system(new_cmd)
-            print(f"Command exited and returned {rtn}")
-            prompts.append(prompt_input())
-            return None
-        if new_cmd == "exit":
-            print("Exited due to 'exit' command")
-            return 0
-        if not as_json:
-            print("############ OPEN-AI QUERY")
-        try:
-            chat_stream: ChatStream = chatbot.query(prompts, no_stream=args.no_stream)
-        except ChatGPTConnectionError as err:
-            print(err)
-            return 1
-        except ChatGPTAuthenticationError as e:
-            print(
-                "Error authenticating with OpenAI, deleting password from config and exiting."
-            )
-            print(e)
-            save_config({})
-            return 1
-        except ChatGPTRateLimitError:
-            print("Rate limit exceeded, set a new key with --set-key")
-            return 1
-        if as_json:
-            print(chat_stream)
-            return 0
-        if chat_stream is None or not chat_stream.success():
-            print("No error response received from OpenAI, response was:")
-            output_stream.write(str(chat_stream.response()))
-            return 1
-        if not args.output:
-            print("############ OPEN-AI RESPONSE\n")
-        response_text = ""
-        for text in chat_stream:
-            if text is None:
-                break
-            response_text += text
-            output_stream.write(response_text)
-        output_stream.write(response_text + "\n")
-        prompts.append(response_text)
-        if not interactive:
-            return None
-        prompts.append(prompt_input())
-        return None
-
     while True:
         try:
-            output_stream = OutStream(args.output)
-            rtn = run_chat_query(prompts, output_stream)
+            rtn: Optional[int] = None
+            output_stream = OutStream(args.output)  # needs to be created every time.
+            new_cmd = prompts[-1].strip().replace("()", "")
+            if new_cmd.startswith("!"):
+                prompts = prompts[0:-1]
+                new_cmd = new_cmd[1:]
+                rtn = os.system(new_cmd)
+                print(f"Command exited and returned {rtn}")
+                prompts.append(prompt_input())
+                continue
+            elif new_cmd == "exit":
+                print("Exited due to 'exit' command")
+                return 0
+            rtn = run_chat_query(chatbot, prompts, output_stream, args)
             if not interactive:
                 return rtn or 0
             if rtn is not None:
                 return rtn
+            # next loop.
+            prompts.append(prompt_input())
         finally:
             output_stream.close()