File size: 1,741 Bytes
93df7ed | 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 | import os
from llm import get_llm_client, build_prompt
from rag.repository_loader import load_repository
def generate_docstring(function_code: str) -> str:
prompt = build_prompt(function_code, "", task_type="docstring")
return get_llm_client().generate(prompt)
def generate_module_docs(file_path: str) -> str:
with open(file_path, "r", encoding="utf-8") as f:
code = f.read()
prompt = build_prompt(
f"Generate documentation for this module: {file_path}",
code,
task_type="qa",
)
return get_llm_client().generate(prompt)
def summarize_repo_structure(root_path: str) -> dict:
documents = load_repository(root_path)
tech_stack = detect_tech_stack(root_path)
file_tree = build_file_tree(documents)
return {
"total_files": len(documents),
"tech_stack": tech_stack,
"file_tree": file_tree,
}
def detect_tech_stack(root_path: str) -> list[str]:
markers = {
"requirements.txt": "Python",
"package.json": "Node.js",
"go.mod": "Go",
"pom.xml": "Java (Maven)",
}
found = []
for marker_file, tech_name in markers.items():
if os.path.exists(os.path.join(root_path, marker_file)):
found.append(tech_name)
return found
def build_file_tree(documents: list) -> list[str]:
return sorted(doc.file_path for doc in documents)
def generate_readme(root_path: str) -> str:
summary = summarize_repo_structure(root_path)
prompt = build_prompt(
f"Generate a README.md for a project with this structure: {summary}",
"",
task_type="qa",
)
return get_llm_client().generate(prompt) |