wehub-resource-sync
e768098d0e
Flake8 Lint / flake8 (push) Waiting to run
Spell check CI / Spell_Check (push) Waiting to run
tools_continuous_delivery / Private PyPI non-main branch release (push) Has been skipped
tools_continuous_delivery / Private PyPI main branch release (push) Failing after 2m42s
Publish Promptflow Doc / Build (push) Has been cancelled
Publish Promptflow Doc / Deploy (push) Has been cancelled
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
import ast
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import sys
|
|
from typing import Union, List
|
|
from promptflow.core import tool
|
|
from azure_open_ai import ChatLLM
|
|
from divider import Divider
|
|
from prompt import docstring_prompt, PromptLimitException
|
|
from promptflow.connections import AzureOpenAIConnection, OpenAIConnection
|
|
|
|
|
|
def get_imports(content):
|
|
tree = ast.parse(content)
|
|
import_statements = []
|
|
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Import):
|
|
for n in node.names:
|
|
import_statements.append(f"import {n.name}")
|
|
elif isinstance(node, ast.ImportFrom):
|
|
module_name = node.module
|
|
for n in node.names:
|
|
import_statements.append(f"from {module_name} import {n.name}")
|
|
|
|
return import_statements
|
|
|
|
|
|
async def async_generate_docstring(divided: List[str]):
|
|
llm = ChatLLM()
|
|
divided = list(reversed(divided))
|
|
all_divided = []
|
|
|
|
# If too many imports result in tokens exceeding the limit, please set an empty string.
|
|
modules = '' # '\n'.join(get_imports(divided[-1]))
|
|
modules_tokens = llm.count_tokens(modules)
|
|
if modules_tokens > 300:
|
|
logging.warning(f'Too many imports, the number of tokens is {modules_tokens}')
|
|
if modules_tokens > 500:
|
|
logging.warning(f'Too many imports, the number of tokens is {modules_tokens}, will set an empty string.')
|
|
modules = ''
|
|
|
|
# Divide the code into two parts if the global class/function is too long.
|
|
while len(divided):
|
|
item = divided.pop()
|
|
try:
|
|
llm.validate_tokens(llm.create_prompt(docstring_prompt(code=item, module=modules)))
|
|
except PromptLimitException as e:
|
|
logging.warning(e.message + ', will divide the code into two parts.')
|
|
divided_tmp = Divider.divide_half(item)
|
|
if len(divided_tmp) > 1:
|
|
divided.extend(list(reversed(divided_tmp)))
|
|
continue
|
|
except Exception as e:
|
|
logging.warning(e)
|
|
all_divided.append(item)
|
|
|
|
tasks = []
|
|
last_code = ''
|
|
for item in all_divided:
|
|
if Divider.has_class_or_func(item):
|
|
tasks.append(llm.async_query(docstring_prompt(last_code=last_code, code=item, module=modules)))
|
|
else: # If the code has not function or class, no need to generate docstring.
|
|
tasks.append(asyncio.sleep(0))
|
|
last_code = item
|
|
res_doc = await asyncio.gather(*tasks)
|
|
new_code = []
|
|
for i in range(len(all_divided)):
|
|
if type(res_doc[i]) is str:
|
|
new_code.append(Divider.merge_doc2code(res_doc[i], all_divided[i]))
|
|
else:
|
|
new_code.append(all_divided[i])
|
|
|
|
return new_code
|
|
|
|
|
|
@tool
|
|
def generate_docstring(divided: List[str],
|
|
connection: Union[AzureOpenAIConnection, OpenAIConnection] = None,
|
|
model: str = None):
|
|
if isinstance(connection, AzureOpenAIConnection):
|
|
os.environ["OPENAI_API_KEY"] = connection.api_key
|
|
os.environ["OPENAI_API_BASE"] = connection.api_base
|
|
os.environ["OPENAI_API_VERSION"] = connection.api_version
|
|
os.environ["API_TYPE"] = connection.api_type
|
|
elif isinstance(connection, OpenAIConnection):
|
|
os.environ["OPENAI_API_KEY"] = connection.api_key
|
|
os.environ["ORGANIZATION"] = connection.organization
|
|
if model:
|
|
os.environ["MODEL"] = model
|
|
|
|
if sys.platform.startswith("win"):
|
|
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
|
return asyncio.run(async_generate_docstring(divided))
|