How I Created an AI Agent to select tests in PRs (Part 1)

_SYSTEM_PROMPT = """\
You are a senior software engineer who selects which pytest tests are relevant \
to review for a given pull request.

You will receive a JSON object with:
- "min_relevance": the minimum relevance level to include ("low", "medium", or "high")
- "changed_symbols": a list of functions/classes/methods (or whole files) modified in the PR,
each with file_path, symbol_name, symbol_type, and lines_changed
- "available_tests": a list of test functions discovered in the repository, each with
test_path, test_name, imports, and source_references

Decide which of the available tests are relevant to the changed symbols by considering:
- Whether a test's source_references directly name a changed symbol
- Whether a test's imports point to the module containing a changed symbol
- One level of transitivity: a test that references something which itself references
a changed symbol (e.g. a helper that calls the changed function)
- The semantic similarity between test names/paths and changed symbol names

Respond with ONLY a JSON array (no prose, no markdown fences, no text outside the array).
Each element must have exactly this shape:
{
"test_path": "<string>",
"test_name": "<string>",
"relevance": "<low|medium|high>",
"confidence_score": <float between 0.0 and 1.0>,
"reason": "<short string explaining why this test is relevant>"
}

Only include tests whose relevance is greater than or equal to the given min_relevance
("low" < "medium" < "high"). If no tests are relevant, return an empty array: []
"""

"""Thin wrappers so AIAgent can call Claude, Gemini, or ChatGPT interchangeably."""

from typing import Protocol

from config import Config

class LLMProvider(Protocol):
"""Something that can turn a system + user prompt into a text completion."""

def complete(self, system_prompt: str, user_prompt: str) -> str: ...

class ClaudeProvider:
"""Calls Anthropic's Claude models."""

def __init__(self, api_key: str, model: str) -> None:
from anthropic import Anthropic

self._client = Anthropic(api_key=api_key)
self._model = model

def complete(self, system_prompt: str, user_prompt: str) -> str:
response = self._client.messages.create(
model=self._model,
max_tokens=4096,
system=system_prompt,
messages=[{"role": "user", "content": user_prompt}],
)
return "".join(block.text for block in response.content if getattr(block, "type", None) == "text")

class GeminiProvider:
"""Calls Google's Gemini models."""

def __init__(self, api_key: str, model: str) -> None:
from google import genai

self._client = genai.Client(api_key=api_key)
self._model = model

def complete(self, system_prompt: str, user_prompt: str) -> str:
from google.genai import types

response = self._client.models.generate_content(
model=self._model,
contents=user_prompt,
config=types.GenerateContentConfig(system_instruction=system_prompt),
)
return response.text or ""

class ChatGPTProvider:
"""Calls OpenAI's ChatGPT models."""

def __init__(self, api_key: str, model: str) -> None:
from openai import OpenAI

self._client = OpenAI(api_key=api_key)
self._model = model

def complete(self, system_prompt: str, user_prompt: str) -> str:
response = self._client.chat.completions.create(
model=self._model,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
)
return response.choices[0].message.content or ""

def create_provider(config: Config) -> LLMProvider:
"""Builds the LLMProvider for config.provider, using the matching API key and model."""
if config.provider == "claude":
assert config.anthropic_api_key is not None
return ClaudeProvider(api_key=config.anthropic_api_key, model=config.claude_model)
if config.provider == "gemini":
assert config.gemini_api_key is not None
return GeminiProvider(api_key=config.gemini_api_key, model=config.gemini_model)
if config.provider == "chatgpt":
assert config.openai_api_key is not None
return ChatGPTProvider(api_key=config.openai_api_key, model=config.chatgpt_model)
raise ValueError(f"Unknown provider: {config.provider!r}")

"""Clones a single branch of a GitHub repository into a temporary directory for local scanning."""

import logging
import shutil
import subprocess
import tempfile

logger = logging.getLogger(__name__)

class RepoFetchError(RuntimeError):
"""Raised when the target repository/branch cannot be cloned."""

class RepoFetcher:
"""Shallow-clones a single branch of a repository so its tests can be scanned locally."""

def __init__(self, token: str | None = None) -> None:
self._token = token
self._tmp_dir: str | None = None

def fetch(self, repo: str, ref: str) -> str:
"""Clones `repo` (owner/name) at branch `ref` and returns the local checkout path."""
self._tmp_dir = tempfile.mkdtemp(prefix="ai-test-selector-")
cmd = [
"git",
"clone",
"--depth",
"1",
"--branch",
ref,
"--single-branch",
self._clone_url(repo),
self._tmp_dir,
]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
shutil.rmtree(self._tmp_dir, ignore_errors=True)
self._tmp_dir = None
raise RepoFetchError(f"Failed to clone {repo}@{ref}: {result.stderr.strip()}")
return self._tmp_dir

def cleanup(self) -> None:
"""Removes the temporary checkout, if one was created."""
if self._tmp_dir:
shutil.rmtree(self._tmp_dir, ignore_errors=True)
self._tmp_dir = None

def _clone_url(self, repo: str) -> str:
if self._token:
return f"https://x-access-token:{self._token}@github.com/{repo}.git"
return f"https://github.com/{repo}.git"