mirror of
https://github.com/arc53/DocsGPT.git
synced 2025-11-30 17:13:15 +00:00
161 lines
4.5 KiB
Python
161 lines
4.5 KiB
Python
from pathlib import Path
|
|
from langchain.text_splitter import CharacterTextSplitter
|
|
import faiss
|
|
from langchain.vectorstores import FAISS
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
from langchain.llms import OpenAI
|
|
from langchain.prompts import PromptTemplate
|
|
import pickle
|
|
import dotenv
|
|
import tiktoken
|
|
import sys
|
|
from argparse import ArgumentParser
|
|
import ast
|
|
|
|
dotenv.load_dotenv()
|
|
|
|
|
|
ps = list(Path("inputs").glob("**/*.py"))
|
|
data = []
|
|
sources = []
|
|
for p in ps:
|
|
with open(p) as f:
|
|
data.append(f.read())
|
|
sources.append(p)
|
|
|
|
|
|
|
|
|
|
|
|
def get_all_functions(source_code):
|
|
tree = ast.parse(source_code)
|
|
functions = {}
|
|
for node in tree.body:
|
|
if isinstance(node, ast.FunctionDef):
|
|
functions[node.name] = ast.unparse(node)
|
|
|
|
return functions
|
|
|
|
def get_all_functions_names(node):
|
|
functions = []
|
|
for child in node.body:
|
|
if isinstance(child, ast.FunctionDef):
|
|
functions.append(child.name)
|
|
return functions
|
|
|
|
|
|
|
|
def get_classes(source_code):
|
|
tree = ast.parse(source_code)
|
|
classes = {}
|
|
for node in tree.body:
|
|
if isinstance(node, ast.ClassDef):
|
|
classes[node.name] = get_all_functions_names(node)
|
|
return classes
|
|
|
|
def get_functions_in_class(source_code, class_name):
|
|
tree = ast.parse(source_code)
|
|
functions = []
|
|
for node in tree.body:
|
|
if isinstance(node, ast.ClassDef):
|
|
if node.name == class_name:
|
|
for function in node.body:
|
|
if isinstance(function, ast.FunctionDef):
|
|
functions.append(function.name)
|
|
return functions
|
|
|
|
|
|
functions_dict = {}
|
|
classes_dict = {}
|
|
c1 = 0
|
|
for code in data:
|
|
functions = get_all_functions(ast.parse(code))
|
|
source = str(sources[c1])
|
|
functions_dict[source] = functions
|
|
classes = get_classes(code)
|
|
classes_dict[source] = classes
|
|
c1 += 1
|
|
|
|
|
|
|
|
if not Path("outputs").exists():
|
|
Path("outputs").mkdir()
|
|
|
|
c1 = len(functions_dict)
|
|
c2 = 0
|
|
functions_dict = {}
|
|
for source, functions in functions_dict.items():
|
|
c2 += 1
|
|
print(f"Processing file {c2}/{c1}")
|
|
f1 = len(functions)
|
|
f2 = 0
|
|
source_w = source.replace("inputs/", "")
|
|
source_w = source_w.replace(".py", ".md")
|
|
# this is how we check subfolders
|
|
if "/" in source_w:
|
|
subfolders = source_w.split("/")
|
|
subfolders = subfolders[:-1]
|
|
subfolders = "/".join(subfolders)
|
|
if not Path(f"outputs/{subfolders}").exists():
|
|
Path(f"outputs/{subfolders}").mkdir(parents=True)
|
|
|
|
for name, function in functions.items():
|
|
f2 += 1
|
|
print(f"Processing function {f2}/{f1}")
|
|
prompt = PromptTemplate(
|
|
input_variables=["code"],
|
|
template="Code: \n{code}, \nDocumentation: ",
|
|
)
|
|
llm = OpenAI(temperature=0)
|
|
response = llm(prompt.format(code=function))
|
|
|
|
if not Path(f"outputs/{source_w}").exists():
|
|
with open(f"outputs/{source_w}", "w") as f:
|
|
f.write(f"# Function name: {name} \n\nFunction: \n```\n{function}\n```, \nDocumentation: \n{response}")
|
|
else:
|
|
with open(f"outputs/{source_w}", "a") as f:
|
|
f.write(f"\n\n# Function name: {name} \n\nFunction: \n```\n{function}\n```, \nDocumentation: \n{response}")
|
|
|
|
|
|
|
|
c1 = len(classes_dict)
|
|
c2 = 0
|
|
|
|
for source, classes in classes_dict.items():
|
|
c2 += 1
|
|
print(f"Processing file {c2}/{c1}")
|
|
f1 = len(classes)
|
|
f2 = 0
|
|
source_w = source.replace("inputs/", "")
|
|
source_w = source_w.replace(".py", ".md")
|
|
|
|
if "/" in source_w:
|
|
subfolders = source_w.split("/")
|
|
subfolders = subfolders[:-1]
|
|
subfolders = "/".join(subfolders)
|
|
if not Path(f"outputs/{subfolders}").exists():
|
|
Path(f"outputs/{subfolders}").mkdir(parents=True)
|
|
|
|
for name, function_names in classes.items():
|
|
print(f"Processing Class {f2}/{f1}")
|
|
f2 += 1
|
|
prompt = PromptTemplate(
|
|
input_variables=["class_name", "functions_names"],
|
|
template="Class name: {class_name} \nFunctions: {functions_names}, \nDocumentation: ",
|
|
)
|
|
llm = OpenAI(temperature=0)
|
|
response = llm(prompt.format(class_name=name, functions_names=function_names))
|
|
|
|
if not Path(f"outputs/{source_w}").exists():
|
|
with open(f"outputs/{source_w}", "w") as f:
|
|
f.write(f"# Class name: {name} \n\nFunctions: \n{function_names}, \nDocumentation: \n{response}")
|
|
else:
|
|
with open(f"outputs/{source_w}", "a") as f:
|
|
f.write(f"\n\n# Class name: {name} \n\nFunctions: \n{function_names}, \nDocumentation: \n{response}")
|
|
|
|
|
|
|
|
|
|
|
|
|