mcp/schemas.py

#!/usr/bin/env python3
"""MedGenesis – **data schemas** powered by `pydantic`.

This module formalises the JSON objects exchanged between the backend
(orchestrator) and the Streamlit UI.  It now aligns with the *enhanced*
pipeline that returns live genes, trials, and Open Targets edges.
"""
from __future__ import annotations

from typing import List, Optional
from pydantic import BaseModel, Field, HttpUrl

# ----------------------------------------------------------------------
# Literature & concepts
# ----------------------------------------------------------------------

class Paper(BaseModel):
    title     : str
    authors   : str
    summary   : str
    link      : HttpUrl
    published : str  = Field(description="ISO date string")
    source    : str  = Field(description="'arxiv' | 'pubmed' | other")

class UMLSConcept(BaseModel):
    term       : str
    cui        : Optional[str]
    name       : Optional[str]
    rootSource : Optional[str]

# ----------------------------------------------------------------------
# Safety & biomedical enrichers
# ----------------------------------------------------------------------

class DrugSafety(BaseModel):
    safety_report_id : Optional[str]
    serious          : Optional[str]
    reactions        : Optional[List[str]]
    receivedate      : Optional[str]

class GeneInfo(BaseModel):
    symbol     : str
    name       : Optional[str]
    summary    : Optional[str]
    entrezgene : Optional[int]
    clinvar    : Optional[dict]

class MeshDefinition(BaseModel):
    term       : str
    definition : str

class GeneDiseaseLink(BaseModel):
    gene_symbol : str
    disease_id  : str
    score       : float

class ClinicalTrial(BaseModel):
    nct_id     : str = Field(alias="nctId")
    brief_title: str = Field(alias="briefTitle")
    phase      : Optional[str]
    status     : Optional[str]
    start_date : Optional[str]

class OTAssociation(BaseModel):
    score       : float
    datatypeId  : str
    datasourceId: str
    disease     : dict  # {id, name}
    target      : dict  # {id, symbol}

# ----------------------------------------------------------------------
# Search I/O
# ----------------------------------------------------------------------

class UnifiedSearchInput(BaseModel):
    query: str

class UnifiedSearchResult(BaseModel):
    papers          : List[Paper]
    umls            : List[UMLSConcept]
    drug_safety     : List[DrugSafety]
    genes           : List[GeneInfo]
    mesh_defs       : List[MeshDefinition]
    gene_disease    : List[GeneDiseaseLink]
    clinical_trials : List[ClinicalTrial]
    ot_associations : Optional[List[OTAssociation]] = []
    ai_summary      : str
    llm_used        : str