Spaces:
Sleeping
Sleeping
File size: 5,402 Bytes
5fdb69e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
from enum import Enum
from pathlib import Path
class Model(Enum):
"""
Enumeration of supported AI models.
"""
OPENAI_MODEL = "gpt-4o"
CLAUDE_MODEL = "claude-3-5-sonnet-20240620"
def get_system_message() -> str:
"""
Generate a system message for AI assistants creating docstrings.
:return: A string containing instructions for the AI assistant.
:rtype: str
"""
system_message = "You are an assistant that creates doc strings in reStructure Text format for an existing python function. "
system_message += "Respond only with an updated python function; use comments sparingly and do not provide any explanation other than occasional comments. "
system_message += "Be sure to include typing annotation for each function argument or key word argument and return object types."
return system_message
def user_prompt_for(python: str) -> str:
"""
Generate a user prompt for rewriting Python functions with docstrings.
:param python: The Python code to be rewritten.
:type python: str
:return: A string containing the user prompt and the Python code.
:rtype: str
"""
user_prompt = "Rewrite this Python function with doc strings in the reStructuredText style."
user_prompt += "Respond only with python code; do not explain your work other than a few comments. "
user_prompt += "Be sure to write a description of the function purpose with typing for each argument and return\n\n"
user_prompt += python
return user_prompt
def messages_for(python: str, system_message: str) -> list:
"""
Create a list of messages for the AI model.
:param python: The Python code to be processed.
:type python: str
:param system_message: The system message for the AI assistant.
:type system_message: str
:return: A list of dictionaries containing role and content for each message.
:rtype: list
"""
return [
{"role": "system", "content": system_message},
{"role": "user", "content": user_prompt_for(python)}
]
def write_output(output: str, file_suffix: str, file_path: Path) -> None:
"""
Write the processed output to a file.
:param output: The processed Python code with docstrings.
:type output: str
:param file_suffix: The suffix to be added to the output file name.
:type file_suffix: str
:param file_path: The path of the input file.
:type file_path: Path
:return: None
"""
code = output.replace("", "").replace("", "")
out_file = file_path.with_name(f"{file_path.stem}{file_suffix if file_suffix else ''}.py")
out_file.write_text(code)
def add_doc_string(client: object, system_message: str, file_path: Path, model: str) -> None:
"""
Add docstrings to a Python file using the specified AI model.
:param client: The AI client object.
:type client: object
:param system_message: The system message for the AI assistant.
:type system_message: str
:param file_path: The path of the input Python file.
:type file_path: Path
:param model: The AI model to be used.
:type model: str
:return: None
"""
if 'gpt' in model:
add_doc_string_gpt(client=client, system_message=system_message, file_path=file_path, model=model)
else:
add_doc_string_claude(client=client, system_message=system_message, file_path=file_path, model=model)
def add_doc_string_gpt(client: object, system_message: str, file_path: Path, model: str = 'gpt-4o') -> None:
"""
Add docstrings to a Python file using GPT model.
:param client: The OpenAI client object.
:type client: object
:param system_message: The system message for the AI assistant.
:type system_message: str
:param file_path: The path of the input Python file.
:type file_path: Path
:param model: The GPT model to be used, defaults to 'gpt-4o'.
:type model: str
:return: None
"""
code_text = file_path.read_text(encoding='utf-8')
stream = client.chat.completions.create(model=model, messages=messages_for(code_text, system_message), stream=True)
reply = ""
for chunk in stream:
fragment = chunk.choices[0].delta.content or ""
reply += fragment
print(fragment, end='', flush=True)
write_output(reply, file_suffix='_gpt', file_path=file_path)
def add_doc_string_claude(client: object, system_message: str, file_path: Path, model: str = 'claude-3-5-sonnet-20240620') -> None:
"""
Add docstrings to a Python file using Claude model.
:param client: The Anthropic client object.
:type client: object
:param system_message: The system message for the AI assistant.
:type system_message: str
:param file_path: The path of the input Python file.
:type file_path: Path
:param model: The Claude model to be used, defaults to 'claude-3-5-sonnet-20240620'.
:type model: str
:return: None
"""
code_text = file_path.read_text(encoding='utf-8')
result = client.messages.stream(
model=model,
max_tokens=2000,
system=system_message,
messages=[{"role": "user", "content": user_prompt_for(code_text)}],
)
reply = ""
with result as stream:
for text in stream.text_stream:
reply += text
print(text, end="", flush=True)
write_output(reply, file_suffix='_claude', file_path=file_path) |