last one

Author: kenxu95

action.yml

@@ -98,10 +98,10 @@ runs:
     - name: Mark the Autodoc Run as completed
       env:
         PYTHONPATH: ${{ github.action_path }}
+        AUTODOC_RUN_ID: ${{ github.autodoc_run_id }}
       run: |
         python ${{ github.action_path }}/scripts/mark_run_completed.py \
         --server-address ${{ inputs.runllm_server_address }} \
         --api-key ${{ inputs.runllm_api_key }} \
-        --run-id ${{ github.run_id }} \
         --pr-url ${{ steps.create_pr.outputs.pull-request-url }}
       shell: bash

scripts/generate_docs.py

@@ -86,7 +86,7 @@ def update_docs(
         mode: AutodocOutputMode,
         diff_by_file_path: Dict[str, str],
         openapi_spec: Optional[str],
-) -> None:
+) -> int:
     """Update the documentation for the given files in the diff.
 
     NOTE: For OpenAPI mode, we currently expect to only be processing a single inut file right now.
@@ -167,6 +167,7 @@ def update_docs(
         raise
 
     print(f"Total completion tokens: {total_token_count}, Cost: ${calculate_cost(total_token_count):.2f}")
+    return run_id
 
 
 # Command line argument parsing
@@ -222,10 +223,14 @@ def update_docs(
         raise Exception("Only 'openapi' mode is currently supported!.")
 
     client = RunLLMClient(args.server_address, args.api_key)
-    update_docs(
+    run_id = update_docs(
         client,
         gh_action_url,
         mode,
         diff_by_file_name,
         args.output_openapi_file,
-    )
+    )
+
+    # Track the Autodoc Run ID for subsequent GH action steps.
+    with open(os.getenv('GITHUB_ENV'), 'a') as env_file:
+        env_file.write(f"AUTODOC_RUN_ID={run_id}\n")

scripts/mark_run_completed.py

@@ -1,14 +1,17 @@
 import argparse
+import os
 
 from scripts.client import RunLLMClient
 
 parser = argparse.ArgumentParser(description="Mark an autodoc run as completed.")
 parser.add_argument("--server-address", type=str, required=True, help="The address of the RunLLM server")
 parser.add_argument("--api-key", type=str, required=True, help="The API key to use for the RunLLM server")
-parser.add_argument("--run-id", type=str, required=True, help="The GH action run id")
 parser.add_argument("--pr-url", type=str, required=True, help="The URL of the pull request")
 args = parser.parse_args()
 
 if __name__ == "__main__":
+    run_id = os.environ.get("AUTODOC_RUN_ID")
+    assert run_id
+
     client = RunLLMClient(args.server_address, args.api_key)
-    client.mark_completed(args.run_id, args.pr_url)
+    client.mark_completed(run_id, args.pr_url)