Fir changes after version was increased to 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  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.`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
+
 # 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.

ai_model.py

@@ -0,0 +1,145 @@
+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

requirements.txt

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

yolo.py

@@ -30,6 +30,8 @@ import yaml
 from termcolor import colored
 from colorama import init
 
+from ai_model import AIModel, GroqModel, OpenAIModel, OllamaModel, AnthropicModel, AzureOpenAIModel
+
 CONFIG_FILE = "yolo.yaml"
 PROMPT_FILE = "yolo.prompt"
 
@@ -87,6 +89,7 @@ def set_openai_api_key(config):
     if not openai.api_key:
         openai.api_key = config["openai_api_key"]
 
+# TODO: Add new configuration paramters
 def print_config(config):
     """
     Print config information.
@@ -135,7 +138,8 @@ def get_os_friendly_name():
 
     return os_name
 
-def get_full_prompt(user_prompt, shell):
+# TODO: Change comment
+def get_system_prompt(user_prompt, shell):
     """
     Constructs a full prompt string by appending the user's prompt to a predefined prompt template
     located in the PROMPT_FILE file.
@@ -164,44 +168,53 @@ def get_full_prompt(user_prompt, shell):
 
     ## Load the prompt and prep it
     prompt_file = os.path.join(prompt_path, PROMPT_FILE)
-    pre_prompt = open(prompt_file,"r").read()
+    system_prompt = open(prompt_file,"r").read()
-    pre_prompt = pre_prompt.replace("{shell}", shell)
+    system_prompt = system_prompt.replace("{shell}", shell)
-    pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
+    system_prompt = system_prompt.replace("{os}", get_os_friendly_name())
-    prompt = pre_prompt + user_prompt
     
-    # Be nice and make it a question.
+    return system_prompt
-    if prompt[-1:] != "?" and prompt[-1:] != ".":
-        prompt+="?"
 
-    return prompt
+def chat_completion(client, query, config):
-
-def call_open_ai(config, query):
     """
-    Do we have a prompt from the user?
+    Generate a chat-based completion for a given query using a specified model.
+    
+    This function sends a user query to a chat model and returns the generated response.
+    
+    Parameters:
+    client (object): The client object to interact with the chat service.
+    query (str): The user's query to send to the chat model.
+    config (dict): Configuration settings for the chat service, which should include:
+        - "shell" (str): Type of shell to use in the system prompt.
+        - "model" (str): The specific model to use for the chat completion.
+        - "temperature" (float): Sampling temperature to use for the response generation (higher values mean the model will take more risks).
+        - "max_tokens" (int): Maximum number of tokens to generate in the chat response.
+
+    Returns:
+    dict: The response from the chat model.
+
+    Raises:
+    SystemExit: If the query is an empty string, the function will print an error message and exit.
     """
     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.
+    system_prompt = get_system_prompt(query, config["shell"])
-    prompt = get_full_prompt(query, config["shell"])
 
-    # Make the first line also the system prompt
+    # Ensure query is a question
-    system_prompt = prompt[1]
+    if query[-1:] != "?" and query[-1:] != ".":
-    #print(prompt)
+        query += "?"
 
-    # Call the ChatGPT API
+    response = client.chat(
-    response = openai.ChatCompletion.create(
         model=config["model"],
         messages=[
             {"role": "system", "content": system_prompt},
-            {"role": "user", "content": prompt}
+            {"role": "user", "content": query}
     ],
     temperature=config["temperature"],
-        max_tokens=config["max_tokens"],
+    max_tokens=config["max_tokens"])
-    )
  
-    return response.choices[0].message.content.strip()
+    return response
 
 def check_for_issue(response):
     """
@@ -355,8 +368,9 @@ def main():
                         help='Print current configuration')
     args = parser.parse_args()
 
-    # Load configuration
+    # Load configurations and set up client
     config = read_yaml_config()
+    client = AIModel.get_model_client(config)
     set_openai_api_key(config)
 
     # Process parameters
@@ -374,12 +388,13 @@ def main():
     # Enable color output on Windows using colorama
     init()
 
-    res_command = call_open_ai(config, user_prompt)
+    result = chat_completion(client, user_prompt, config)
-    check_for_issue(res_command)
+    check_for_issue(result)
-    check_for_markdown(res_command)
+    check_for_markdown(result)
-    user_input = prompt_user_input(config, res_command)
+ 
+    user_input = prompt_user_input(config, result)
     print()
-    evaluate_input(config, user_input, res_command)
+    evaluate_input(config, user_input, result)
 
 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: If set to False, commands from LLM run *without* prompting the user.
-safety: True
+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
+# 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: