File size: 4,111 Bytes
d12bff5 |
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 |
import os
import zipfile
import tempfile
import google.generativeai as genai
from dotenv import load_dotenv
from doc_generator import generate_requirements_txt
load_dotenv()
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
model = genai.GenerativeModel("models/gemini-2.0-flash")
ANNOTATIONS = {
"app.py": "← Gradio + MCP server",
"README.md": "← With demo + tag \"mcp-server-track\"",
"demo_video.mp4": "← Link embedded in README"
}
PROMPT = """You are an expert software project documentation assistant.
You will write a clear, complete, and well-structured `README.md` file for a source code repository with the following files and content excerpts:
{file_summaries}
The README must contain:
1. A title
2. A short project description
3. An "Installation" section
4. A "Usage" section
5. A "Features" section
6. An "Authors" section (write "To be completed" if not detected)
7. A "License" section (write "To be completed" if not detected)
Respond only with the README.md content, without markdown ``` tags.
"""
def summarize_files(dir_path, max_files=20, max_chars=5000):
summaries = []
for root, _, files in os.walk(dir_path):
for file in files:
if file.endswith((".py", ".js", ".ts", ".java", ".md", ".json", ".txt")):
try:
with open(os.path.join(root, file), "r", encoding="utf-8") as f:
content = f.read()
rel_path = os.path.relpath(os.path.join(root, file), dir_path)
summaries.append(f"### {rel_path}\n```\n{content[:1000]}\n```")
if len("".join(summaries)) > max_chars:
break
except Exception:
continue
if len(summaries) >= max_files:
break
return "\n\n".join(summaries)
def generate_readme_from_zip(zip_file_path: str, output_dir: str) -> (str, str):
with tempfile.TemporaryDirectory() as tempdir:
with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
zip_ref.extractall(tempdir)
file_summaries = summarize_files(tempdir)
prompt = PROMPT.format(file_summaries=file_summaries)
response = model.generate_content(prompt)
readme_content = response.text.strip()
readme_path = os.path.join(output_dir, "README.md")
index_path = os.path.join(output_dir, "index.md")
os.makedirs(output_dir, exist_ok=True)
# Clean markdown code blocks if they exist
lines = readme_content.splitlines()
if len(lines) > 2:
lines = lines[1:-1] # remove the first and last lines
readme_content = "\n".join(lines)
else:
# if less than 3 lines, empty or keep as needed
readme_content = ""
with open(readme_path, "w", encoding="utf-8") as f:
f.write(readme_content)
# ✅ Generate index from tempdir (correct location of extracted files)
write_index_file(tempdir, index_path)
return readme_path, index_path
def generate_tree_structure(path: str, prefix: str = "") -> str:
entries = sorted(os.listdir(path))
lines = []
dir_name = os.path.basename(os.path.abspath(path))
lines.append(f"📁 repo/")
for idx, entry in enumerate(entries):
full_path = os.path.join(path, entry)
connector = "├── "
comment = f" {ANNOTATIONS.get(entry, '')}".rstrip()
lines.append(prefix + connector + (entry + "/" if os.path.isdir(full_path) else entry) + comment)
if os.path.isdir(full_path):
extension_prefix = "│ "
subtree = generate_tree_structure(full_path, prefix + extension_prefix)
lines.extend(subtree.splitlines()[1:]) # skip repeated dir name
lines.extend(["├── README.md",
"└── index.md"])
return "\n".join(lines)
def write_index_file(project_path: str, output_path: str):
structure = generate_tree_structure(project_path)
with open(output_path, "w", encoding="utf-8") as f:
f.write(structure)
|