add ability to select markdown theme

Author: npiv

README.md

@@ -227,8 +227,8 @@ chatblade can be used with an Azure OpenAI endpoint, in which case in addition t
 ### Help
 
 ```
-usage: Chatblade [-h] [--openai-api-key key] [--temperature t] [-c {3.5,4}] [-i] [-s] [-t] [-p name] [-e] [-r] [-n] [-o] [-l] [-S sess] [--session-list] [--session-path]
-                 [--session-dump] [--session-delete] [--session-rename newsess]
+usage: Chatblade [-h] [--openai-api-key key] [--temperature t] [-c {3.5,4}] [-i] [-s] [-t] [-p name] [-e] [-r] [-n] [-o] [--theme theme] [-l] [-S sess] [--session-list]
+                 [--session-path] [--session-dump] [--session-delete] [--session-rename newsess]
                  [query ...]
 
 a CLI Swiss Army Knife for ChatGPT
@@ -244,13 +244,14 @@ options:
   -i, --interactive               start an interactive chat session. This will implicitly continue the conversation
   -s, --stream                    Stream the incoming text to the terminal
   -t, --tokens                    display what *would* be sent, how many tokens, and estimated costs
-  -p name, --prompt-file name     prompt name - will load the prompt at ~/.config/chatblade/name as system msg
+  -p name, --prompt-file name     prompt name - will load the prompt with that name at ~/.config/chatblade/name or a path to a file
 
 result formatting options:
   -e, --extract                   extract content from response if possible (either json or code block)
   -r, --raw                       print session as pure text, don't pretty print or format
   -n, --no-format                 do not add pretty print formatting to output
   -o, --only                      Only display the response, omit query
+  --theme theme                   Set the theme for syntax highlighting see https://pygments.org/styles/, can also be set with CHATBLADE_THEME
 
 session options:
   -l, --last                      alias for '-S last', the default session if none is specified

chatblade/parser.py

@@ -19,6 +19,13 @@ def get_openai_key(options):
     else:
         return None
 
+def get_theme(options):
+    if options["theme"]:
+        return options["theme"]
+    elif "CHATBLADE_THEME" in os.environ:
+        return os.environ["CHATBLADE_THEME"]
+    else:
+        return None
 
 def extract_query(query):
     """The query comes from both the query and any piped input
@@ -40,6 +47,7 @@ def extract_query(query):
 def extract_options(options):
     options = vars(options)  # to map
     options["openai_api_key"] = get_openai_key(options)
+    options["theme"] = get_theme(options)
     options["model"] = {"3.5": "gpt-3.5-turbo", "4": "gpt-4"}[options["chat_gpt"]]
     del options["query"]
     del options["chat_gpt"]
@@ -137,6 +145,12 @@ def parse(args):
         help="Only display the response, omit query",
         action="store_true",
     )
+    display_opts.add_argument(
+        "--theme",
+        metavar="theme",
+        type=str,
+        help="Set the theme for syntax highlighting see https://pygments.org/styles/, can also be set with CHATBLADE_THEME",
+    )
 
     session_opts = parser.add_argument_group("session options")
     session_opts.add_argument(

chatblade/printer.py

@@ -62,7 +62,7 @@ def print_message(message, args):
     printable = message.content
     if not args.raw:
         printable = detect_and_format_message(
-            message.content, cutoff=1000 if message.role == "user" else None
+            message.content, cutoff=1000 if message.role == "user" else None, theme=args.theme
         )
     if not args.no_format:
       console.print(Rule(message.role, style=COLORS[message.role]))
@@ -86,7 +86,7 @@ def extract_messages(messages, args):
         print(message.content.strip())
 
 
-def detect_and_format_message(msg, cutoff=None):
+def detect_and_format_message(msg, cutoff=None, theme=None):
     if cutoff and len(msg) > cutoff:
         msg = "... **text shortened** ... " + msg[-cutoff:]
         return msg
@@ -95,7 +95,8 @@ def detect_and_format_message(msg, cutoff=None):
         return JSON(extract_json(msg))
     elif looks_like_markdown(msg):
         utils.debug(detected="markdown")
-        return Markdown(msg)
+        theme = "monokai" if theme is None else theme
+        return Markdown(msg,code_theme=theme)
     else:
         utils.debug(detected="regular")
         return msg

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = chatblade
-version = 0.2.3
+version = 0.3.0
 description = CLI Swiss Army Knife for ChatGPT
 long_description = file: README.md
 long_description_content_type=text/markdown