diff --git a/rhodecode/apps/ai_agents/ai_service.py b/rhodecode/apps/ai_agents/ai_service.py index 2b1dc4f5..21364beb 100644 --- a/rhodecode/apps/ai_agents/ai_service.py +++ b/rhodecode/apps/ai_agents/ai_service.py @@ -12,6 +12,13 @@ class AIServiceError(Exception): pass +@dataclass +class Response: + model: str + message: str + error: bool = False + + def wrap_ai_exceptions(f): @wraps(f) def wrapper(*args, **kwargs): @@ -49,8 +56,8 @@ class AIServiceBase: # TODO: maybe just response pass - def _transform(self, resp): - return resp + def _transform(self, resp, model=None): + return Response(model=model if model else self.model_settings.model_name.value, message=resp) @wrap_ai_exceptions @@ -68,13 +75,6 @@ def get_gpt_service(api_key: str, model_name=AIModelName.GPT, version=GPTVersion raise ValueError("Unknown model name: %s" % model_name) -@dataclass -class Response: - model: str - message: str - error: bool = False - - class GPTService(AIServiceBase): def __init__(self, model_settings: AISettings): super().__init__(model_settings) @@ -94,11 +94,15 @@ class GPTService(AIServiceBase): def _transform(self, resp): for item in resp.output: if item.type == "function_call" and item.name == "ping_pong": - return json.loads(item.arguments) + return super()._transform( + resp=json.loads(item.arguments), + model=resp.model, + ) def _get_function(self, name): """ - this is a GPT-specific function - AKA fine-tuning prompt + this is a GPT-specific function - AKA instruction for GPT + doc: https://platform.openai.com/docs/guides/function-calling """ match name: case "ping":