158 lines
6.8 KiB
Python
158 lines
6.8 KiB
Python
"""
|
|
Configuration for dynamic model selection rules.
|
|
This makes model selection configurable rather than hardcoded.
|
|
"""
|
|
|
|
from typing import Dict, List, Any
|
|
from modules.datamodels.datamodelAi import OperationTypeEnum, ModelCapabilitiesEnum, PriorityEnum, SelectionRule
|
|
|
|
|
|
class ModelSelectionConfig:
|
|
"""Configuration for model selection rules."""
|
|
|
|
def __init__(self):
|
|
self.rules = self._loadDefaultRules()
|
|
self.fallbackModels = self._loadFallbackModels()
|
|
|
|
def _loadDefaultRules(self) -> List[SelectionRule]:
|
|
"""Load default selection rules."""
|
|
return [
|
|
# High quality for planning and analysis
|
|
SelectionRule(
|
|
name="highQualityAnalysis",
|
|
condition="Planning or analysis operations requiring high quality",
|
|
weight=10.0,
|
|
operationTypes=[OperationTypeEnum.PLAN, OperationTypeEnum.ANALYSE],
|
|
priority=PriorityEnum.QUALITY,
|
|
capabilities=[ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.REASONING, ModelCapabilitiesEnum.ANALYSIS],
|
|
minQualityRating=8
|
|
),
|
|
|
|
# Fast processing for basic operations
|
|
SelectionRule(
|
|
name="fastBasicProcessing",
|
|
condition="Basic operations requiring speed",
|
|
weight=8.0,
|
|
operationTypes=[OperationTypeEnum.GENERAL],
|
|
priority=PriorityEnum.SPEED,
|
|
capabilities=[ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.CHAT],
|
|
minQualityRating=5
|
|
),
|
|
|
|
# Cost-effective for high-volume operations
|
|
SelectionRule(
|
|
name="costEffectiveProcessing",
|
|
condition="High-volume operations where cost matters",
|
|
weight=7.0,
|
|
operationTypes=[OperationTypeEnum.GENERAL, OperationTypeEnum.GENERATE],
|
|
priority=PriorityEnum.COST,
|
|
capabilities=[ModelCapabilitiesEnum.TEXT_GENERATION],
|
|
maxCost=0.01 # $0.01 per 1k tokens
|
|
),
|
|
|
|
# Image analysis specific
|
|
SelectionRule(
|
|
name="imageAnalyse",
|
|
condition="Image analysis operations",
|
|
weight=10.0,
|
|
operationTypes=[OperationTypeEnum.IMAGE_ANALYSE],
|
|
priority=PriorityEnum.QUALITY,
|
|
capabilities=[ModelCapabilitiesEnum.VISION, ModelCapabilitiesEnum.MULTIMODAL],
|
|
minQualityRating=8
|
|
),
|
|
|
|
# Web research specific
|
|
SelectionRule(
|
|
name="webResearch",
|
|
condition="Web research operations",
|
|
weight=9.0,
|
|
operationTypes=[OperationTypeEnum.WEB_RESEARCH],
|
|
priority=PriorityEnum.BALANCED,
|
|
capabilities=[ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.ANALYSIS, ModelCapabilitiesEnum.WEB_SEARCH],
|
|
minQualityRating=7
|
|
),
|
|
|
|
# Large context requirements
|
|
SelectionRule(
|
|
name="largeContext",
|
|
condition="Operations requiring large context",
|
|
weight=8.0,
|
|
operationTypes=[OperationTypeEnum.GENERAL, OperationTypeEnum.ANALYSE],
|
|
priority=PriorityEnum.BALANCED,
|
|
capabilities=[ModelCapabilitiesEnum.TEXT_GENERATION],
|
|
minContextLength=100000 # 100k tokens
|
|
)
|
|
]
|
|
|
|
def _loadFallbackModels(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Load fallback model selection criteria."""
|
|
return {
|
|
OperationTypeEnum.GENERAL: {
|
|
"priorityOrder": ["speed", "quality", "cost"],
|
|
"operationTypes": [ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.CHAT],
|
|
"minQualityRating": 5,
|
|
"maxCostPer1k": 0.01
|
|
},
|
|
OperationTypeEnum.IMAGE_ANALYSE: {
|
|
"priorityOrder": ["quality", "speed"],
|
|
"operationTypes": [ModelCapabilitiesEnum.VISION, ModelCapabilitiesEnum.MULTIMODAL],
|
|
"minQualityRating": 8,
|
|
"maxCostPer1k": 0.1
|
|
},
|
|
OperationTypeEnum.IMAGE_GENERATE: {
|
|
"priorityOrder": ["quality", "speed"],
|
|
"operationTypes": [ModelCapabilitiesEnum.IMAGE_GENERATE, ModelCapabilitiesEnum.ART, ModelCapabilitiesEnum.VISUAL_CREATION],
|
|
"minQualityRating": 8,
|
|
"maxCostPer1k": 0.1
|
|
},
|
|
OperationTypeEnum.WEB_RESEARCH: {
|
|
"priorityOrder": ["quality", "speed", "cost"],
|
|
"operationTypes": [ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.ANALYSIS],
|
|
"preferredTags": [ModelCapabilitiesEnum.WEB_SEARCH],
|
|
"minQualityRating": 7,
|
|
"maxCostPer1k": 0.02
|
|
},
|
|
OperationTypeEnum.PLAN: {
|
|
"priorityOrder": ["quality", "speed"],
|
|
"operationTypes": [ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.REASONING, ModelCapabilitiesEnum.ANALYSIS],
|
|
"preferredTags": [PriorityEnum.QUALITY],
|
|
"minQualityRating": 8,
|
|
"maxCostPer1k": 0.1
|
|
},
|
|
OperationTypeEnum.ANALYSE: {
|
|
"priorityOrder": ["quality", "speed"],
|
|
"operationTypes": [ModelCapabilitiesEnum.TEXT_GENERATION, ModelCapabilitiesEnum.ANALYSIS, ModelCapabilitiesEnum.REASONING],
|
|
"preferredTags": [PriorityEnum.QUALITY],
|
|
"minQualityRating": 8,
|
|
"maxCostPer1k": 0.1
|
|
}
|
|
}
|
|
|
|
def getRulesForOperation(self, operationType: str) -> List[SelectionRule]:
|
|
"""Get rules that apply to a specific operation type."""
|
|
return [rule for rule in self.rules if operationType in rule.operationTypes]
|
|
|
|
def getFallbackCriteria(self, operationType: str) -> Dict[str, Any]:
|
|
"""Get fallback selection criteria for a specific operation type."""
|
|
return self.fallbackModels.get(operationType, self.fallbackModels[OperationTypeEnum.GENERAL])
|
|
|
|
def addRule(self, rule: SelectionRule):
|
|
"""Add a new selection rule."""
|
|
self.rules.append(rule)
|
|
|
|
def removeRule(self, ruleName: str):
|
|
"""Remove a selection rule by name."""
|
|
self.rules = [rule for rule in self.rules if rule.name != ruleName]
|
|
|
|
def updateRule(self, ruleName: str, **kwargs):
|
|
"""Update an existing rule."""
|
|
for rule in self.rules:
|
|
if rule.name == ruleName:
|
|
for key, value in kwargs.items():
|
|
if hasattr(rule, key):
|
|
setattr(rule, key, value)
|
|
break
|
|
|
|
|
|
# Global configuration instance
|
|
model_selection_config = ModelSelectionConfig()
|