Script to generate .ipynb doc files

release/4.3a0
p-zach 2025-04-03 16:38:20 -04:00
parent c03dba8579
commit ef31675431
1 changed files with 105 additions and 0 deletions

View File

@ -0,0 +1,105 @@
import os
import time
import requests
import argparse
import nbformat as nbf
from openai import OpenAI
_output_folder = "output"
_gtsam_gh_base = "https://raw.githubusercontent.com/borglab/gtsam/refs/heads/develop/"
_asst_id = "asst_na7wYBtXyGU0x5t2RdcnpxzP"
_request_text = "Document the file found at {}."
def is_url_valid(url):
"""Verify that the supplied URL does not return a 404."""
try:
response = requests.head(url, allow_redirects=True)
return response.status_code != 404
except requests.RequestException:
return False
def save_ipynb(text: str, file_path: str):
"""Save text to a single Markdown cell in a new .ipynb file."""
script_dir = os.path.dirname(os.path.abspath(__file__))
output_dir = os.path.join(script_dir, _output_folder)
os.makedirs(output_dir, exist_ok=True)
output_file = os.path.splitext(os.path.basename(file_path))[0] + ".ipynb"
output_full_path = os.path.join(output_dir, output_file)
nb = nbf.v4.new_notebook()
new_cell = nbf.v4.new_markdown_cell(text)
nb['cells'].append(new_cell)
with open(output_full_path, 'w', encoding='utf-8') as file:
nbf.write(nb, file)
return output_file
def generate_ipynb(file_path: str, openai_client):
"""Generate an interactive Python notebook for the given GTSAM header file.
Args:
file_path (str): The fully-qualified path from the root of the gtsam
repository to the header file that will be documented.
openai_client (openai.OpenAI): The OpenAI client to use.
"""
# Create the URL to get the header file from.
url = _gtsam_gh_base + file_path
if not is_url_valid(url):
print(f"{url} was not found on the server, or an error occurred.")
return
print(f"Sending request to OpenAI to document {url}.")
# Create a new thread and send the request
thread = openai_client.beta.threads.create()
openai_client.beta.threads.messages.create(
thread_id=thread.id, role="user", content=_request_text.format(url))
run = openai_client.beta.threads.runs.create(thread_id=thread.id,
assistant_id=_asst_id)
print("Waiting for the assistant to process the request...")
# Wait for request to be processed
while True:
run_status = openai_client.beta.threads.runs.retrieve(
thread_id=thread.id, run_id=run.id)
if run_status.status == "completed":
break
time.sleep(2)
print("Request processed. Retrieving response...")
# Fetch messages
messages = openai_client.beta.threads.messages.list(thread_id=thread.id)
# Retrieve response text and strip ```markdown ... ```
text = messages.data[0].content[0].text.value.strip('`').strip('markdown')
# Write output to file
output_filename = save_ipynb(text, file_path)
print(f"Response retrieved. Find output in {output_filename}.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
prog="gpt_generate",
description=
"Generates .ipynb documentation files given paths to GTSAM header files."
)
parser.add_argument(
"file_paths",
nargs='+',
help="The paths to the header files from the root gtsam directory.")
args = parser.parse_args()
# Retrieves API key from environment variable OPENAI_API_KEY
client = OpenAI()
for file_path in args.file_paths:
generate_ipynb(file_path, client)