Support multiple model types for different reasoning tasks

This commit is contained in:
mrT23
2025-04-27 08:50:03 +03:00
parent 60a887ffe1
commit f53bd524c5
5 changed files with 27 additions and 17 deletions

View File

@ -12,7 +12,7 @@ from pr_agent.algo.git_patch_processing import (
from pr_agent.algo.language_handler import sort_files_by_main_languages
from pr_agent.algo.token_handler import TokenHandler
from pr_agent.algo.types import EDIT_TYPE, FilePatchInfo
from pr_agent.algo.utils import ModelType, clip_tokens, get_max_tokens, get_weak_model
from pr_agent.algo.utils import ModelType, clip_tokens, get_max_tokens, get_model
from pr_agent.config_loader import get_settings
from pr_agent.git_providers.git_provider import GitProvider
from pr_agent.log import get_logger
@ -339,7 +339,9 @@ async def retry_with_fallback_models(f: Callable, model_type: ModelType = ModelT
def _get_all_models(model_type: ModelType = ModelType.REGULAR) -> List[str]:
if model_type == ModelType.WEAK:
model = get_weak_model()
model = get_model('model_weak')
elif model_type == ModelType.REASONING:
model = get_model('model_reasoning')
else:
model = get_settings().config.model
fallback_models = get_settings().config.fallback_models

View File

@ -30,12 +30,13 @@ from pr_agent.config_loader import get_settings, global_settings
from pr_agent.log import get_logger
def get_weak_model() -> str:
if get_settings().get("config.model_weak"):
def get_model(model_type: str = "model_weak") -> str:
if model_type == "model_weak" and get_settings().get("config.model_weak"):
return get_settings().config.model_weak
elif model_type == "model_reasoning" and get_settings().get("config.model_reasoning"):
return get_settings().config.model_reasoning
return get_settings().config.model
class Range(BaseModel):
line_start: int # should be 0-indexed
line_end: int
@ -45,6 +46,7 @@ class Range(BaseModel):
class ModelType(str, Enum):
REGULAR = "regular"
WEAK = "weak"
REASONING = "reasoning"
class PRReviewHeader(str, Enum):
REGULAR = "## PR Reviewer Guide"