models.py

from enum import Enum
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field
from openenv.core.env_server.types import Action, Observation

class ActionType(str, Enum):
    enrich_lead = "enrich_lead"
    check_crm = "check_crm"
    update_lead_score = "update_lead_score"
    route_to_rep = "route_to_rep"
    merge_with_account = "merge_with_account"
    flag_reengagement = "flag_reengagement"
    disqualify = "disqualify"
    add_to_nurture = "add_to_nurture"
    request_more_info = "request_more_info"
    escalate_to_manager = "escalate_to_manager"

class RevOpsAction(Action):
    action_type: ActionType
    score: Optional[int] = Field(None, description="Score to assign (0-100) used with update_lead_score")
    rep_id: Optional[str] = Field(None, description="ID of the rep to route to")
    account_id: Optional[str] = Field(None, description="ID of the account to merge with")
    opportunity_id: Optional[str] = Field(None, description="ID of the related opportunity")
    disqualification_reason: Optional[str] = Field(None, description="Reason for disqualification")
    note: Optional[str] = Field(None, description="Additional notes or context")

class LeadData(BaseModel):
    id: str
    name: Optional[str] = None
    title: Optional[str] = None
    email: Optional[str] = None
    company: Optional[str] = None
    message: Optional[str] = None
    source: Optional[str] = None
    score: Optional[int] = None

class EnrichmentData(BaseModel):
    enriched: bool = False
    company_size: Optional[str] = None
    revenue: Optional[str] = None
    industry: Optional[str] = None
    region: Optional[str] = None
    tier: Optional[str] = None

class OpportunityHistory(BaseModel):
    opportunity_id: str
    status: str
    amount: int
    assigned_rep_id: str
    closed_date: Optional[str] = None

class CRMState(BaseModel):
    checked: bool = False
    existing_accounts: List[Dict[str, Any]] = []
    opportunities: List[OpportunityHistory] = []
    
class Rep(BaseModel):
    rep_id: str
    name: str
    segment: str
    territory: str
    capacity: str

class RevOpsObservation(Observation):
    task_id: str = "task_easy"
    lead: LeadData = Field(default_factory=lambda: LeadData(id="dummy"))
    enrichment: EnrichmentData = Field(default_factory=EnrichmentData)
    crm: CRMState = Field(default_factory=CRMState)
    reps: List[Rep] = Field(default_factory=list)
    icp_criteria: str = ""
    sla_time_remaining_minutes: int = 60
    last_action_feedback: Optional[str] = None

class RevOpsReward(BaseModel):
    value: float
    message: str

class RevOpsState(BaseModel):
    observation: RevOpsObservation
    action_history: List[RevOpsAction]
    accumulated_reward: float
    is_done: bool
    grader_score: Optional[float] = None