refactor: extract helper functions, move scripts to scripts/agents

Author: jescalada

.github/workflows/pr-quality-check.yml

@@ -30,4 +30,4 @@ jobs:
           PR_NUMBER: ${{ github.event.pull_request.number }}
           PR_TITLE: ${{ github.event.pull_request.title }}
           REPO_NAME: ${{ github.repository }}
-        run: python scripts/pr_checker_agent.py
+        run: python scripts/agents/pr_checker_agent.py

.github/workflows/security-review.yml

@@ -42,4 +42,4 @@ jobs:
           PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
           REPO_NAME: ${{ github.repository }}
           TRIGGER: ${{ github.event_name }}
-        run: python scripts/security_review_agent.py
+        run: python scripts/agents/security_review_agent.py

.github/workflows/triage.yml

@@ -30,4 +30,4 @@ jobs:
           ISSUE_NUMBER: ${{ github.event.issue.number }}
           ISSUE_TITLE: ${{ github.event.issue.title }}
           REPO_NAME: ${{ github.repository }}
-        run: python scripts/triage_agent.py
+        run: python scripts/agents/triage_agent.py

scripts/agents/helpers.py

@@ -0,0 +1,33 @@
+def validate_api_keys():
+    valid_api_keys = ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY"]
+    if not any(os.environ.get(k) for k in valid_api_keys):
+        raise ValueError("No API key is set")
+
+
+def validate_env_vars(env_vars: list[str]):
+    for env_var in env_vars:
+        if not os.environ.get(env_var):
+            raise ValueError(f"{env_var} is not set")
+
+
+def run_agent(messages: list, tools: list, handle_tool_call, model: str):
+    while True:
+        response = litellm.completion(
+            model=model, messages=messages, tools=tools, temperature=0
+        )
+        message = response.choices[0].message
+        if message.content:
+            print(f"[agent] {message.content}")
+        messages.append(message.model_dump(exclude_none=True))
+        if response.choices[0].finish_reason == "stop" or not message.tool_calls:
+            break
+        tool_results = []
+        for tool_call in message.tool_calls:
+            inputs = json.loads(tool_call.function.arguments)
+            result = handle_tool_call(tool_call.function.name, inputs)
+            tool_results.append({
+                "role": "tool",
+                "tool_call_id": tool_call.id,
+                "content": result,
+            })
+        messages.extend(tool_results)

scripts/pr_checker_agent.py scripts/agents/pr_checker_agent.pyscripts/pr_checker_agent.py renamed to scripts/agents/pr_checker_agent.py

@@ -2,6 +2,7 @@
 import json
 import litellm
 from github import Github, Auth
+from helpers import validate_env_vars, validate_api_keys, run_agent
 
 # Setup
 
@@ -11,14 +12,8 @@
 author = os.environ["AUTHOR_USERNAME"]
 
 MODEL = os.environ["MODEL"]
-for env_var in ["GITHUB_TOKEN", "REPO_NAME", "PR_NUMBER", "AUTHOR_USERNAME", "MODEL"]:
-    if not os.environ[env_var]:
-        raise ValueError(f"{env_var} is not set")
-
-valid_api_keys = ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY"]
-if not any(os.environ.get(api_key) for api_key in valid_api_keys):
-    raise ValueError("No API key is set")
-
+validate_env_vars(["GITHUB_TOKEN", "REPO_NAME", "PR_NUMBER", "AUTHOR_USERNAME", "MODEL"])
+validate_api_keys()
 
 # Tools
 
@@ -126,36 +121,7 @@ def run_pr_review_agent():
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": build_initial_message()},
     ]
-
-    while True:
-        response = litellm.completion(
-            model=MODEL,
-            messages=messages,
-            tools=TOOLS,
-            temperature=0,
-        )
-
-        message = response.choices[0].message
-
-        if message.content:
-            print(f"[agent] {message.content}")
-
-        messages.append(message.model_dump(exclude_none=True))
-
-        if response.choices[0].finish_reason == "stop" or not message.tool_calls:
-            break
-
-        tool_results = []
-        for tool_call in message.tool_calls:
-            inputs = json.loads(tool_call.function.arguments)
-            result = handle_tool_call(tool_call.function.name, inputs)
-            tool_results.append({
-                "role": "tool",
-                "tool_call_id": tool_call.id,
-                "content": result,
-            })
-
-        messages.extend(tool_results)
+    run_agent(messages, TOOLS, handle_tool_call, MODEL)
 
 
 if __name__ == "__main__":

