307 lines
12 KiB
Python
307 lines
12 KiB
Python
from abc import abstractmethod
|
|
from dataclasses import dataclass
|
|
from functools import wraps
|
|
from typing import Optional, Iterable, Any
|
|
|
|
from rhodecode.apps.ai_agents.ai_settings import AISettings
|
|
from rhodecode.lib.codeblocks import DiffSet
|
|
from rhodecode.lib.vcs.nodes import FileNode
|
|
|
|
|
|
FILE = "FILE:"
|
|
|
|
|
|
class AIServiceError(Exception):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class Response:
|
|
model: str
|
|
message: str
|
|
error: bool = False
|
|
|
|
|
|
@dataclass
|
|
class Request:
|
|
content: str
|
|
|
|
|
|
@dataclass
|
|
class PingRequest(Request):
|
|
function_name: str
|
|
role: str
|
|
|
|
|
|
@dataclass
|
|
class Review:
|
|
role: str
|
|
msg: str
|
|
|
|
|
|
@dataclass
|
|
class CodeReviewRequest(Request):
|
|
content: list[Review]
|
|
function_name: str
|
|
|
|
|
|
@dataclass
|
|
class CustomFunctions:
|
|
Ping: str = "ping_pong"
|
|
Code_review: str = "return_code_review"
|
|
|
|
|
|
def wrap_ai_exceptions(f):
|
|
@wraps(f)
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
return f(*args, **kwargs)
|
|
except Exception as e:
|
|
exc = AIServiceError(str(e))
|
|
# to not lose the original traceback
|
|
exc.__cause__ = e
|
|
raise exc.with_traceback(e.__traceback__)
|
|
|
|
return wrapper
|
|
|
|
|
|
class AIServiceBase:
|
|
DEFAULT_BASIC_REVIEW_POINTS = [
|
|
"Correctness and edge cases (logic errors, boundary conditions, invalid inputs).",
|
|
"Error handling and resilience (fail-fast where appropriate, clear propagation, retries/backoff, cleanup).",
|
|
"Security (input validation/sanitization, injection risks, unsafe eval/exec, authn/z, secret handling, serialization).",
|
|
"Readability and maintainability (clear names, comments/docs where helpful, remove dead code, consistent formatting).",
|
|
"Interface & API design (encapsulation, stable contracts, backward compatibility, minimal surface area).",
|
|
"Performance & complexity (hot paths, unnecessary allocations/work, algorithmic complexity, scalability).",
|
|
"Concurrency & asynchrony (race conditions, synchronization, thread/process safety, async/await or equivalents).",
|
|
"Resource management (files, network, DBs; timeouts; quotas; connection pooling; deterministic cleanup).",
|
|
"Observability (useful logging, metrics, tracing; avoid sensitive data in logs; actionable error messages).",
|
|
"Testability & testing (deterministic seams, unit/integration tests, fixtures/mocks, meaningful coverage).",
|
|
"Dependency & supply-chain hygiene (version constraints, provenance, minimal deps, portability).",
|
|
"Portability & interoperability (standards compliance, platform differences, encoding/locale issues).",
|
|
]
|
|
|
|
def __init__(self, model_settings: AISettings):
|
|
self.model_settings = model_settings
|
|
self._validate_mandatory_settings()
|
|
|
|
@wrap_ai_exceptions
|
|
def _validate_mandatory_settings(self):
|
|
api_key = self.model_settings.api_key
|
|
assert api_key is not None and api_key, "API key is required"
|
|
|
|
name = self.model_settings.model_name
|
|
assert name is not None and name, "Model name is required"
|
|
|
|
version = self.model_settings.model_version
|
|
assert version is not None and version, "Model version is required"
|
|
|
|
@wrap_ai_exceptions
|
|
def ping(self) -> Response:
|
|
resp = self._get_response(self._get_ping_request())
|
|
return self._transform(resp)
|
|
|
|
@wrap_ai_exceptions
|
|
def code_review(self, pr_diffset: DiffSet, *args, **kwargs) -> Response:
|
|
resp = self._get_response(self._get_review_requests(pr_diffset, *args, **kwargs))
|
|
return self._transform(resp)
|
|
|
|
def get_code_review_instructions(self, basic_points, review_content, custom_instructions=None):
|
|
user_msg = (
|
|
"BASIC REVIEW INSTRUCTIONS:\n"
|
|
+ "\n".join(f"- {point}" for point in basic_points)
|
|
+ ("\n\nADDITIONAL CUSTOM INSTRUCTIONS:\n" + custom_instructions if custom_instructions else "")
|
|
+ "\n\nIMPORTANT OUTPUT RULES:\n"
|
|
f"- Use the function call to return `{CustomFunctions.Code_review}`: a list of tuples "
|
|
" [line_number:int (1-based, exactly same index as it was in the request per file), line_text:str, suggestion:str].\n"
|
|
"- Only include tuples for lines with a concrete, actionable suggestion.\n"
|
|
"- Keep suggestions concise and specific.\n"
|
|
"- If many lines share the same issue, include a representative subset and list other line numbers.\n"
|
|
"- Do not include overall prose; only return via the function.\n"
|
|
+ "\nFILES (numbered per file):\n"
|
|
+ self.numbered_code_block(review_content)
|
|
)
|
|
return user_msg
|
|
|
|
def numbered_code_block(self, review_content: list[dict]) -> str:
|
|
"""
|
|
Render code with explicit 1-based line numbers so the model can reference them precisely.
|
|
"""
|
|
parts: list[str] = []
|
|
for content in review_content:
|
|
full_code = content["content"].splitlines()
|
|
only_include_lines: list[int] = content["changed_lines"]
|
|
width = max(2, len(str(len(full_code)))) if full_code else 2
|
|
part_template = f"{FILE} %s\nTYPE: %s"
|
|
parts.append(part_template % (content["file_name"], content["file_type"]))
|
|
if full_code:
|
|
for i, line in enumerate(full_code, start=1):
|
|
if i in only_include_lines:
|
|
# include only changes without losing index
|
|
parts.append(f"{str(i).rjust(width)} | {line}")
|
|
else:
|
|
parts.append("(empty file)")
|
|
parts.append("") # blank line between files
|
|
return "\n".join(parts).rstrip()
|
|
|
|
def _get_response_input(self, request: PingRequest | CodeReviewRequest) -> list[dict[str, str]]:
|
|
_input = None
|
|
if isinstance(request, CodeReviewRequest):
|
|
return [
|
|
{
|
|
"role": review.role,
|
|
"content": review.msg,
|
|
}
|
|
for review in request.content
|
|
]
|
|
|
|
if isinstance(request, PingRequest):
|
|
return [
|
|
{
|
|
"role": request.role,
|
|
"content": request.content,
|
|
}
|
|
]
|
|
|
|
raise ValueError("Unknown request type: %s" % type(request))
|
|
|
|
def _extract_changes(self, pr_diffset: DiffSet) -> list[Any]:
|
|
review_content = []
|
|
for f in pr_diffset.files:
|
|
target_file_node: FileNode = f.get("target_filenode", None)
|
|
if not target_file_node:
|
|
self.log.warning("No target file node, skipping it")
|
|
continue
|
|
|
|
file_name = f.get("patch", {}).get("filename", "unknown")
|
|
if not target_file_node.str_content:
|
|
self.log.debug("No content, skipping file: %s", file_name)
|
|
continue
|
|
|
|
modified_lines_nums = []
|
|
for hunk in f.get("hunks", []):
|
|
for line in hunk.get("lines", []):
|
|
action = line.get("modified", {}).get("action", "").strip()
|
|
if action and action != "-":
|
|
modified_lines_nums.append(line.get("modified", {}).get("lineno"))
|
|
|
|
file_content = {
|
|
"content": target_file_node.str_content,
|
|
"changed_lines": modified_lines_nums,
|
|
"file_name": file_name,
|
|
"file_type": f.target_file_type,
|
|
}
|
|
|
|
review_content.append(file_content)
|
|
return review_content
|
|
|
|
def _get_system_message(self) -> str:
|
|
return (
|
|
"You are a meticulous senior code reviewer. "
|
|
"You must ONLY return your findings by calling the provided function "
|
|
f"`{CustomFunctions.Code_review}` with the required JSON shape. "
|
|
"Each tuple = [line_number:int (1-based, GLOBAL across all files), line_text:str, suggestion:str]. "
|
|
"Only include lines that genuinely warrant a suggestion. "
|
|
"The input may mix languages; use each file's TYPE header and the line's syntax to tailor advice. "
|
|
"If a suggestion spans lines or files, attach it to the most relevant line and mention the related line numbers. "
|
|
"Treat 'Plain Text' as docs/logs/config and suggest clarity/safety where relevant."
|
|
)
|
|
|
|
def _get_model_name(self, model_name=None):
|
|
return model_name if model_name else self._get_model_name()
|
|
|
|
def _get_function(self, name):
|
|
"""
|
|
this is a GPT/Claude-specific function - AKA instruction for GPT
|
|
doc: https://platform.openai.com/docs/guides/function-calling
|
|
|
|
Claude uses the same tools method
|
|
"""
|
|
match name:
|
|
case CustomFunctions.Ping:
|
|
return {
|
|
"type": "function",
|
|
"name": CustomFunctions.Ping,
|
|
"description": "Always return pong",
|
|
"strict": True,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"response": {"type": "string", "enum": ["pong"]}},
|
|
"required": ["response"],
|
|
"additionalProperties": False,
|
|
},
|
|
}
|
|
case CustomFunctions.Code_review:
|
|
return {
|
|
"type": "function",
|
|
"name": CustomFunctions.Code_review,
|
|
"description": (
|
|
"Return structured code review per file. Each file uses per-file 1-based "
|
|
"line numbers. Only include lines that warrant a concrete suggestion."
|
|
),
|
|
"strict": True,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"response": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"file_name": {"type": "string"},
|
|
"review": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"line_number": {"type": "integer", "minimum": 1},
|
|
"line_code": {"type": "string"},
|
|
"suggestion": {"type": "string"},
|
|
},
|
|
"required": ["line_number", "line_code", "suggestion"],
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
},
|
|
"required": ["file_name", "review"],
|
|
"additionalProperties": False,
|
|
},
|
|
}
|
|
},
|
|
"required": ["response"],
|
|
"additionalProperties": False,
|
|
},
|
|
}
|
|
case _:
|
|
raise ValueError("Unknown function name: %s" % name)
|
|
|
|
def _get_ping_request(self) -> Request:
|
|
return PingRequest(content="ping", function_name=CustomFunctions.Ping, role="user")
|
|
|
|
def _get_review_requests(
|
|
self,
|
|
pr_diffset: DiffSet,
|
|
basic_instructions: Optional[Iterable[str]] = None,
|
|
custom_instructions: Optional[str] = None,
|
|
*args,
|
|
**kwargs,
|
|
) -> CodeReviewRequest:
|
|
review_content = self._extract_changes(pr_diffset)
|
|
|
|
system_msg = self._get_system_message()
|
|
|
|
basic_points = list(basic_instructions) if basic_instructions else self.DEFAULT_BASIC_REVIEW_POINTS
|
|
|
|
user_msg = self.get_code_review_instructions(basic_points, review_content, custom_instructions)
|
|
|
|
return CodeReviewRequest(
|
|
content=[Review(role="system", msg=system_msg), Review(role="user", msg=user_msg)],
|
|
function_name=CustomFunctions.Code_review,
|
|
)
|
|
|
|
@abstractmethod
|
|
def _get_response(self, request: PingRequest | CodeReviewRequest) -> Response:
|
|
raise NotImplementedError
|
|
|
|
def _transform(self, resp, model=None) -> Response:
|
|
return Response(model=self._get_model_name(model_name=model), message=resp)
|