136 lines
5.6 KiB
Python
136 lines
5.6 KiB
Python
from typing import Optional, List, Dict, Any, Literal
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
# Operation Types
|
|
class OperationType:
|
|
GENERAL = "general"
|
|
GENERATE_PLAN = "generate_plan"
|
|
ANALYSE_CONTENT = "analyse_content"
|
|
GENERATE_CONTENT = "generate_content"
|
|
WEB_RESEARCH = "web_research"
|
|
IMAGE_ANALYSIS = "image_analysis"
|
|
IMAGE_GENERATION = "image_generation"
|
|
|
|
|
|
# Processing Modes
|
|
class ProcessingMode:
|
|
BASIC = "basic"
|
|
ADVANCED = "advanced"
|
|
DETAILED = "detailed"
|
|
|
|
|
|
# Priority Levels
|
|
class Priority:
|
|
SPEED = "speed"
|
|
QUALITY = "quality"
|
|
COST = "cost"
|
|
BALANCED = "balanced"
|
|
|
|
|
|
# Model Tags
|
|
class ModelTags:
|
|
# Core capabilities
|
|
TEXT = "text"
|
|
CHAT = "chat"
|
|
REASONING = "reasoning"
|
|
ANALYSIS = "analysis"
|
|
IMAGE = "image"
|
|
VISION = "vision"
|
|
MULTIMODAL = "multimodal"
|
|
WEB = "web"
|
|
SEARCH = "search"
|
|
CRAWL = "crawl"
|
|
EXTRACT = "extract"
|
|
CONTENT = "content"
|
|
INFORMATION = "information"
|
|
|
|
# Quality indicators
|
|
HIGH_QUALITY = "high_quality"
|
|
FAST = "fast"
|
|
COST_EFFECTIVE = "cost_effective"
|
|
GENERAL = "general"
|
|
|
|
# Specialized capabilities
|
|
IMAGE_GENERATION = "image_generation"
|
|
ART = "art"
|
|
VISUAL = "visual"
|
|
VARIATIONS = "variations"
|
|
API = "api"
|
|
INFO = "info"
|
|
MODELS = "models"
|
|
|
|
|
|
# Operation Type to Required Tags Mapping
|
|
OPERATION_TAG_MAPPING = {
|
|
OperationType.GENERAL: [ModelTags.TEXT, ModelTags.CHAT, ModelTags.REASONING],
|
|
OperationType.GENERATE_PLAN: [ModelTags.TEXT, ModelTags.REASONING, ModelTags.ANALYSIS],
|
|
OperationType.ANALYSE_CONTENT: [ModelTags.TEXT, ModelTags.ANALYSIS, ModelTags.REASONING],
|
|
OperationType.GENERATE_CONTENT: [ModelTags.TEXT, ModelTags.CHAT, ModelTags.REASONING],
|
|
OperationType.WEB_RESEARCH: [ModelTags.TEXT, ModelTags.ANALYSIS, ModelTags.REASONING],
|
|
OperationType.IMAGE_ANALYSIS: [ModelTags.IMAGE, ModelTags.VISION, ModelTags.MULTIMODAL],
|
|
OperationType.IMAGE_GENERATION: [ModelTags.IMAGE_GENERATION, ModelTags.ART, ModelTags.VISUAL],
|
|
}
|
|
|
|
|
|
# Processing Mode to Priority Mapping
|
|
PROCESSING_MODE_PRIORITY_MAPPING = {
|
|
ProcessingMode.BASIC: Priority.SPEED,
|
|
ProcessingMode.ADVANCED: Priority.BALANCED,
|
|
ProcessingMode.DETAILED: Priority.QUALITY,
|
|
}
|
|
|
|
|
|
class ModelCapabilities(BaseModel):
|
|
"""Model capabilities and characteristics for dynamic selection."""
|
|
|
|
name: str = Field(description="Model name/identifier")
|
|
maxTokens: int = Field(description="Maximum token limit for this model")
|
|
capabilities: List[str] = Field(description="List of capabilities: text, image, vision, reasoning, analysis, etc.")
|
|
costPerToken: float = Field(default=0.0, description="Cost per token (if available)")
|
|
processingTime: float = Field(default=1.0, description="Average processing time multiplier")
|
|
isAvailable: bool = Field(default=True, description="Whether model is currently available")
|
|
|
|
|
|
class AiCallOptions(BaseModel):
|
|
"""Options for centralized AI processing with clear operation types and tags."""
|
|
|
|
operationType: str = Field(default="general", description="Type of operation: general, generate_plan, analyse_content, generate_content, web_research")
|
|
priority: str = Field(default="balanced", description="speed|quality|cost|balanced")
|
|
compressPrompt: bool = Field(default=True, description="Whether to compress the prompt")
|
|
compressContext: bool = Field(default=True, description="If False: process each chunk; If True: summarize and work on summary")
|
|
processDocumentsIndividually: bool = Field(default=True, description="If True, process each document separately; else pool docs")
|
|
maxContextBytes: Optional[int] = Field(default=None, description="Hard cap for extracted context size passed to the model")
|
|
maxCost: Optional[float] = Field(default=None, description="Max cost budget")
|
|
maxProcessingTime: Optional[int] = Field(default=None, description="Max processing time in seconds")
|
|
requiredTags: Optional[List[str]] = Field(default=None, description="Required model tags for selection")
|
|
processingMode: str = Field(default="basic", description="Processing mode: basic, advanced, detailed")
|
|
resultFormat: Optional[str] = Field(default=None, description="Expected result format: txt, json, csv, xml, etc.")
|
|
|
|
# New fields for dynamic strategy
|
|
callType: Literal["planning", "text"] = Field(default="text", description="Call type: planning or text")
|
|
safetyMargin: float = Field(default=0.1, ge=0.0, le=0.5, description="Safety margin for token limits (0.0-0.5)")
|
|
modelCapabilities: Optional[List[str]] = Field(default=None, description="Required model capabilities for filtering")
|
|
|
|
# Model generation parameters
|
|
temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0, description="Temperature for response generation (0.0-2.0, lower = more consistent)")
|
|
maxTokens: Optional[int] = Field(default=None, ge=1, le=32000, description="Maximum tokens in response")
|
|
|
|
|
|
class AiCallRequest(BaseModel):
|
|
"""Centralized AI call request payload for interface use."""
|
|
|
|
prompt: str = Field(description="The user prompt")
|
|
context: Optional[str] = Field(default=None, description="Optional external context (e.g., extracted docs)")
|
|
options: AiCallOptions = Field(default_factory=AiCallOptions)
|
|
|
|
|
|
class AiCallResponse(BaseModel):
|
|
"""Standardized AI call response."""
|
|
|
|
content: str = Field(description="AI response content")
|
|
modelName: str = Field(description="Selected model name")
|
|
usedTokens: Optional[int] = Field(default=None, description="Estimated used tokens")
|
|
costEstimate: Optional[float] = Field(default=None, description="Estimated cost of the call")
|
|
|
|
|