models.py

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


class TaskType(str, Enum):
    """Types of identity verification tasks the environment supports."""
    DOCUMENT_VERIFICATION = "document_verification"
    CREDENTIAL_AUTHENTICATION = "credential_authentication"
    PROFILE_MANAGEMENT = "profile_management"
    FRAUD_DETECTION = "fraud_detection"


class DocumentType(str, Enum):
    """Types of identity documents."""
    PASSPORT = "passport"
    DRIVERS_LICENSE = "drivers_license"
    NATIONAL_ID = "national_id"
    BIRTH_CERTIFICATE = "birth_certificate"


class IdentityDocument(BaseModel):
    """Identity document data for verification."""
    document_type: DocumentType = Field(description="Type of document")
    document_number: str = Field(description="Document identifier number")
    full_name: str = Field(description="Full name as on document")
    date_of_birth: str = Field(description="Date of birth (YYYY-MM-DD)")
    expiry_date: Optional[str] = Field(default=None, description="Document expiry date")
    issuing_country: str = Field(description="Country that issued the document")
    additional_fields: dict[str, Any] = Field(default_factory=dict, description="Additional document data")


class UserCredentials(BaseModel):
    """User credentials for authentication."""
    username: str = Field(description="Username or email")
    password: str = Field(description="Password")
    two_factor_code: Optional[str] = Field(default=None, description="2FA code if enabled")


class UserProfile(BaseModel):
    """User profile information."""
    user_id: str = Field(description="Unique user identifier")
    full_name: str = Field(description="Full name")
    email: str = Field(description="Email address")
    phone: Optional[str] = Field(default=None, description="Phone number")
    date_of_birth: Optional[str] = Field(default=None, description="Date of birth")
    address: Optional[str] = Field(default=None, description="Physical address")
    attributes: dict[str, Any] = Field(default_factory=dict, description="Additional attributes")


class VerificationResult(BaseModel):
    """Result of identity verification."""
    verified: bool = Field(description="Whether verification passed")
    confidence: float = Field(default=0.0, description="Confidence score 0-1")
    issues: list[str] = Field(default_factory=list, description="List of verification issues")
    checks_performed: dict[str, bool] = Field(default_factory=dict, description="Individual check results")


class CoreIdentityObservation(BaseModel):
    """Observation returned by the environment after reset or step."""
    task_id: str = Field(description="Unique identifier for the task")
    task_type: TaskType = Field(description="Type of identity task")
    task_name: str = Field(description="Human-readable name of the task")
    task_description: str = Field(description="Description of what needs to be done")
    difficulty: str = Field(default="medium", description="Difficulty level: easy, medium, hard")
    
    document: Optional[IdentityDocument] = Field(default=None, description="Document to verify")
    credentials: Optional[UserCredentials] = Field(default=None, description="Credentials to authenticate")
    profile: Optional[UserProfile] = Field(default=None, description="User profile to manage")
    
    expected_verification: Optional[bool] = Field(default=None, description="Expected verification result")
    challenge_data: dict[str, Any] = Field(default_factory=dict, description="Additional challenge data")
    max_steps: int = Field(default=10, description="Maximum allowed steps")


class CoreIdentityAction(BaseModel):
    """Action to perform in the environment."""
    verification: VerificationResult = Field(description="The verification result from the agent")
    submit: bool = Field(default=False, description="Whether to submit the final result")


class CoreIdentityReward(BaseModel):
    """Reward components for the environment."""
    accuracy: float = Field(default=0.0, description="Accuracy of verification")
    completeness: float = Field(default=0.0, description="Completeness of checks performed")
    efficiency: float = Field(default=0.0, description="Efficiency bonus for fewer steps")
    total: float = Field(default=0.0, description="Total reward value")