131 lines
3.9 KiB
Python
Executable File
131 lines
3.9 KiB
Python
Executable File
#!/home/brickman/anaconda3/bin/python
|
|
import requests
|
|
import json
|
|
import requests
|
|
from markdown import markdown
|
|
from bs4 import BeautifulSoup
|
|
from rich.console import Console
|
|
from rich.markdown import Markdown
|
|
import sys, os
|
|
import tempfile
|
|
|
|
|
|
try:
|
|
if len(sys.argv[1]) > 1:
|
|
question = sys.argv[1]
|
|
else:
|
|
print("Error: Please specifice a question")
|
|
exit()
|
|
except:
|
|
print("Error: Please specifice a question")
|
|
exit()
|
|
|
|
def codeparse(text_code):
|
|
html = markdown(text_code, output_format="html5")
|
|
soup = BeautifulSoup(html, "html.parser")
|
|
for code in soup.findAll('code'):
|
|
code = str(code)
|
|
code = code.replace("<code>bash\n","")
|
|
code = code.replace("<code>","")
|
|
code = code.replace("</code>","")
|
|
print(code)
|
|
|
|
|
|
def check_substring(main_string, substring):
|
|
return substring in main_string
|
|
|
|
def generate_text(model, prompt, system):
|
|
url = "http://localhost:11434/api/generate"
|
|
data = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"system": system,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": 0.6,
|
|
}
|
|
}
|
|
response = requests.post(url, json=data)
|
|
text = json.loads(response.text)
|
|
return text["response"]
|
|
|
|
out = ""
|
|
|
|
count = 0
|
|
|
|
while not check_substring(out, "CODETRUE"):
|
|
#print('Thinking....')
|
|
text = generate_text("codellama:7b", f"You are Linux Assistant. I will ask you a question in natural language and you will return a single bash code block. Make sure that you keep the result as brief as posible.\n\nPrompt:\n{question}", """You are Linux Assistant. I will ask you a question in natural language and you will return a single bash code block. The format should be as follows:
|
|
Discription here
|
|
```bash
|
|
code goes here
|
|
```
|
|
""")
|
|
|
|
#text = 'The `sudo` command takes an optional argument `-h` which stands for "help". Running `sudo -h | more` will display the help message for the `sudo` command, which includes a description of its usage and options.'
|
|
|
|
out = generate_text("codellama:7b", f"""Check the following prompt and see if it lines up with the code.
|
|
Look at the prompt and then the resulting bash code. If the result is in the correct format as speciffied below
|
|
return CODETRUE if the format is wrong or if the code is incorrect return CODEFALSE. Do not write any code or return CODE_TRUE if you write code that fixes the problem.
|
|
|
|
Correct format:
|
|
Discription here
|
|
```bash
|
|
code goes here
|
|
```
|
|
|
|
Prompt:
|
|
{question}
|
|
|
|
Bash:
|
|
{text}""", "You are a code expert.")
|
|
count+=1
|
|
if count > 4:
|
|
print('\n\nError: Could not find an answer to your problem')
|
|
exit()
|
|
|
|
# Showing llm output
|
|
console = Console()
|
|
md = Markdown("# LLM Output\n" + text)
|
|
console.print(md)
|
|
|
|
|
|
# printing code blocks
|
|
html = markdown(text, output_format="html5")
|
|
soup = BeautifulSoup(html, "html.parser")
|
|
for code in soup.findAll('code'):
|
|
code = str(code)
|
|
code = code.replace("<code>bash\n","")
|
|
code = code.replace("<code>","")
|
|
code = code.replace("</code>","")
|
|
code = code.replace("&", "&")
|
|
if len(code) < 4:
|
|
continue
|
|
|
|
print("\n\n------------------------CODE----------------------------")
|
|
console.print(Markdown(f"```bash\n{code}\n```"))
|
|
print("--------------------------------------------------------")
|
|
|
|
answer = input("Would you like to run this code [N/y/e/q]:").lower()
|
|
if answer == "y" or answer == "yes":
|
|
print("\n")
|
|
os.system(code)
|
|
print("")
|
|
|
|
elif answer == "e" or answer == "edit":
|
|
file_content = code
|
|
|
|
t = tempfile.NamedTemporaryFile(delete=False)
|
|
t.write(file_content.encode())
|
|
t.close()
|
|
try:
|
|
editor = os.environ['EDITOR']
|
|
except KeyError:
|
|
editor = 'nano'
|
|
os.system(f"{editor} {t.name}")
|
|
os.system(f"bash {t.name}")
|
|
os.system(f"rm -rf {t.name}")
|
|
elif answer == "q" or answer == "quit":
|
|
exit()
|
|
|