Skip to content

Commit 531f547

Browse files
GangGreenTemperTatumGangGreenTemperTatum
authored
Merge pull request #30 from dreadnode/ads/refactor-rigging-pr-decorator
chore: refactor rigging pr decorator
2 parents 54161a6 + b2c277c commit 531f547

2 files changed

Lines changed: 82 additions & 165 deletions

File tree

‎.github/scripts/rigging_pr_decorator.py‎

Lines changed: 68 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -1,145 +1,77 @@
1+
# /// script
2+
# requires-python = ">=3.10"
3+
# dependencies = [
4+
# "rigging",
5+
# "typer",
6+
# ]
7+
# ///
8+
19
import asyncio
2-
import base64
3-
import os
10+
import subprocess
411
import typing as t
512

6-
from pydantic import ConfigDict, StringConstraints
13+
import typer
714

815
import rigging as rg
9-
from rigging import logger
10-
from rigging.generator import GenerateParams, Generator, register_generator
11-
12-
logger.enable("rigging")
13-
14-
MAX_TOKENS = 8000
15-
TRUNCATION_WARNING = "\n\n**Note**: Due to the large size of this diff, some content has been truncated."
16-
str_strip = t.Annotated[str, StringConstraints(strip_whitespace=True)]
17-
18-
19-
class PRDiffData(rg.Model):
20-
"""XML model for PR diff data"""
21-
22-
content: str_strip = rg.element()
23-
24-
@classmethod
25-
def xml_example(cls) -> str:
26-
return """<diff><content>example diff content</content></diff>"""
27-
28-
29-
class PRDecorator(Generator):
30-
"""Generator for creating PR descriptions"""
31-
32-
model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
33-
34-
api_key: str = ""
35-
max_tokens: int = MAX_TOKENS
36-
37-
def __init__(self, model: str, params: rg.GenerateParams) -> None:
38-
api_key = params.extra.get("api_key")
39-
if not api_key:
40-
raise ValueError("api_key is required in params.extra")
41-
42-
super().__init__(model=model, params=params, api_key=api_key)
43-
self.api_key = api_key
44-
self.max_tokens = params.max_tokens or MAX_TOKENS
45-
46-
async def generate_messages(
47-
self,
48-
messages: t.Sequence[t.Sequence[rg.Message]],
49-
params: t.Sequence[GenerateParams],
50-
) -> t.Sequence[rg.GeneratedMessage]:
51-
responses = []
52-
for message_seq, p in zip(messages, params):
53-
base_generator = rg.get_generator(self.model, params=p)
54-
llm_response = await base_generator.generate_messages([message_seq], [p])
55-
responses.extend(llm_response)
56-
return responses
57-
58-
59-
register_generator("pr_decorator", PRDecorator)
60-
61-
62-
async def generate_pr_description(diff_text: str) -> str:
63-
"""Generate a PR description from the diff text"""
64-
diff_tokens = len(diff_text) // 4
65-
if diff_tokens >= MAX_TOKENS:
66-
char_limit = (MAX_TOKENS * 4) - len(TRUNCATION_WARNING)
67-
diff_text = diff_text[:char_limit] + TRUNCATION_WARNING
68-
69-
diff_data = PRDiffData(content=diff_text)
70-
params = rg.GenerateParams(
71-
extra={
72-
"api_key": os.environ["OPENAI_API_KEY"],
73-
"diff_text": diff_text,
74-
},
75-
temperature=0.7,
76-
max_tokens=500,
77-
)
78-
79-
generator = rg.get_generator("pr_decorator!gpt-4-turbo-preview", params=params)
80-
prompt = f"""You are a helpful AI that generates clear and concise PR descriptions.
81-
Analyze the provided diff between {PRDiffData.xml_example()} tags and create a summary using exactly this format:
82-
83-
### PR Summary
84-
85-
#### Overview of Changes
86-
<overview paragraph>
87-
88-
#### Key Modifications
89-
1. **<modification title>**: <description>
90-
2. **<modification title>**: <description>
91-
3. **<modification title>**: <description>
92-
(continue as needed)
93-
94-
#### Potential Impact
95-
- <impact point 1>
96-
- <impact point 2>
97-
- <impact point 3>
98-
(continue as needed)
99-
100-
Here is the PR diff to analyze:
101-
{diff_data.to_xml()}"""
102-
103-
chat = await generator.chat(prompt).run()
104-
return chat.last.content.strip()
105-
106-
107-
async def main():
108-
"""Main function for CI environment"""
109-
if not os.environ.get("OPENAI_API_KEY"):
110-
raise ValueError("OPENAI_API_KEY environment variable must be set")
111-
112-
try:
113-
diff_text = os.environ.get("GIT_DIFF", "")
114-
if not diff_text:
115-
raise ValueError("No diff found in GIT_DIFF environment variable")
11616