scripts/security_review_agent.py scripts/agents/security_review_agent.pyscripts/security_review_agent.py renamed to scripts/agents/security_review_agent.py

@@ -2,6 +2,7 @@
 import json
 import litellm
 from github import Github, Auth
+from helpers import validate_env_vars, validate_api_keys, run_agent
 
 # Setup
 
@@ -10,13 +11,8 @@
 pr = repo.get_pull(int(os.environ["PR_NUMBER"]))
 
 MODEL = os.environ["MODEL"]
-for env_var in ["GITHUB_TOKEN", "REPO_NAME", "PR_NUMBER", "MODEL"]:
-    if not os.environ[env_var]:
-        raise ValueError(f"{env_var} is not set")
-
-valid_api_keys = ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY"]
-if not any(os.environ.get(api_key) for api_key in valid_api_keys):
-    raise ValueError("No API key is set")
+validate_env_vars(["GITHUB_TOKEN", "REPO_NAME", "PR_NUMBER", "MODEL"])
+validate_api_keys()
 
 IGNORED_FILENAMES = set(os.environ.get(
     "IGNORED_FILENAMES",
@@ -178,37 +174,7 @@ def run_security_review_agent():
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": build_initial_message()},
     ]
-
-    while True:
-        response = litellm.completion(
-            model=MODEL,
-            messages=messages,
-            tools=TOOLS,
-            temperature=0,
-        )
-
-        message = response.choices[0].message
-
-        if message.content:
-            print(f"[agent] {message.content}")
-
-        messages.append(message.model_dump(exclude_none=True))
-
-        if response.choices[0].finish_reason == "stop" or not message.tool_calls:
-            break
-
-        tool_results = []
-        for tool_call in message.tool_calls:
-            inputs = json.loads(tool_call.function.arguments)
-            result = handle_tool_call(tool_call.function.name, inputs)
-            print(f"[tool] {tool_call.function.name}: {result}")
-            tool_results.append({
-                "role": "tool",
-                "tool_call_id": tool_call.id,
-                "content": result,
-            })
-
-        messages.extend(tool_results)
+    run_agent(messages, TOOLS, handle_tool_call, MODEL)
 
 
 if __name__ == "__main__":

scripts/triage_agent.py scripts/agents/triage_agent.pyscripts/triage_agent.py renamed to scripts/agents/triage_agent.py

@@ -2,6 +2,7 @@
 import json
 import litellm
 from github import Github, Auth
+from helpers import validate_env_vars, validate_api_keys, run_agent
 
 # Setup
 
@@ -13,13 +14,8 @@
 AVAILABLE_LABELS = os.environ.get("AVAILABLE_LABELS", "bug,enhancement,question,documentation,needs-info")
 MODEL = os.environ["MODEL"]
 
-for env_var in ["GITHUB_TOKEN", "REPO_NAME", "ISSUE_NUMBER", "ISSUE_TITLE", "ISSUE_BODY", "MODEL"]:
-    if not os.environ[env_var]:
-        raise ValueError(f"{env_var} is not set")
-
-valid_api_keys = ["ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY"]
-if not any(os.environ.get(api_key) for api_key in valid_api_keys):
-    raise ValueError("No API key is set")
+validate_env_vars(["GITHUB_TOKEN", "REPO_NAME", "ISSUE_NUMBER", "ISSUE_TITLE", "ISSUE_BODY", "MODEL"])
+validate_api_keys()
 
 # Tools
 
@@ -224,33 +220,7 @@ def run_triage_agent():
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": build_initial_message()},
     ]
-
-    while True:
-        response = litellm.completion(
-            model=MODEL,
-            messages=messages,
-            tools=TOOLS,
-            temperature=0,
-        )
-
-        message = response.choices[0].message
-        messages.append(message.model_dump(exclude_none=True))
-
-        finish_reason = response.choices[0].finish_reason
-        if finish_reason == "stop" or not message.tool_calls:
-            break
-
-        tool_results = []
-        for tool_call in message.tool_calls:
-            inputs = json.loads(tool_call.function.arguments)
-            result = handle_tool_call(tool_call.function.name, inputs)
-            tool_results.append({
-                "role": "tool",
-                "tool_call_id": tool_call.id,
-                "content": result,
-            })
-
-        messages.extend(tool_results)
+    run_agent(messages, TOOLS, handle_tool_call, MODEL)
 
 
 if __name__ == "__main__":