api/models.py

import os, sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from typing import List, Optional, Dict, Any
from datetime import datetime
from pydantic import BaseModel


class SourceDocument(BaseModel):
    institution: str
    document_title: str
    url: str
    date: Optional[str] = None
    credibility: str = "official_government"


class SearchResult(BaseModel):
    entity_id: str
    entity_type: str
    name: str
    state: Optional[str] = None
    party: Optional[str] = None
    risk_score: Optional[int] = None
    # BUG-12 FIX: risk_level was missing -- frontend Components.ResultCard renders
    # class="risk-badge risk-badge--${entity.risk_level}" so without this field
    # every badge was always risk-badge--undefined which rendered as LOW visually.
    risk_level: Optional[str] = None
    sources: List[SourceDocument] = []


class SearchResponse(BaseModel):
    query: str
    total: int
    results: List[SearchResult]
    generated_at: str


class RiskFactor(BaseModel):
    name: str
    score: float
    weight: float
    description: str
    evidence: List[str] = []


class RiskResponse(BaseModel):
    entity_id: str
    entity_name: str
    risk_score: int
    risk_level: str
    factors: List[RiskFactor]
    explanation: str
    sources: List[SourceDocument] = []
    generated_at: str


class GraphNode(BaseModel):
    id: str
    label: str
    name: str
    properties: Dict[str, Any] = {}


class GraphEdge(BaseModel):
    source: str
    target: str
    relationship: str
    properties: Dict[str, Any] = {}
    source_document: Optional[str] = None


class GraphResponse(BaseModel):
    entity_id: str
    depth: int
    nodes: List[GraphNode]
    edges: List[GraphEdge]
    generated_at: str


class ProfileSection(BaseModel):
    section: str
    data: List[Dict[str, Any]]
    sources: List[SourceDocument] = []


class ProfileResponse(BaseModel):
    """BUG-22 FIX: added sources field -- profile.py was returning it but model lacked it,
    causing Pydantic ValidationError silently on some entity types."""
    entity_id: str
    entity_type: str
    name: str
    overview: Dict[str, Any] = {}
    sections: List[ProfileSection] = []
    sources: List[SourceDocument] = []
    risk_score: Optional[int] = None
    generated_at: str


class HealthResponse(BaseModel):
    status: str
    neo4j_connected: bool
    version: str
    generated_at: str


class StatsResponse(BaseModel):
    nodes: Dict[str, int]
    relationships: Dict[str, int]
    last_pipeline_run: Optional[str]
    generated_at: str


class FeedItem(BaseModel):
    headline: str
    entity_names: List[str]
    risk_level: str
    source: str
    url: Optional[str]
    detected_at: str
    summary: str