117-
try:
118-
diff_text = base64.b64decode(diff_text).decode("utf-8")
119-
except Exception:
120-
padding = 4 - (len(diff_text) % 4)
121-
if padding != 4:
122-
diff_text += "=" * padding
123-
diff_text = base64.b64decode(diff_text).decode("utf-8")
124-
125-
logger.debug(f"Processing diff of length: {len(diff_text)}")
126-
description = await generate_pr_description(diff_text)
127-
128-
with open(os.environ["GITHUB_OUTPUT"], "a") as f:
129-
f.write("content<<EOF\n")
130-
f.write(description)
131-
f.write("\nEOF\n")
132-
f.write(f"debug_diff_length={len(diff_text)}\n")
133-
f.write(f"debug_description_length={len(description)}\n")
134-
debug_preview = description[:500]
135-
f.write("debug_preview<<EOF\n")
136-
f.write(debug_preview)
137-
f.write("\nEOF\n")
138-
139-
except Exception as e:
140-
logger.error(f"Error in main: {e}")
141-
raise
17+
TRUNCATION_WARNING = "\n---\n**Note**: Due to the large size of this diff, some content has been truncated."
18+
19+
20+
@rg.prompt
21+
def generate_pr_description(diff: str) -> t.Annotated[str, rg.Ctx("markdown")]: # type: ignore[empty-body]
22+
"""
23+
Analyze the provided git diff and create a PR description in markdown format.
24+
25+
<guidance>
26+
- Keep the summary concise and informative.
27+
- Use bullet points to structure important statements.
28+
- Focus on key modifications and potential impact - if any.
29+
- Do not add in general advice or best-practice information.
30+
- Write like a developer who authored the changes.
31+
- Prefer flat bullet lists over nested.
32+
- Do not include any title structure.
33+
</guidance>
34+
"""
35+
36+
37+
def get_diff(target_ref: str, source_ref: str) -> str:
38+
"""
39+
Get the git diff between two branches.
40+
"""
41+
42+
merge_base = subprocess.run(
43+
["git", "merge-base", source_ref, target_ref],
44+
capture_output=True,
45+
text=True,
46+
check=True,
47+
).stdout.strip()
48+
diff_text = subprocess.run(
49+
["git", "diff", merge_base],
50+
capture_output=True,
51+
text=True,
52+
check=True,
53+
).stdout
54+
return diff_text
55+
56+
57+
def main(
58+
target_ref: str,
59+
source_ref: str = "HEAD",
60+
generator_id: str = "openai/gpt-4o-mini",
61+
max_diff_lines: int = 1000,
62+
) -> None:
63+
"""
64+
Use rigging to generate a PR description from a git diff.
65+
"""
66+
67+
diff = get_diff(target_ref, source_ref)
68+
diff_lines = diff.split("\n")
69+
if len(diff_lines) > max_diff_lines:
70+
diff = "\n".join(diff_lines[:max_diff_lines]) + TRUNCATION_WARNING
71+
72+
description = asyncio.run(generate_pr_description.bind(generator_id)(diff))
73+
print(description)
14274

14375

14476
if __name__ == "__main__":
145-
asyncio.run(main())
77+
typer.run(main)

‎.github/workflows/rigging_pr_description.yml‎

Lines changed: 14 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@ name: Update PR Description with Rigging
22

33
on:
44
pull_request:
5-
types:
6-
- edited # Trigger when the PR is updated (e.g., title, description, or labels)
7-
- reopened # Trigger when the PR is reopened
5+
types: [opened, synchronize]
86

97
jobs:
108
update-description:
@@ -14,50 +12,37 @@ jobs:
1412
contents: read
1513

1614
steps:
17-
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 #v4.2.2
15+
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
1816
with:
19-
fetch-depth: 0
17+
fetch-depth: 0 # full history for proper diffing
2018

21-
# Get the diff first
22-
- name: Get Diff
23-
id: diff
24-
run: |
25-
git fetch origin ${{ github.base_ref }}
26-
MERGE_BASE=$(git merge-base HEAD origin/${{ github.base_ref }})
27-
# Encode the diff as base64 to preserve all characters
28-
DIFF=$(git diff $MERGE_BASE..HEAD | base64 -w 0)
29-
echo "diff=$DIFF" >> $GITHUB_OUTPUT
30-
31-
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b #v5.0.3
19+
- name: Set up Python
20+
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.0.3
3221
with:
33-
python-version: "3.11"
22+
python-version: "3.10"
3423

35-
- name: Install dependencies
24+
- name: Install uv
3625
run: |
3726
python -m pip install --upgrade pip
38-
pip cache purge
39-
pip install pydantic==2.9.1
40-
pip install rigging[all]
27+
pip install uv
4128
42-
# Generate the description using the diff
4329
- name: Generate PR Description
4430
id: description
4531
env:
46-
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
4732
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
48-
PR_NUMBER: ${{ github.event.pull_request.number }}
49-
GIT_DIFF: ${{ steps.diff.outputs.diff }}
5033
run: |
51-
python .github/scripts/rigging_pr_decorator.py
34+
DESCRIPTION=$(uv run --no-project .github/scripts/rigging_pr_decorator.py origin/${{ github.base_ref }})
35+
echo "description<<EOF" >> $GITHUB_OUTPUT
36+
echo "$DESCRIPTION" >> $GITHUB_OUTPUT
37+
echo "EOF" >> $GITHUB_OUTPUT
5238
53-
# Update the PR description
5439
- name: Update PR Description
55-
uses: nefrob/pr-description@4dcc9f3ad5ec06b2a197c5f8f93db5e69d2fdca7 #v1.2.0
40+
uses: nefrob/pr-description@4dcc9f3ad5ec06b2a197c5f8f93db5e69d2fdca7 # v1.2.0
5641
with:
5742
content: |
5843
## AI-Generated Summary
5944
60-
${{ steps.description.outputs.content }}
45+
${{ steps.description.outputs.description }}
6146
6247
---
6348

0 commit comments

Comments
 (0)