ai-legal-assistant/backend/app/schemas/phase2.py

"""Schemas for Phase 2 features."""
from datetime import datetime
from typing import List, Optional
from enum import Enum

from pydantic import BaseModel, Field


# ============ Risk Analysis ============
class RiskLevelEnum(str, Enum):
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class RiskTypeEnum(str, Enum):
    VAGUE = "vague"
    UNFAIR = "unfair"
    ILLEGAL = "illegal"
    MISSING = "missing"
    CONFLICT = "conflict"
    AMBIGUOUS = "ambiguous"


class RiskAnalysisRequest(BaseModel):
    contract_id: int
    contract_content: str


class RiskResponse(BaseModel):
    id: int
    contract_id: int
    clause_text: str
    risk_type: RiskTypeEnum
    risk_level: RiskLevelEnum
    description: str
    suggestion: str
    created_at: datetime

    class Config:
        from_attributes = True


# ============ Case Prediction ============
class PredictionRequest(BaseModel):
    case_description: str = Field(..., min_length=10)


class PredictionResponse(BaseModel):
    id: int
    user_id: int
    case_description: str
    predicted_outcome: str
    win_probability: float
    similar_cases: Optional[List[dict]] = None
    key_factors: Optional[List[str]] = None
    confidence: float
    created_at: datetime

    class Config:
        from_attributes = True


# ============ Knowledge Graph ============
class EntityTypeEnum(str, Enum):
    LAW = "law"
    ARTICLE = "article"
    CASE = "case"
    CONCEPT = "concept"
    ORGANIZATION = "organization"


class RelationTypeEnum(str, Enum):
    REFERENCES = "references"
    AMENDS = "amends"
    REPLACES = "replaces"
    INTERPRETS = "interprets"
    APPLIES = "applies"
    DEFINES = "defines"


class EntityCreate(BaseModel):
    name: str = Field(..., max_length=200)
    entity_type: EntityTypeEnum
    properties: Optional[dict] = None


class EntityResponse(BaseModel):
    id: int
    name: str
    entity_type: EntityTypeEnum
    properties: Optional[dict] = None
    created_at: datetime

    class Config:
        from_attributes = True


class RelationCreate(BaseModel):
    source_id: int
    target_id: int
    relation_type: RelationTypeEnum
    weight: float = 1.0


class RelationResponse(BaseModel):
    id: int
    source_id: int
    target_id: int
    relation_type: RelationTypeEnum
    weight: float
    created_at: datetime

    class Config:
        from_attributes = True


# ============ Translation ============
class TranslationRequest(BaseModel):
    source_text: str
    source_lang: str = Field(..., max_length=10)
    target_lang: str = Field(..., max_length=10)


class TranslationResponse(BaseModel):
    id: int
    source_text: str
    source_lang: str
    target_text: Optional[str] = None
    target_lang: str
    status: str
    created_at: datetime
    completed_at: Optional[datetime] = None

    class Config:
        from_attributes = True


# ============ Lawyer Matching ============
class LawyerCreate(BaseModel):
    name: str = Field(..., max_length=100)
    specialties: Optional[List[str]] = None
    experience_years: int = 0


class LawyerResponse(BaseModel):
    id: int
    name: str
    specialties: Optional[List[str]] = None
    experience_years: int
    success_rate: float
    cases_count: int
    rating: float
    created_at: datetime

    class Config:
        from_attributes = True


class RecommendationResponse(BaseModel):
    id: int
    case_description: str
    lawyer_id: int
    match_score: float
    match_reasons: Optional[List[str]] = None
    created_at: datetime

    class Config:
        from_attributes = True