Updated to version 0.5

README.md

@@ -2,6 +2,28 @@
 
 ![Animated GIF](https://github.com/wunderwuzzi23/blog/raw/master/static/images/2023/yolo-shell-anim-gif.gif)
 
+# Update Yolo v0.5 - Support for Claude and other providers
+
+* Added Claude support. Can an API key from Anthropic, current model `claude-3-5-sonnet-20240620`. 
+* ai_model.py to abstract model usage and allow adding new providers more easily
+* Rewrote some logic to simplify and generalize support for various new APIs (like Ollama, Claude)
+
+# Update Yolo v0.4 - Support for Groq
+
+* Added groq support. You can get an API key at `https://console.groq.com` and set mode to for instance `llama3-8b-8192`. groq is lightning fast. 
+* Simplified and improved default `prompt.txt`, 
+* Note: Testing shows that model `gpt-4o` gives the best results.
+
+
+# Update Yolo v0.3 - Support for Azure OpenAI
+
+* Key changes are upgrades to the latest OpenAI libraries and support for Azure OpenAI. There is an `api` key in the `yolo.yaml` that can be set to `azure_openai` and then you can provide all the parameters accordingly in the yaml file as well (`api-version`, your `azure-endpoint`,...). The api key for azure is called `AZURE_OPENAI_API_KEY` by the way. It can be set via environment variable and config file.
+* It's now possible to change the color of the suggested command via config file
+* The "modify prompt" feature is now optional and can be toggled via config file.
+* Minor bug fixes (like copy to clipboard should work on macOS)
+
+Tested on macOS and Linux. Windows hopefully still works also.
+
 # Update Yolo v0.2 - Support for GPT-4 API
 
 This update introduces the `yolo.yaml` configuration file. In this file you can specify which OpenAI model you want to query, and other settings. The safety switch also moved into this configuration file.
@@ -9,16 +31,18 @@ This update introduces the `yolo.yaml` configuration file. In this file you can
 For now the default model is still `gpt-3.5-turbo`, but you can update to `gpt-4` if you have gotten access already!
 
 ```
-Yolo v0.2 - by @wunderwuzzi23
+Yolo v0.3 - by @wunderwuzzi23
 
 Usage: yolo [-a] list the current directory information
 Argument: -a: Prompt the user before running the command (only useful when safety is off)
 
 Current configuration per yolo.yaml:
-* Model        : gpt-3.5-turbo
+* API          : openai
+* Model        : gpt-4-turbo-preview
 * Temperature  : 0
 * Max. Tokens  : 500
-* Safety       : on
+* Safety       : True
+* Command Color: blue
 ```
 
 Happy Hacking!
@@ -41,7 +65,18 @@ yolo show me some funny unicode characters
 There are three ways to configure the key on Linux and macOS:
 - You can either `export OPENAI_API_KEY=<yourkey>`, or have a `.env` file in the same directory as `yolo.py` with `OPENAI_API_KEY="<yourkey>"` as a line
 - Create a file at `~/.openai.apikey` with the key in it
-- Add the key to the `yolo.yaml` configuration file
+- Set the key in the `yolo.yaml` configuration file
+
+### Azure OpenAI Key configuration
+There are three ways to configure the key on Linux and macOS:
+- You can either `export AZURE_OPENAI_API_KEY=<yourkey>`, or have a `.env` file in the same directory as `yolo.py` with `AZURE_OPENAI_API_KEY="<yourkey>"` as a line
+- Create a file at `~/.azureopenai.apikey` with the key in it
+- Set the key in the `yolo.yaml` configuration file
+
+### Groq Configuration
+- Grab an API key from `console.groq.com` 
+- You can either `export GROQ_API_KEY=<yourkey>`, or have a `.env` file in the same directory as `yolo.py` with `GROQ_API_KEY="<yourkey>"` as a line
+- Set `api` and `model` (e.g llama3-8b-8192) in `yolo.yaml` configuration file
 
 ## Aliases
 
@@ -74,6 +109,22 @@ You also have the option to:
 
 That's it basically.
 
+
+### PowerShell Command 
+
+If you want a PowerShell style command you can use something like this (check that the path is to yolo.py is correct) and you can add this function to your PowerShell profile. 
+
+```
+function Invoke-Yolo { 
+    param([Parameter(ValueFromRemainingArguments=$true)]$Commands)
+    $AllCommands = $Commands -join " "
+    $YoloPath = Join-Path -Path $env:USERPROFILE -ChildPath "\yolo-ai-cmdbot\yolo.py"
+    python.exe $YoloPath $AllCommands
+}
+```
+
+You can find the profile file with `$PROFILE` in a PowerShell Terminal. I might add a better installation in future.
+
 ## OpenAI API Key Configuration on Windows
 
 On Windows `export OPENAI_API_KEY=<yourkey>` will not work instead:
@@ -83,6 +134,8 @@ On Windows `export OPENAI_API_KEY=<yourkey>` will not work instead:
 
 Optionally (since v.0.2), the key can also be stored in `yolo.yaml`.
 
+If you want to use Azure, the the key is called `AZURE_OPENAI_API_KEY`.
+
 ## Running yolo on Windows 
 
 Windows is less tested, it does work though and will use PowerShell.

ai_model.py

@@ -0,0 +1,149 @@
+# MIT License
+# Copyright (c) 2023-2024 wunderwuzzi23
+# Greetings from Seattle! 
+
+from abc import ABC, abstractmethod
+from openai import OpenAI
+from groq import Groq
+from ollama import Client
+from openai import AzureOpenAI 
+from anthropic import Anthropic
+import os
+
+class AIModel(ABC):
+    @abstractmethod
+    def chat(self, model, messages):
+        pass
+
+    @abstractmethod
+    def moderate(self, message):
+        pass
+
+    @staticmethod
+    def get_model_client(config):
+        api_provider=config["api"]
+
+        if api_provider == "" or api_provider==None:
+            api_provider = "groq"
+        
+        if api_provider == "groq":
+            return GroqModel(api_key=os.environ.get("GROQ_API_KEY"))
+        
+        elif api_provider == "openai":
+            api_key = os.getenv("OPENAI_API_KEY")
+            if not api_key:  
+                api_key=config["openai_api_key"]       
+            if not api_key:  #If statement to avoid "invalid filepath" error
+                home_path = os.path.expanduser("~")   
+                api_key=open(os.path.join(home_path,".openai.apikey"), "r").readline().strip()
+                api_key = api_key
+
+            return OpenAIModel(api_key=api_key)
+        
+        elif api_provider == "azure":
+            api_key = os.getenv("AZURE_OPENAI_API_KEY")
+            if not api_key:  
+                api_key=config["azure_openai_api_key"]
+            if not api_key: 
+                home_path = os.path.expanduser("~")   
+                api_key=open(os.path.join(home_path,".azureopenai.apikey"), "r").readline().strip()
+
+            return AzureOpenAIModel(
+                    api_key=api_key,
+                    azure_endpoint=config["azure_endpoint"], 
+                    api_version=config["azure_api_version"])
+        
+        elif api_provider == "ollama":
+            ollama_api   = os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434")
+            #ollama_model = os.environ.get("OLLAMA_MODEL", "llama3-8b-8192")
+            return OllamaModel(ollama_api)
+
+        if api_provider == "anthropic":
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key: 
+                api_key=config["anthropic_api_key"]
+            return AnthropicModel(api_key=api_key) 
+        else:
+            raise ValueError(f"Invalid AI model provider: {api_provider}")
+
+class GroqModel(AIModel):
+    def __init__(self, api_key):
+        self.client = Groq(api_key=api_key)
+
+    def chat(self, messages, model, temperature, max_tokens):
+        resp = self.client.chat.completions.create(model=model, 
+                                                   messages=messages, 
+                                                   temperature=temperature, 
+                                                   max_tokens=max_tokens)
+        return resp.choices[0].message.content
+    
+    def moderate(self, message):
+        pass
+
+class OpenAIModel(AIModel):
+    def __init__(self, api_key):
+        self.client = OpenAI(api_key=api_key)
+
+    def chat(self, messages, model, temperature, max_tokens):
+        resp = self.client.chat.completions.create(model=model, 
+                                                   messages=messages, 
+                                                   temperature=temperature,
+                                                   max_tokens=max_tokens)
+        
+        return resp.choices[0].message.content
+    
+    def moderate(self, message):
+        return self.client.moderations.create(input=message)
+
+class OllamaModel(AIModel):
+    def __init__(self, host):
+        self.client = Client(host=host)
+    
+    def chat(self, messages, model, temperature, max_tokens):
+        resp = self.client.chat(model=model, 
+                                messages=messages)
+        return resp["message"]["content"]
+    
+    def moderate(self, message):
+        pass
+
+
+class AzureOpenAIModel(AIModel):
+    def __init__(self, azure_endpoint, api_key, api_version):
+        self.client = AzureOpenAI(azure_endpoint=azure_endpoint, api_key=api_key, api_version=api_version)
+
+    def chat(self, messages, model, temperature, max_tokens):
+
+        resp = self.client.chat.completions.create(model=model, 
+                        messages=messages, 
+                        temperature=temperature, 
+                        max_tokens=max_tokens)
+        
+        return resp.choices[0].message.content
+    
+    def moderate(self, message):
+        return self.client.moderations.create(input=message)
+
+class AnthropicModel(AIModel):
+    def __init__(self, api_key):
+        self.client = Anthropic(api_key=api_key)
+
+    def chat(self, messages, model, temperature, max_tokens):
+        ## Anthropic requires the system prompt to be passed separately
+        ## Hence extracting system prompt role from the messages
+        ## and then passing the messages without the system role 
+        ## messages is not subscriptable, so we need to convert it to a list
+        system_prompt = next((m.get("content", "") for m in messages if m.get("role") == "system"), "")
+
+        # Remove system messages from the list
+        user_messages = [m for m in messages if m.get("role") != "system"]
+        resp = self.client.messages.create(model=model, 
+                                    system=system_prompt,
+                                    messages=user_messages,
+                                    temperature=temperature, 
+                                    max_tokens=max_tokens)
+        
+        return resp.content[0].text
+    
+    def moderate(self, message):
+        pass

install.bat

@@ -5,6 +5,7 @@ setlocal enabledelayedexpansion
 if not exist %~dp0\yolo.py ( echo `yolo.py` missing in %~dp0 cannot install & goto :choice_default_3 )
 if not exist %~dp0\prompt.txt ( echo `prompt.txt` missing in %~dp0 cannot install & goto :choice_default_3 ) 
 if not exist %~dp0\yolo.yaml ( echo `yolo.yaml` missing in %~dp0 cannot install & goto :choice_default_3 ) 
+if not exist %~dp0\ai_model.py ( echo `ai_model.py` missing in %~dp0 cannot install & goto :choice_default_3 ) 
 
 
 :: Note: "~" or %HOME% is equivalent to "%HOMEDRIVE%%HOMEPATH%\" but the latter is set in VM environments (from what I can tell)
@@ -137,6 +138,7 @@ mkdir !TARGET_DIR!
 copy  %~dp0\yolo.py !TARGET_DIR!
 copy  %~dp0\prompt.txt !TARGET_DIR!
 copy  %~dp0\yolo.yaml !TARGET_DIR!
+copy  %~dp0\ai_model.py !TARGET_DIR!
 goto :EOF
 
 :: Create yolo.bat and input code linking to created directory
@@ -312,4 +314,6 @@ echo   -If you run PowerShell as administrator you can then run `setx OPENAI_API
 echo   -Go to `Start` and search `edit environment variables for your account` and manually create the variable with name `OPENAI_API_KEY` and value `[yourkey]`
 echo (4) Another option is to put the API key in the yolo.yaml configuration file (since v.0.2)
 echo.
+echo Yolo also supports Azure OpenAI, and many other LLMs now. Configure settings in yolo.yaml accordingly.
+echo.
 goto :EOF

install.sh

@@ -7,7 +7,7 @@ TARGET_FULLPATH=$TARGET_DIR/yolo.py
 mkdir -p $TARGET_DIR
 
 echo "- Copying files..."
-cp yolo.py prompt.txt yolo.yaml $TARGET_DIR
+cp yolo.py prompt.txt yolo.yaml ai_model.py $TARGET_DIR
 chmod +x $TARGET_FULLPATH
 
 # Creates two aliases for use
@@ -34,9 +34,11 @@ fi
 echo
 echo "Done."
 echo
-echo "Make sure you have the OpenAI API key set via one of these options:" 
+echo "Make sure you have your LLM key (e.g. OpenAI API) set via one of these options:" 
 echo "  - environment variable"
-echo "  - .env or an ~/.openai.apikey file or in"
+echo "  - .env or in"
 echo "  - yolo.yaml"
 echo
+echo "Yolo also supports Azure OpenAI, Ollama, groq, Claude now. Change settings in yolo.yaml accordingly."
+echo
 echo "Have fun!"

prompt.txt

@@ -1,31 +1,26 @@
-Act as a natural language to {shell} command translation engine on {os}.
+You are Yolo, a natural language to {shell} command translation engine for {os}. You are an expert in {shell} on {os} and translate the question at the end to valid command line syntax.
 
-You are an expert in {shell} on {os} and translate the question at the end to valid syntax.
+Rules:
+* No code style markdown output, ever.
+* Construct valid {shell} command to solve the question
+* Leverage help and man pages to ensure valid syntax and an optimal solution
+* Be concise, think step by step, and show just final commands in plain text
+* Only show a single answer, but you can always chain commands together 
+* Create valid syntax of {shell} on {os}, include comments if useful
+* If python or python3 is installed you can use it to solve problems
+* Even if there is a lack of details, find the most logical solution by going about it step by step
+* Do not return multiple solutions
+* Do not show html, styled, colored formatting 
+* Do not create invalid syntax or cause syntax errors
+* Do not add unnecessary text in the response 
+* Do not add notes or intro sentences 
+* Do not show multiple distinct solutions to the question
+* Do not add explanations on what the commands do
+* Do not return what the question was 
+* Do not repeat or paraphrase the question in your response 
+* Do not rush to a conclusion
+* Never start a response with ```
 
-Follow these rules:
-Construct valid {shell} command that solve the question
-Leverage help and man pages to ensure valid syntax and an optimal solution
-Be concise 
-Just show the commands 
-Return only plaintext 
-Only show a single answer, but you can always chain commands together 
-Think step by step
-Only create valid syntax (you can use comments if it makes sense)
-If python is installed you can use it to solve problems
-if python3 is installed you can use it to solve problems
-Even if there is a lack of details, attempt to find the most logical solution by going about it step by step
-Do not return multiple solutions
-Do not show html, styled, colored formatting 
-Do not creating invalid syntax 
-Do not add unnecessary text in the response 
-Do not add notes or intro sentences 
-Do not show multiple distinct solutions to the question
-Do not add explanations on what the commands do
-Do not return what the question was 
-Do not repeat or paraphrase the question in your response 
-Do not cause syntax errors
-Do not rush to a conclusion
-
-Follow all of the above rules. This is important you MUST follow the above rules. There are no exceptions to these rules. You must always follow them. No exceptions.
+Follow above rules. There are no exceptions to these rules. 
 
 Question: 

requirements.txt

@@ -1,7 +1,10 @@
-openai==0.27
-termcolor==2.2.0
-colorama==0.4.4
-python-dotenv==1.0.0
-distro==1.7.0
-PyYAML==5.4.1
-pyperclip==1.8.2
+ollama==0.2.1
+openai==1.35.7
+termcolor==2.4.0
+colorama==0.4.6
+python-dotenv==1.0.1
+distro==1.9.0
+PyYAML==6.0.1
+pyperclip==1.9.0
+groq==0.9.0
+anthropic==0.30.0

yolo.py

@@ -1,24 +1,22 @@
 #!/usr/bin/env python3
 
 # MIT License
-# Copyright (c) 2023 wunderwuzzi23
+# Copyright (c) 2023-2024 wunderwuzzi23
 # Greetings from Seattle! 
 
 import os
 import platform
-import openai
+from ai_model import AIModel, GroqModel, OpenAIModel, OllamaModel, AnthropicModel, AzureOpenAIModel
 import sys
 import subprocess
 import dotenv 
 import distro
 import yaml
 import pyperclip
-
 from termcolor import colored
 from colorama import init
 
-def read_config() -> any:
-
+def read_config():
   ## Find the executing directory (e.g. in case an alias is set)
   ## So we can find the config file
   yolo_path = os.path.abspath(__file__)
@@ -28,9 +26,7 @@ def read_config() -> any:
   with open(config_file, 'r') as file:
     return yaml.safe_load(file)
 
-# Construct the prompt
-def get_full_prompt(user_prompt, shell):
-
+def get_system_prompt(shell):
   ## Find the executing directory (e.g. in case an alias is set)
   ## So we can find the prompt.txt file
   yolo_path = os.path.abspath(__file__)
@@ -38,36 +34,35 @@ def get_full_prompt(user_prompt, shell):
 
   ## Load the prompt and prep it
   prompt_file = os.path.join(prompt_path, "prompt.txt")
-  pre_prompt = open(prompt_file,"r").read()
-  pre_prompt = pre_prompt.replace("{shell}", shell)
-  pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
-  prompt = pre_prompt + user_prompt
-  
-  # be nice and make it a question
+  system_prompt = open(prompt_file,"r").read()
+  system_prompt = system_prompt.replace("{shell}", shell)
+  system_prompt = system_prompt.replace("{os}", get_os_friendly_name())
+
+  return system_prompt
+
+def ensure_prompt_is_question(prompt):
   if prompt[-1:] != "?" and prompt[-1:] != ".":
     prompt+="?"
-
   return prompt
 
-def print_usage():
-  print("Yolo v0.2.1 - by @wunderwuzzi23")
+def print_usage(config):
+  print("Yolo v0.5 - by @wunderwuzzi23 (June 29, 2024)")
   print()
   print("Usage: yolo [-a] list the current directory information")
   print("Argument: -a: Prompt the user before running the command (only useful when safety is off)")
   print()
 
   print("Current configuration per yolo.yaml:")
+  print("* API          : " + str(config["api"]))
   print("* Model        : " + str(config["model"]))
   print("* Temperature  : " + str(config["temperature"]))
   print("* Max. Tokens  : " + str(config["max_tokens"]))
   print("* Safety       : " + str(bool(config["safety"])))
-
+  print("* Command Color: " + str(config["suggested_command_color"]))
 
 def get_os_friendly_name():
-  
-  # Get OS Name
   os_name = platform.system()
-  
+
   if os_name == "Linux":
     return "Linux/"+distro.name(pretty=True)
   elif os_name == "Windows":
@@ -77,31 +72,94 @@ def get_os_friendly_name():
   else:
     return os_name
 
+def chat_completion(client, query, config, shell):
+  if query == "":
+      print ("No user prompt specified.")
+      sys.exit(-1)
+ 
+  system_prompt = get_system_prompt(shell)
 
-def set_api_key():
-  # Two options for the user to specify they openai api key.
-  #1. Place a ".env" file in same directory as this with the line:
-  #   OPENAI_API_KEY="<yourkey>"
-  #   or do `export OPENAI_API_KEY=<yourkey>` before use
-  dotenv.load_dotenv()
-  openai.api_key = os.getenv("OPENAI_API_KEY")
+  response = client.chat(
+    model=config["model"],
+    messages=[
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": query}
+    ],
+    temperature=config["temperature"],
+    max_tokens=config["max_tokens"])
+ 
+  return response
+
+def check_for_issue(response):
+  prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
+  if response.lower().startswith(prefixes):
+    print(colored("There was an issue: "+response, 'red'))
+    sys.exit(-1)
+
+def check_for_markdown(response):
+  if response.count("```",2):
+    print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
+    sys.exit(-1)
+
+def missing_posix_display():
+  return 'DISPLAY' not in os.environ or not os.environ["DISPLAY"]
+
+def prompt_user_for_action(config, ask_flag, response):
+  print("Command: " + colored(response, config["suggested_command_color"], attrs=['bold']))
   
-  #2. Place a ".openai.apikey" in the home directory that holds the line:
-  #   <yourkey>
-  #   Note: This options will likely be removed in the future
-  if not openai.api_key:  #If statement to avoid "invalid filepath" error
-    home_path = os.path.expanduser("~")    
-    openai.api_key_path = os.path.join(home_path,".openai.apikey")
+  modify_snippet = ""
+  if bool(config["modify"]) == True:
+    modify_snippet = " [m]odify"
+    
+  copy_to_clipboard_snippet = " [c]opy to clipboard"
+  if os.name == "posix" and missing_posix_display():
+    if get_os_friendly_name() != "Darwin/macOS":
+      copy_to_clipboard_snippet = ""
 
-  #3. Final option is the key might be in the yolo.yaml config file
-  #   openai_apikey: <yourkey>
-  if not openai.api_key:  
-    openai.api_key = config["openai_api_key"]
+  if bool(config["safety"]) == True or ask_flag == True:
+    prompt_text = f"Execute command? [Y]es [n]o{modify_snippet}{copy_to_clipboard_snippet} ==> "
+    print(prompt_text, end = '')
+    user_input = input()
+    return user_input 
+  
+  if bool(config["safety"]) == False:
+     return "Y"
 
-if __name__ == "__main__":
+def eval_user_intent_and_execute(client, config, user_input, command, shell, ask_flag):
+  if user_input.upper() not in ["", "Y", "C", "M"]:
+    print("No action taken.")
+    return
+  
+  if user_input.upper() == "Y" or user_input == "":
+    if shell == "powershell.exe":
+      subprocess.run([shell, "/c", command], shell=False)  
+    else: 
+      # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
+      subprocess.run([shell, "-c", command], shell=False)
+  
+  if bool(config["modify"]) and user_input.upper() == "M":
+    print("Modify prompt: ", end = '')
+    modded_query = input()
+    modded_response = chat_completion(client, modded_query, config, shell)
+    check_for_issue(modded_response)
+    check_for_markdown(modded_response)
+    user_intent = prompt_user_for_action(config, ask_flag, modded_response)
+    print()
+    eval_user_intent_and_execute(client, config, user_intent, modded_response, shell, ask_flag)
+  
+  if user_input.upper() == "C":
+      if os.name == "posix" and missing_posix_display():
+        if get_os_friendly_name() != "Darwin/macOS":
+          return
+      pyperclip.copy(command) 
+      print("Copied command to clipboard.")
+
+def main():
+  init()  #Enable color output on Windows using colorama
+  dotenv.load_dotenv()
 
   config = read_config()
-  set_api_key()
+  client = AIModel.get_model_client(config)
 
   # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
   shell = os.environ.get("SHELL", "powershell.exe") 
@@ -112,7 +170,7 @@ if __name__ == "__main__":
 
   # Parse arguments and make sure we have at least a single word
   if len(sys.argv) < 2:
-    print_usage()
+    print_usage(config)
     sys.exit(-1)
 
   # Safety switch via argument -a (local override of global setting)
@@ -121,100 +179,20 @@ if __name__ == "__main__":
     ask_flag = True
     command_start_idx = 2
 
-  # To allow easy/natural use we don't require the input to be a 
-  # single string. So, the user can just type yolo what is my name?
-  # without having to put the question between ''
+  # To allow easy/natural use we don't require the input to be a single string.
+  # User can just type yolo what is my name? without having to put the question between ''
   arguments = sys.argv[command_start_idx:]
   user_prompt = " ".join(arguments)
 
-def call_open_ai(query):
-  # do we have a prompt from the user?
-  if query == "":
-      print ("No user prompt specified.")
-      sys.exit(-1)
- 
-  # Load the correct prompt based on Shell and OS and append the user's prompt
-  prompt = get_full_prompt(query, shell)
+  ## core prompting loop logic
+  result = chat_completion(client, user_prompt, config, shell) 
+  check_for_issue(result)
+  check_for_markdown(result)
 
-  # Make the first line also the system prompt
-  system_prompt = prompt[1]
-  #print(prompt)
+  users_intent = prompt_user_for_action(config, ask_flag, result)
+  print()
+  eval_user_intent_and_execute(client, config, users_intent, result, shell, ask_flag)
 
-  # Call the ChatGPT API
-  response = openai.ChatCompletion.create(
-    model=config["model"],
-    messages=[
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": prompt}
-    ],
-    temperature=config["temperature"],
-    max_tokens=config["max_tokens"],
-  )
- 
-  return response.choices[0].message.content.strip()
-
-
-#Enable color output on Windows using colorama
-init() 
-
-def check_for_issue(response):
-  prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
-  if response.lower().startswith(prefixes):
-    print(colored("There was an issue: "+response, 'red'))
-    sys.exit(-1)
-
-def check_for_markdown(response):
-  # odd corner case, sometimes ChatCompletion returns markdown
-  if response.count("```",2):
-    print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
-    sys.exit(-1)
-
-def missing_posix_display():
-  display = subprocess.check_output("echo $DISPLAY", shell=True)
-  return display == b'\n'
-
-def prompt_user_input(response):
-  print("Command: " + colored(response, 'blue'))
-  #print(config["safety"])
-
-  if bool(config["safety"]) == True or ask_flag == True:
-    prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
-    if os.name == "posix" and missing_posix_display():
-        prompt_text =  "Execute command? [Y]es [n]o [m]odify ==> "
-    print(prompt_text, end = '')
-    user_input = input()
-    return user_input 
-  
-  if config["safety"] == False:
-     return "Y"
-
-def evaluate_input(user_input, command):
-  if user_input.upper() == "Y" or user_input == "":
-    if shell == "powershell.exe":
-      subprocess.run([shell, "/c", command], shell=False)  
-    else: 
-      # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
-      subprocess.run([shell, "-c", command], shell=False)
-  
-  if user_input.upper() == "M":
-    print("Modify prompt: ", end = '')
-    modded_query = input()
-    modded_response = call_open_ai(modded_query)
-    check_for_issue(modded_response)
-    check_for_markdown(modded_response)
-    modded_user_input = prompt_user_input(modded_response)
-    print()
-    evaluate_input(modded_user_input, modded_response)
-  
-  if user_input.upper() == "C":
-      if os.name == "posix" and missing_posix_display():
-        return
-      pyperclip.copy(command) 
-      print("Copied command to clipboard.")
-
-res_command = call_open_ai(user_prompt) 
-check_for_issue(res_command)
-check_for_markdown(res_command)
-user_input = prompt_user_input(res_command)
-print()
-evaluate_input(user_input, res_command)
+if __name__ == "__main__":
+  main()
+  

yolo.yaml

@@ -1,9 +1,22 @@
-model: gpt-3.5-turbo  # If you have access to gpt-4 API already, you can update this.
+api: openai   # openai, azure, groq, ollama, anthropic         
+model: gpt-4o # if azure this is the deployment name
+              # other options: gpt-4o, llama3-8b-8192, or claude-3-5-sonnet-20240620
+
+# Azure specific (only needed if api: azure-openai)
+azure_endpoint: https://<name>.openai.azure.com
+azure_api_version: 2024-02-15-preview
+
+# Completion parameters
 temperature: 0
 max_tokens: 500
 
-# Safety: If set to False, commands returned from the AI will be run *without* prompting the user.
-safety: True
+safety: True                  # Safety: If set to False, commands from LLM run *without* prompting the user.
+modify: False                 # Enable prompt modify feature
+suggested_command_color: blue # Suggested Command Color
 
-# Open AI API Key (optional): The key can aso be provided via environment variable (OPENAI_API_KEY), .env, or ~/.openai.apikey file
-openai_api_key:
+# API Keys (optional): Preferred to use environment variables 
+# OPENAI_API_KEY, AZURE_OPENAI_API_KEY, ANTHROPIC_API_KEY or GROQ_API_KEY (.env file is also supported)
+azure_openai_api_key: 
+openai_api_key:
+groq_api_key:
+anthropic_api_key: