Update mcp/schemas.py

mcp/schemas.py

@@ -1,34 +1,88 @@
-# 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 pydantic import BaseModel
 from typing import List, Optional
+from pydantic import BaseModel, Field, HttpUrl
+
+# ----------------------------------------------------------------------
+# Literature & concepts
+# ----------------------------------------------------------------------
 
 class Paper(BaseModel):
-    title: str
-    authors: str
-    summary: str
-    link: str
-    published: str
-    source: str
+    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]
-    definition: Optional[str]
+    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]
+    safety_report_id : Optional[str]
+    serious          : Optional[str]
+    reactions        : Optional[List[str]]
+    receivedate      : Optional[str]
 
-class UnifiedSearchResult(BaseModel):
-    papers: List[Paper]
-    umls: List[UMLSConcept]
-    drug_safety: List[DrugSafety]
-    ai_summary: str
-    suggested_reading: List[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