unrhodecode/rhodecode/apps/ai_agents/models/base.py
2025-08-29 09:01:40 +02:00

131 lines
5.5 KiB
Python

from abc import abstractmethod
from dataclasses import dataclass
from functools import wraps
from rhodecode.apps.ai_agents.ai_settings import AISettings
from rhodecode.lib.codeblocks import DiffSet
class AIServiceError(Exception):
pass
@dataclass
class Response:
model: str
message: str
error: bool = False
@dataclass
class Request:
content: str
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))
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"
"- Use the function call to return `return_code_review`: a list of tuples "
" [line_number:int (1-based, GLOBAL), 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, str]]) -> str:
"""
Render code with explicit 1-based line numbers so the model can reference them precisely.
We keep original content for returning 'line_text' in suggestions.
"""
parts: list[str] = []
for content in review_content:
lines = content["changes"].splitlines()
width = max(2, len(str(len(lines)))) if lines else 2
parts.append("FILE: %s\nTYPE: %s" % (content["file_name"], content["file_type"]))
if lines:
for i, line in enumerate(lines, start=1):
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_model_name(self, model_name=None):
return model_name if model_name else self._get_model_name()
@abstractmethod
def _get_ping_request(self) -> Request:
raise NotImplementedError
@abstractmethod
def _get_review_requests(self, pr_diffset: DiffSet, *args, **kwargs) -> list[Request]:
raise NotImplementedError
@abstractmethod
def _get_response(self, request: Request | list[Request]) -> Response:
raise NotImplementedError
def _transform(self, resp, model=None) -> Response:
return Response(model=self._get_model_name(model_name=model), message=resp)