215 lines
9.8 KiB
Python
215 lines
9.8 KiB
Python
from typing import Optional, List, Dict, Any, Literal, Callable, TYPE_CHECKING
|
|
from pydantic import BaseModel, Field
|
|
from enum import Enum
|
|
|
|
if TYPE_CHECKING:
|
|
from modules.datamodels.datamodelExtraction import ContentPart
|
|
|
|
# Operation Types
|
|
class OperationTypeEnum(str, Enum):
|
|
GENERAL = "general"
|
|
PLAN = "plan"
|
|
ANALYSE = "analyse"
|
|
GENERATE = "generate"
|
|
EXTRACT = "extract"
|
|
WEB_RESEARCH = "webResearch"
|
|
IMAGE_ANALYSE = "imageAnalyse"
|
|
IMAGE_GENERATE = "imageGenerate"
|
|
|
|
|
|
# Processing Modes
|
|
class ProcessingModeEnum(str, Enum):
|
|
BASIC = "basic"
|
|
ADVANCED = "advanced"
|
|
DETAILED = "detailed"
|
|
|
|
# Priority Levels
|
|
class PriorityEnum(str, Enum):
|
|
SPEED = "speed"
|
|
QUALITY = "quality"
|
|
COST = "cost"
|
|
BALANCED = "balanced"
|
|
|
|
|
|
# Model Capabilities Enumeration
|
|
class ModelCapabilitiesEnum(str, Enum):
|
|
# Text generation capabilities
|
|
TEXT_GENERATION = "text_generation"
|
|
CHAT = "chat"
|
|
REASONING = "reasoning"
|
|
ANALYSIS = "analysis"
|
|
|
|
# Image capabilities
|
|
IMAGE_ANALYSE = "imageAnalyse"
|
|
IMAGE_GENERATE = "imageGenerate"
|
|
VISION = "vision"
|
|
MULTIMODAL = "multimodal"
|
|
ART = "art"
|
|
VISUAL_CREATION = "visual_creation"
|
|
|
|
# Web capabilities
|
|
WEB_SEARCH = "web_search"
|
|
WEB_CRAWLING = "web_crawling"
|
|
CONTENT_EXTRACTION = "content_extraction"
|
|
TEXT_EXTRACTION = "text_extraction"
|
|
INFORMATION_RETRIEVAL = "information_retrieval"
|
|
URL_DISCOVERY = "url_discovery"
|
|
MAPPING = "mapping"
|
|
|
|
# Research capabilities
|
|
RESEARCH = "research"
|
|
QUESTION_ANSWERING = "question_answering"
|
|
INFORMATION_GATHERING = "information_gathering"
|
|
NEWS = "news"
|
|
CURRENT_EVENTS = "current_events"
|
|
|
|
|
|
class AiModel(BaseModel):
|
|
"""Enhanced AI model definition with dynamic capabilities."""
|
|
|
|
# Core identification
|
|
name: str = Field(description="Unique model identifier")
|
|
displayName: str = Field(description="Human-readable model name")
|
|
connectorType: str = Field(description="Type of connector (openai, anthropic, perplexity, tavily, etc.)")
|
|
|
|
# Token and context limits
|
|
maxTokens: int = Field(description="Maximum tokens this model can generate")
|
|
contextLength: int = Field(description="Maximum context length this model can handle")
|
|
|
|
# Cost information
|
|
costPer1kTokensInput: float = Field(default=0.0, description="Cost per 1000 input tokens")
|
|
costPer1kTokensOutput: float = Field(default=0.0, description="Cost per 1000 output tokens")
|
|
|
|
# Performance ratings
|
|
speedRating: int = Field(ge=1, le=10, description="Speed rating (1-10, higher = faster)")
|
|
qualityRating: int = Field(ge=1, le=10, description="Quality rating (1-10, higher = better)")
|
|
|
|
# Function reference (not serialized)
|
|
functionCall: Optional[Callable] = Field(default=None, exclude=True, description="Function to call for this model")
|
|
calculatePriceUsd: Optional[Callable] = Field(default=None, exclude=True, description="Function to calculate price in USD")
|
|
|
|
# Selection criteria
|
|
capabilities: List[ModelCapabilitiesEnum] = Field(description="List of model capabilities. See ModelCapabilitiesEnum enum for available values.")
|
|
priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Default priority for this model. See PriorityEnum for available values.")
|
|
processingMode: ProcessingModeEnum = Field(default=ProcessingModeEnum.BASIC, description="Default processing mode. See ProcessingModeEnum for available values.")
|
|
operationTypes: List[OperationTypeEnum] = Field(default=[], description="Operation types this model should avoid")
|
|
minContextLength: Optional[int] = Field(default=None, description="Minimum context length required")
|
|
isAvailable: bool = Field(default=True, description="Whether model is currently available")
|
|
|
|
# Metadata
|
|
version: Optional[str] = Field(default=None, description="Model version")
|
|
lastUpdated: Optional[str] = Field(default=None, description="Last update timestamp")
|
|
|
|
class Config:
|
|
arbitraryTypesAllowed = True # Allow Callable type
|
|
|
|
|
|
class SelectionRule(BaseModel):
|
|
"""A rule for model selection."""
|
|
name: str = Field(description="Rule name identifier")
|
|
condition: str = Field(description="Description of when this rule applies")
|
|
weight: float = Field(description="Weight for scoring (higher = more important)")
|
|
operationTypes: List[OperationTypeEnum] = Field(description="Operation types this rule applies to")
|
|
priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Priority level for this rule")
|
|
capabilities: List[ModelCapabilitiesEnum] = Field(default=[], description="Required capabilities for this rule")
|
|
minQualityRating: Optional[int] = Field(default=None, description="Minimum quality rating")
|
|
maxCost: Optional[float] = Field(default=None, description="Maximum cost threshold")
|
|
minContextLength: Optional[int] = Field(default=None, description="Minimum context length required")
|
|
|
|
|
|
class AiCallOptions(BaseModel):
|
|
"""Options for centralized AI processing with clear operation types and tags."""
|
|
operationType: OperationTypeEnum = Field(default=OperationTypeEnum.GENERAL, description="Type of operation")
|
|
priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Priority level")
|
|
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")
|
|
processingMode: ProcessingModeEnum = Field(default=ProcessingModeEnum.BASIC, description="Processing mode")
|
|
resultFormat: Optional[str] = Field(default=None, description="Expected result format: txt, json, csv, xml, etc.")
|
|
|
|
safetyMargin: float = Field(default=0.1, ge=0.0, le=0.5, description="Safety margin for token limits (0.0-0.5)")
|
|
capabilities: Optional[List[ModelCapabilitiesEnum]] = 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")
|
|
maxParts: Optional[int] = Field(default=1000, ge=1, le=1000, description="Maximum number of continuation parts to fetch")
|
|
|
|
|
|
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)
|
|
contentParts: Optional[List['ContentPart']] = None # NEW: Content parts for model-aware chunking
|
|
|
|
|
|
class AiCallResponse(BaseModel):
|
|
"""Standardized AI call response."""
|
|
|
|
content: str = Field(description="AI response content")
|
|
modelName: str = Field(description="Selected model name")
|
|
priceUsd: float = Field(default=0.0, description="Calculated price in USD")
|
|
processingTime: float = Field(default=0.0, description="Duration in seconds")
|
|
bytesSent: int = Field(default=0, description="Input data size in bytes")
|
|
bytesReceived: int = Field(default=0, description="Output data size in bytes")
|
|
errorCount: int = Field(default=0, description="0 for success, 1+ for errors")
|
|
|
|
|
|
class EnhancedAiCallOptions(AiCallOptions):
|
|
"""Enhanced options for improved document processing with chunk mapping."""
|
|
|
|
# Parallel processing
|
|
enableParallelProcessing: bool = Field(
|
|
default=True,
|
|
description="Enable parallel processing of chunks"
|
|
)
|
|
maxConcurrentChunks: int = Field(
|
|
default=5,
|
|
ge=1,
|
|
le=20,
|
|
description="Maximum number of chunks to process concurrently"
|
|
)
|
|
|
|
# Chunk mapping
|
|
preserveChunkMetadata: bool = Field(
|
|
default=True,
|
|
description="Preserve chunk metadata during processing"
|
|
)
|
|
chunkSeparator: str = Field(
|
|
default="\n\n---\n\n",
|
|
description="Separator between chunks in merged output"
|
|
)
|
|
|
|
|
|
class AiModelCall(BaseModel):
|
|
"""Standardized input for AI model calls."""
|
|
|
|
messages: List[Dict[str, Any]] = Field(description="Messages in OpenAI format (role, content)")
|
|
model: Optional[AiModel] = Field(default=None, description="The AI model being called")
|
|
options: Dict[str, Any] = Field(default_factory=dict, description="Additional model-specific options")
|
|
|
|
class Config:
|
|
arbitraryTypesAllowed = True
|
|
|
|
|
|
class AiModelResponse(BaseModel):
|
|
"""Standardized output from AI model calls."""
|
|
|
|
content: str = Field(description="The AI response content")
|
|
success: bool = Field(default=True, description="Whether the call was successful")
|
|
error: Optional[str] = Field(default=None, description="Error message if success=False")
|
|
|
|
# Optional metadata that models can include
|
|
modelId: Optional[str] = Field(default=None, description="Model identifier used")
|
|
processingTime: Optional[float] = Field(default=None, description="Processing time in seconds")
|
|
tokensUsed: Optional[Dict[str, int]] = Field(default=None, description="Token usage (input, output, total)")
|
|
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional model-specific metadata")
|
|
|
|
class Config:
|
|
arbitraryTypesAllowed = True
|
|
|