69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
from abc import abstractmethod
|
|
from dataclasses import dataclass
|
|
from functools import wraps
|
|
|
|
from rhodecode.apps.ai_agents.ai_settings import AISettings
|
|
|
|
|
|
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:
|
|
raise AIServiceError(e) from e
|
|
|
|
return wrapper
|
|
|
|
|
|
class AIServiceBase:
|
|
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)
|
|
|
|
@abstractmethod
|
|
def _get_ping_request(self) -> Request:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def code_review(self, code):
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def _get_response(self, request: Request):
|
|
raise NotImplementedError
|
|
|
|
def _transform(self, resp, model=None) -> Response:
|
|
return Response(model=model if model else self.model_settings.model_name.value, message=resp)
|