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

"""Phase 2 models: Risk, Prediction, Knowledge Graph, Translation, Lawyer."""
import enum
from datetime import datetime
from typing import Optional, List

from sqlalchemy import String, Text, DateTime, Enum as SQLEnum, Integer, ForeignKey, Float, JSON
from sqlalchemy.orm import Mapped, mapped_column

from app.core.database import Base


# ============ Feature 1: Contract Risk ============
class RiskLevel(str, enum.Enum):
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class RiskType(str, enum.Enum):
    VAGUE = "vague"               # 模糊条款
    UNFAIR = "unfair"             # 不公平条款
    ILLEGAL = "illegal"           # 违法条款
    MISSING = "missing"           # 缺失条款
    CONFLICT = "conflict"         # 冲突条款
    AMBIGUOUS = "ambiguous"       # 歧义条款


class ContractRisk(Base):
    """Contract risk analysis result."""
    __tablename__ = "contract_risks"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    contract_id: Mapped[int] = mapped_column(Integer, ForeignKey("contracts.id"), index=True)
    clause_text: Mapped[str] = mapped_column(Text)
    risk_type: Mapped[RiskType] = mapped_column(SQLEnum(RiskType))
    risk_level: Mapped[RiskLevel] = mapped_column(SQLEnum(RiskLevel))
    description: Mapped[str] = mapped_column(Text)
    suggestion: Mapped[str] = mapped_column(Text)
    position: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        contract_id: int,
        clause_text: str,
        risk_type: RiskType,
        risk_level: RiskLevel,
        description: str,
        suggestion: str,
        position: Optional[dict] = None,
        **kwargs
    ):
        self.contract_id = contract_id
        self.clause_text = clause_text
        self.risk_type = risk_type
        self.risk_level = risk_level
        self.description = description
        self.suggestion = suggestion
        self.position = position
        self.created_at = datetime.utcnow()


# ============ Feature 2: Case Prediction ============
class CasePrediction(Base):
    """Case prediction analysis."""
    __tablename__ = "case_predictions"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    user_id: Mapped[int] = mapped_column(Integer, ForeignKey("users.id"), index=True)
    case_description: Mapped[str] = mapped_column(Text)
    predicted_outcome: Mapped[str] = mapped_column(Text)
    win_probability: Mapped[float] = mapped_column(Float)
    similar_cases: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    key_factors: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    confidence: Mapped[float] = mapped_column(Float)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        user_id: int,
        case_description: str,
        predicted_outcome: str,
        win_probability: float,
        confidence: float,
        similar_cases: Optional[list] = None,
        key_factors: Optional[list] = None,
        **kwargs
    ):
        self.user_id = user_id
        self.case_description = case_description
        self.predicted_outcome = predicted_outcome
        self.win_probability = win_probability
        self.confidence = confidence
        self.similar_cases = similar_cases
        self.key_factors = key_factors
        self.created_at = datetime.utcnow()


# ============ Feature 3: Knowledge Graph ============
class EntityType(str, enum.Enum):
    LAW = "law"
    ARTICLE = "article"
    CASE = "case"
    CONCEPT = "concept"
    ORGANIZATION = "organization"


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


class LegalEntity(Base):
    """Legal knowledge graph entity."""
    __tablename__ = "legal_entities"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    name: Mapped[str] = mapped_column(String(200), index=True)
    entity_type: Mapped[EntityType] = mapped_column(SQLEnum(EntityType))
    properties: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
    source_id: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        name: str,
        entity_type: EntityType,
        properties: Optional[dict] = None,
        source_id: Optional[int] = None,
        **kwargs
    ):
        self.name = name
        self.entity_type = entity_type
        self.properties = properties
        self.source_id = source_id
        self.created_at = datetime.utcnow()


class LegalRelation(Base):
    """Legal knowledge graph relation."""
    __tablename__ = "legal_relations"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    source_id: Mapped[int] = mapped_column(Integer, ForeignKey("legal_entities.id"), index=True)
    target_id: Mapped[int] = mapped_column(Integer, ForeignKey("legal_entities.id"), index=True)
    relation_type: Mapped[RelationType] = mapped_column(SQLEnum(RelationType))
    weight: Mapped[float] = mapped_column(Float, default=1.0)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        source_id: int,
        target_id: int,
        relation_type: RelationType,
        weight: float = 1.0,
        **kwargs
    ):
        self.source_id = source_id
        self.target_id = target_id
        self.relation_type = relation_type
        self.weight = weight
        self.created_at = datetime.utcnow()


# ============ Feature 4: Translation ============
class TranslationRecord(Base):
    """Translation record."""
    __tablename__ = "translation_records"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    document_id: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    source_text: Mapped[str] = mapped_column(Text)
    source_lang: Mapped[str] = mapped_column(String(10))
    target_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    target_lang: Mapped[str] = mapped_column(String(10))
    status: Mapped[str] = mapped_column(String(20), default="pending")
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)
    completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)

    def __init__(
        self,
        source_text: str,
        source_lang: str,
        target_lang: str,
        document_id: Optional[int] = None,
        status: str = "pending",
        **kwargs
    ):
        self.source_text = source_text
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.document_id = document_id
        self.status = status
        self.created_at = datetime.utcnow()


# ============ Feature 5: Lawyer Matching ============
class Lawyer(Base):
    """Lawyer profile."""
    __tablename__ = "lawyers"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    name: Mapped[str] = mapped_column(String(100))
    specialties: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    experience_years: Mapped[int] = mapped_column(Integer, default=0)
    success_rate: Mapped[float] = mapped_column(Float, default=0.0)
    cases_count: Mapped[int] = mapped_column(Integer, default=0)
    rating: Mapped[float] = mapped_column(Float, default=0.0)
    bio: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    contact: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        name: str,
        specialties: Optional[list] = None,
        experience_years: int = 0,
        **kwargs
    ):
        self.name = name
        self.specialties = specialties or []
        self.experience_years = experience_years
        self.created_at = datetime.utcnow()


class LawyerRecommendation(Base):
    """Lawyer recommendation for a case."""
    __tablename__ = "lawyer_recommendations"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    case_description: Mapped[str] = mapped_column(Text)
    lawyer_id: Mapped[int] = mapped_column(Integer, ForeignKey("lawyers.id"), index=True)
    match_score: Mapped[float] = mapped_column(Float)
    match_reasons: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)

    def __init__(
        self,
        case_description: str,
        lawyer_id: int,
        match_score: float,
        match_reasons: Optional[list] = None,
        **kwargs
    ):
        self.case_description = case_description
        self.lawyer_id = lawyer_id
        self.match_score = match_score
        self.match_reasons = match_reasons
        self.created_at = datetime.utcnow()