Use selection_model field from autoselect configuration

- Add selection_model field to AutoselectConfig model
- Update _get_model_selection to use autoselect_config.selection_model instead of hardcoded 'general'
- Update handle_autoselect_request to log selection_model from config
- Update handle_autoselect_streaming_request to log selection_model from config
- Allows flexible configuration of which rotation to use for model selection
parent dde30272
...@@ -46,6 +46,7 @@ class AutoselectModelInfo(BaseModel): ...@@ -46,6 +46,7 @@ class AutoselectModelInfo(BaseModel):
class AutoselectConfig(BaseModel): class AutoselectConfig(BaseModel):
model_name: str model_name: str
description: str description: str
selection_model: str
fallback: str fallback: str
available_models: List[AutoselectModelInfo] available_models: List[AutoselectModelInfo]
......
...@@ -475,12 +475,12 @@ class AutoselectHandler: ...@@ -475,12 +475,12 @@ class AutoselectHandler:
return match.group(1).strip() return match.group(1).strip()
return None return None
async def _get_model_selection(self, prompt: str) -> str: async def _get_model_selection(self, prompt: str, autoselect_config) -> str:
"""Send the autoselect prompt to a model and get the selection""" """Send the autoselect prompt to a model and get the selection"""
import logging import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
logger.info(f"=== AUTOSELECT MODEL SELECTION START ===") logger.info(f"=== AUTOSELECT MODEL SELECTION START ===")
logger.info(f"Using 'general' rotation for model selection") logger.info(f"Using '{autoselect_config.selection_model}' rotation for model selection")
# Use the first available provider/model for the selection # Use the first available provider/model for the selection
# This is a simple implementation - could be enhanced to use a specific selection model # This is a simple implementation - could be enhanced to use a specific selection model
...@@ -499,10 +499,10 @@ class AutoselectHandler: ...@@ -499,10 +499,10 @@ class AutoselectHandler:
logger.info(f" Max tokens: 100 (short response expected)") logger.info(f" Max tokens: 100 (short response expected)")
logger.info(f" Stream: False") logger.info(f" Stream: False")
# Use the fallback rotation for the selection # Use the configured selection rotation for the selection
try: try:
logger.info(f"Sending selection request to rotation handler...") logger.info(f"Sending selection request to rotation handler...")
response = await rotation_handler.handle_rotation_request("general", selection_request) response = await rotation_handler.handle_rotation_request(autoselect_config.selection_model, selection_request)
logger.info(f"Selection response received") logger.info(f"Selection response received")
content = response.get('choices', [{}])[0].get('message', {}).get('content', '') content = response.get('choices', [{}])[0].get('message', {}).get('content', '')
...@@ -542,6 +542,7 @@ class AutoselectHandler: ...@@ -542,6 +542,7 @@ class AutoselectHandler:
logger.info(f"Available models for selection: {len(autoselect_config.available_models)}") logger.info(f"Available models for selection: {len(autoselect_config.available_models)}")
for model_info in autoselect_config.available_models: for model_info in autoselect_config.available_models:
logger.info(f" - {model_info.model_id}: {model_info.description}") logger.info(f" - {model_info.model_id}: {model_info.description}")
logger.info(f"Selection model: {autoselect_config.selection_model}")
logger.info(f"Fallback model: {autoselect_config.fallback}") logger.info(f"Fallback model: {autoselect_config.fallback}")
# Extract the user prompt from the request # Extract the user prompt from the request
...@@ -572,7 +573,7 @@ class AutoselectHandler: ...@@ -572,7 +573,7 @@ class AutoselectHandler:
# Get the model selection # Get the model selection
logger.info(f"Requesting model selection from AI...") logger.info(f"Requesting model selection from AI...")
selected_model_id = await self._get_model_selection(autoselect_prompt) selected_model_id = await self._get_model_selection(autoselect_prompt, autoselect_config)
# Validate the selected model # Validate the selected model
logger.info(f"=== MODEL VALIDATION ===") logger.info(f"=== MODEL VALIDATION ===")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment