From 256b616f0244af77d2c423e723b3e341707bb5ea Mon Sep 17 00:00:00 2001 From: Heiko Joerg Schick Date: Mon, 19 Aug 2024 19:09:47 +0200 Subject: [PATCH] Fir changes after version was increased to 0.5 --- README.md | 22 +++++++ ai_model.py | 145 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 17 +++--- yolo.py | 83 ++++++++++++++++----------- yolo.yaml | 23 ++++++-- 5 files changed, 244 insertions(+), 46 deletions(-) create mode 100644 ai_model.py diff --git a/README.md b/README.md index 9badad6..b236954 100644 --- a/README.md +++ b/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. diff --git a/ai_model.py b/ai_model.py new file mode 100644 index 0000000..f946cf5 --- /dev/null +++ b/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 diff --git a/requirements.txt b/requirements.txt index 0a8582d..c897057 100644 --- a/requirements.txt +++ b/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 diff --git a/yolo.py b/yolo.py index f9273f0..bd85cb4 100755 --- a/yolo.py +++ b/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() - pre_prompt = pre_prompt.replace("{shell}", shell) - pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name()) - prompt = pre_prompt + user_prompt + 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 - # Be nice and make it a question. - if prompt[-1:] != "?" and prompt[-1:] != ".": - prompt+="?" - - return prompt - -def call_open_ai(config, query): +def chat_completion(client, query, config): """ - 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) + + system_prompt = get_system_prompt(query, config["shell"]) - # Load the correct prompt based on shell and OS and append the user's prompt. - prompt = get_full_prompt(query, config["shell"]) + # Ensure query is a question + if query[-1:] != "?" and query[-1:] != ".": + query += "?" - # Make the first line also the system prompt - system_prompt = prompt[1] - #print(prompt) - - # Call the ChatGPT API - response = openai.ChatCompletion.create( + response = client.chat( 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() + {"role": "user", "content": query} + ], + temperature=config["temperature"], + max_tokens=config["max_tokens"]) + + 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) - check_for_issue(res_command) - check_for_markdown(res_command) - user_input = prompt_user_input(config, res_command) + result = chat_completion(client, user_prompt, config) + check_for_issue(result) + check_for_markdown(result) + + 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() diff --git a/yolo.yaml b/yolo.yaml index d92a15b..3ea6a84 100644 --- a/yolo.yaml +++ b/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://.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: \ No newline at end of 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: \ No newline at end of file