summaryrefslogblamecommitdiffstats
path: root/etc/tool/improve_code.py
blob: b2e36f866aa3cd545b53839c2782bdad2be5fe40 (plain) (tree)
1
2
3
4
5
6
7
8
9
10









                                                         
                                                                             











                                  



                                  






                                                   
                






                                    
                               

                                 

import sys, re
from pathlib import Path
from os import path

sys.path.append(str(Path(__file__).parent.parent.parent))

import g4f

def read_code(text):
    if match := re.search(r"```(python|py|)\n(?P<code>[\S\s]+?)\n```", text):
        return match.group("code")
    
path = input("Path: ")

with open(path, "r") as file:
    code = file.read()

prompt = f"""
Improve the code in this file:
```py
{code}
```
Don't remove anything.
Add typehints if possible.
Don't add any typehints to kwargs.
Don't remove license comments.
"""

print("Create code...")
response = []
for chunk in g4f.ChatCompletion.create(
    model=g4f.models.gpt_35_long,
    messages=[{"role": "user", "content": prompt}],
    timeout=300,
    stream=True
):
    response.append(chunk)
    print(chunk, end="", flush=True)
print()
response = "".join(response)

if code := read_code(response):
    with open(path, "w") as file:
        file.write(code)