schihei/yolo-ai-cmdbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fir changes after version was increased to 0.5

Browse Source
This commit is contained in:
Heiko Joerg Schick
2024-08-19 19:09:47 +02:00
parent 84a8876016
commit 256b616f02
5 changed files with 244 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+22
View File
@@ -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
+145
View File
@@ -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
+10 -7
View File
@@ -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
+49 -34
View File
@@ -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()
yolo.yaml
+18 -5
View File
@@ -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: