Skip to content

Commit d5b47a4

Browse files
cristipufucristipufu
committed
fix: agent sample
1 parent 738043c commit d5b47a4

File tree

2 files changed

+96
-85
lines changed

2 files changed

+96
-85
lines changed

‎samples/github-helper-agent/main.py‎

Lines changed: 95 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import json
22
import os
3+
import re
34
from contextlib import asynccontextmanager
4-
from typing import Optional
5+
from typing import List, Optional
56

67
import dotenv
78
from langchain_anthropic import ChatAnthropic
@@ -15,6 +16,8 @@
1516

1617
dotenv.load_dotenv()
1718

19+
AI_GENERATED_LABEL = "_/ai generated"
20+
1821

1922
class PullRequestInfo(BaseModel):
2023
"""Input parameters with Pull Request details"""
@@ -26,6 +29,16 @@ class PullRequestInfo(BaseModel):
2629
command: str = Field(default="review")
2730

2831

32+
class PullRequestComment(BaseModel):
33+
"""Human or AI message extracted from Pull Request reviews/comments/issues"""
34+
35+
id: int
36+
body: str
37+
role: str
38+
in_reply_to: Optional[str]
39+
created_at: Optional[str]
40+
41+
2942
class GraphState(AgentState):
3043
"""Graph state"""
3144

@@ -36,6 +49,39 @@ class GraphState(AgentState):
3649
command: str
3750

3851

52+
def process_comment(comment) -> PullRequestComment:
53+
"""Process a Pull Request GitHub comment and return a PullRequestMessage."""
54+
55+
in_reply_to = None
56+
created_at = comment.get("created_at") or comment.get("submitted_at")
57+
if comment["body"].startswith(AI_GENERATED_LABEL):
58+
# Parse in_reply_to from the AI label
59+
match = re.search(r"\[(\d+)\]", comment["body"])
60+
if match:
61+
in_reply_to = match.group(1)
62+
return PullRequestComment(
63+
body=comment["body"].strip(),
64+
role="assistant",
65+
created_at=created_at,
66+
id=comment["id"],
67+
in_reply_to=in_reply_to,
68+
)
69+
else:
70+
# /help confuses the LLM
71+
message = comment["body"].replace("/help", "").strip()
72+
path = comment.get("path")
73+
line = comment.get("line")
74+
if path and line:
75+
message = f"Comment on {path} line {line}: {message}"
76+
return PullRequestComment(
77+
body=message,
78+
role="user",
79+
created_at=created_at,
80+
id=comment["id"],
81+
in_reply_to=None,
82+
)
83+
84+
3985
@asynccontextmanager
4086
async def make_graph():
4187
async with sse_client(
@@ -52,7 +98,7 @@ async def make_graph():
5298
async def hydrate_history(input: PullRequestInfo) -> GraphState:
5399
"""Fetch PR context at the start of the workflow."""
54100

55-
context_messages = []
101+
pr_history: List[PullRequestComment] = []
56102

57103
# Fetch PR details
58104
tool_result = await session.call_tool(
@@ -65,15 +111,19 @@ async def hydrate_history(input: PullRequestInfo) -> GraphState:
65111
)
66112

67113
pr_details = json.loads(tool_result.content[0].text)
68-
pr_body = pr_details["body"] if "body" in pr_details else ""
69-
70-
# Add PR details as a system message
71-
context_messages.append(
72-
{
73-
"role": "system",
74-
"content": f"Pull Request #{input.pullNumber} by {pr_details['user']['login']}\nTitle: {pr_details['title']}\nDescription: {pr_body}",
75-
}
114+
pr_body = pr_details.get("body") or ""
115+
116+
# Add PR details as the first human message
117+
pr_history.append(
118+
PullRequestComment(
119+
body=f"Pull Request #{input.pullNumber} by {pr_details['user']['login']}\nTitle: {pr_details['title']}\nDescription: {pr_body}",
120+
role="user",
121+
created_at=pr_details["created_at"],
122+
id=pr_details["id"],
123+
in_reply_to=None,
124+
)
76125
)
126+
77127
# Fetch PR comments
78128
tool_result = await session.call_tool(
79129
"get_pull_request_comments",
@@ -83,26 +133,9 @@ async def hydrate_history(input: PullRequestInfo) -> GraphState:
83133
"pullNumber": input.pullNumber,
84134
},
85135
)
86-
87136
comments = json.loads(tool_result.content[0].text)
88-
89-
# Add each comment as a user or assistant message
90137
for comment in comments:
91-
# Bot comments
92-
if comment["body"].startswith("_/ai generated_\n"):
93-
context_messages.append(
94-
{
95-
"role": "assistant",
96-
"content": comment["body"]
97-
.replace("_/ai generated_\n", "")
98-
.strip(),
99-
}
100-
)
101-
# User comments
102-
else:
103-
context_messages.append(
104-
{"role": "user", "content": comment["body"].strip()}
105-
)
138+
pr_history.append(process_comment(comment))
106139

107140
# Fetch PR review comments
108141
tool_result = await session.call_tool(
@@ -113,71 +146,39 @@ async def hydrate_history(input: PullRequestInfo) -> GraphState:
113146
"pullNumber": input.pullNumber,
114147
},
115148
)
116-
117149
review_comments = json.loads(tool_result.content[0].text)
118-
119-
# Add review comments as messages
120150
for comment in review_comments:
121-
if comment["body"].startswith("_/ai generated_\n"):
122-
context_messages.append(
123-
{
124-
"role": "assistant",
125-
"content": comment["body"]
126-
.replace("_/ai generated_\n", "")
127-
.strip(),
128-
}
129-
)
130-
else:
131-
context_messages.append(
132-
{
133-
"role": "user",
134-
"content": f"Comment on {comment['path']} line {comment['line']}: {comment['body']}",
135-
}
136-
)
151+
pr_history.append(process_comment(comment))
137152

138-
# Fetch PR review comments
153+
# Fetch issue comments
139154
tool_result = await session.call_tool(
140155
"get_issue_comments",
141156
{
142157
"owner": input.owner,
143158
"repo": input.repo,
144159
"issue_number": input.pullNumber,
145160
"page": 1,
146-
"per_page": 100
161+
"per_page": 100,
147162
},
148163
)
149-
150164
issue_comments = json.loads(tool_result.content[0].text)
151-
152-
# Add review comments as messages
153165
for comment in issue_comments:
154-
if comment["body"].startswith("_/ai generated_\n"):
155-
context_messages.append(
156-
{
157-
"role": "assistant",
158-
"content": comment["body"]
159-
.replace("_/ai generated_\n", "")
160-
.strip(),
161-
}
162-
)
163-
else:
164-
context_messages.append(
165-
{"role": "user", "content": comment["body"].strip()}
166-
)
167-
168-
# Add the input message as the last user message
169-
if input.commentNumber:
170-
context_messages.append(
171-
{
172-
"role": "user",
173-
"content": f"The command is {input.command} with review comment id #{input.commentNumber}",
174-
}
175-
)
176-
else:
177-
context_messages.append(
166+
pr_history.append(process_comment(comment))
167+
168+
# Sort chat items by created_at timestamp
169+
pr_history.sort(key=lambda item: item.created_at)
170+
171+
messages = []
172+
for item in pr_history:
173+
messages.append(
178174
{
179-
"role": "user",
180-
"content": f"Please {input.command} this PR and provide detailed feedback.",
175+
"role": item.role,
176+
"content": item.body,
177+
"metadata": {
178+
"id": item.id,
179+
"created_at": item.created_at,
180+
"in_reply_to": item.in_reply_to,
181+
},
181182
}
182183
)
183184

@@ -187,18 +188,26 @@ async def hydrate_history(input: PullRequestInfo) -> GraphState:
187188
"repo": input.repo,
188189
"pull_number": input.pullNumber,
189190
"in_reply_to": input.commentNumber,
190-
"messages": context_messages,
191+
"messages": messages,
191192
}
192193

193194
def pr_prompt(state: GraphState) -> GraphState:
194-
"""Create a prompt that incorporates PR data and the requested command."""
195+
in_reply_to = state.get("in_reply_to")
196+
if in_reply_to:
197+
label = f"{AI_GENERATED_LABEL} [{in_reply_to}]_\n"
198+
command_message = f"The CURRENT command is '{state['command']}' from review comment id #{in_reply_to}. EXECUTE the CURRENT command and provide detailed feedback."
199+
else:
200+
label = f"{AI_GENERATED_LABEL}_\n"
201+
command_message = f"The CURRENT command is '{state['command']}. EXECUTE the CURRENT command and provide detailed feedback."
202+
203+
"""Create a prompt that incorporates PR data."""
195204
system_message = f"""You are a professional developer with experience in code reviews and GitHub pull requests.
196205
You are working with repo: {state["owner"]}/{state["repo"]}, PR #{state["pull_number"]}.
197206
198207
IMPORTANT INSTRUCTIONS:
199208
1. ALWAYS get the contents of the changed files in the current PR
200209
2. ALWAYS use the contents of the changed files as context when replying to a user command.
201-
3. Always start your responses with "_/ai generated_\n" to properly tag your comments.
210+
3. ALWAYS start your responses with "{label}" to properly tag your comments.
202211
4. When you reply to a comment, make sure to address the specific request.
203212
5. When reviewing code, be thorough but constructive - point out both issues and good practices.
204213
6. You MUST use the available tools to post your response as a PR comment or perform the PR code review.
@@ -215,11 +224,13 @@ def pr_prompt(state: GraphState) -> GraphState:
215224
2. Analyze the available PR data and understand what the user is asking for
216225
3. Use the appropriate tools to gather any additional information needed
217226
4. Prepare your response based on the request
218-
5. Based on the user's command:
219-
- if the command has a review comment id, REPLY TO THE REVIEW COMMENT WITH THE SPECIFIED ID
220-
- else POST PULL REQUEST REVIEW
227+
5. [IMPORTANT] Based on the user's command:
228+
- if the command has a review comment id, REPLY TO THE REVIEW COMMENT WITH THE SPECIFIED ID using tool add_pull_request_review_comment
229+
- else POST PULL REQUEST REVIEW using tool create_pull_request_review
221230
222231
Remember: The user wants specific, actionable feedback and help with their code.
232+
233+
{command_message}
223234
"""
224235

225236
return [{"role": "system", "content": system_message}] + state[

‎samples/github-helper-agent/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "github-helper-agent"
3-
version = "0.0.5"
3+
version = "0.0.8"
44
description = "An automated agent that reviews GitHub pull requests and provides feedback"
55
authors = [{ name = "Cristi Pufu", email = "[email protected]" }]
66
dependencies = [

0 commit comments

Comments
 (0)