app/models/ai_models.py

from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Any
from enum import Enum

class AnalysisFocus(str, Enum):
    FUND_SELECTION = "Fund Selection"
    RISK_ASSESSMENT = "Risk Assessment"
    GOAL_ALIGNMENT = "Goal Alignment"
    TAX_OPTIMIZATION = "Tax Optimization"
    REBALANCING = "Rebalancing"

class MarketView(str, Enum):
    BULLISH = "Bullish"
    NEUTRAL = "Neutral"
    BEARISH = "Bearish"
    VOLATILE = "Volatile"

class InvestmentHorizon(str, Enum):
    SHORT_TERM = "Short Term (1-3 years)"
    MEDIUM_TERM = "Medium Term (3-7 years)"
    LONG_TERM = "Long Term (7+ years)"

class AIAnalysisRequest(BaseModel):
    client_profile: Dict[str, Any]
    portfolio_data: Dict[str, Any]
    goals_data: Dict[str, Any]
    analysis_focus: List[AnalysisFocus]
    market_conditions: MarketView
    investment_horizon: InvestmentHorizon

class AISwarmAgent(BaseModel):
    agent_name: str
    content: Optional[str] = None

class AIAnalysisResponse(BaseModel):
    success: bool
    error: Optional[str] = None
    output: Optional[List[AISwarmAgent]] = None