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)