113 lines
3.9 KiB
Python
113 lines
3.9 KiB
Python
import logging
|
|
import math
|
|
from typing import Optional, Iterable
|
|
|
|
from anthropic import Anthropic
|
|
from anthropic.types import ToolParam
|
|
|
|
from rhodecode.apps.ai_agents.ai_settings import AISettings, AIModelName, ClaudeVersion
|
|
from rhodecode.apps.ai_agents.models.base import (
|
|
AIServiceBase,
|
|
Response,
|
|
AIServiceError,
|
|
PingRequest,
|
|
CodeReviewRequest,
|
|
Review,
|
|
CustomFunctions,
|
|
FILE,
|
|
)
|
|
from rhodecode.lib.codeblocks import DiffSet
|
|
|
|
|
|
class ClaudeService(AIServiceBase):
|
|
def __init__(self, model_settings: AISettings):
|
|
super().__init__(model_settings)
|
|
self._client = Anthropic(api_key=self.model_settings.api_key)
|
|
self.log = logging.getLogger(ClaudeService.__name__)
|
|
|
|
def _get_response(self, request: PingRequest | CodeReviewRequest):
|
|
# maximum number of output tokens Claude can generate before stopping (does not include input tokens)
|
|
max_tokens = self._calculate_max_tokens(request)
|
|
self.log.debug("estimated max output tokens: %s", max_tokens)
|
|
|
|
_input = self._get_response_input(request)
|
|
|
|
model_full_name = self._get_current_model_version()
|
|
|
|
function = self._adapt_function(self._get_function(request.function_name))
|
|
|
|
return self._client.messages.create(
|
|
model=model_full_name,
|
|
messages=_input,
|
|
max_tokens=max_tokens,
|
|
tools=[
|
|
function,
|
|
],
|
|
)
|
|
|
|
def _adapt_function(self, function_dict: dict) -> ToolParam:
|
|
return ToolParam(
|
|
name=function_dict["name"],
|
|
description=function_dict["description"],
|
|
input_schema=function_dict["parameters"],
|
|
)
|
|
|
|
def _transform(self, resp, model=None) -> Response:
|
|
for r in resp.content:
|
|
self.log.debug("response: %s", r)
|
|
if r.type == "tool_use" and r.name in [CustomFunctions.Code_review, CustomFunctions.Ping]:
|
|
return Response(model=resp.model, message=r.input)
|
|
raise ValueError("Response has incorrect return type, can't parses it.")
|
|
|
|
def _get_current_model_version(self):
|
|
models = self._client.models.list()
|
|
semi_key = "%s-%s" % (
|
|
self.model_settings.model_name.lower().strip(),
|
|
self.model_settings.model_version.lower().strip(),
|
|
)
|
|
for model in models:
|
|
if semi_key in model.id:
|
|
return model.id
|
|
raise AIServiceError("Unknown model version: %s" % semi_key)
|
|
|
|
def _calculate_max_tokens(self, request: PingRequest | CodeReviewRequest):
|
|
default = 256
|
|
|
|
if isinstance(request, CodeReviewRequest):
|
|
user_msg: Review = next(filter(lambda r: FILE in r.msg, request.content), None)
|
|
if user_msg:
|
|
user_msg_code_lines = user_msg.msg.count(FILE)
|
|
return self._estimate_max_tokens(user_msg_code_lines)
|
|
|
|
return default
|
|
|
|
def _estimate_max_tokens(
|
|
self,
|
|
lines: int,
|
|
estimated_sentences_per_line: float = 3.0,
|
|
average_words_per_english_sentence: float = 20.0,
|
|
rough_heuristic_tokens_per_word: float = 4,
|
|
safety_buffer_percent: float = 0.15,
|
|
) -> int:
|
|
"""
|
|
Estimate a safe `max_tokens` (output only) for code review.
|
|
"""
|
|
raw = (
|
|
lines * estimated_sentences_per_line * average_words_per_english_sentence * rough_heuristic_tokens_per_word
|
|
)
|
|
return math.ceil(raw * (1.0 + safety_buffer_percent))
|
|
|
|
def _get_review_requests(
|
|
self,
|
|
pr_diffset: DiffSet,
|
|
basic_instructions: Optional[Iterable[str]] = None,
|
|
custom_instructions: Optional[str] = None,
|
|
*args,
|
|
**kwargs,
|
|
) -> CodeReviewRequest:
|
|
# Claude has only 2 roles: user, assistant
|
|
requests = super()._get_review_requests(pr_diffset, basic_instructions, custom_instructions, *args, **kwargs)
|
|
for r in requests.content:
|
|
r.role = "user"
|
|
|
|
return requests
|