From bb5f29653e0718492a8c8c23a44d3dab043b0e00 Mon Sep 17 00:00:00 2001 From: engkinandatama Date: Sat, 10 Oct 2026 10:05:04 +0700 Subject: [PATCH] feat(graph): implement Knowledge Graph data model, in-memory store, builder, and MCP tool (v0.8.0) --- docs/API.md | 3 +- pyproject.toml | 2 +- src/biocontext/cli.py | 16 +- src/biocontext/config.py | 13 ++ src/biocontext/graph/__init__.py | 23 +++ src/biocontext/graph/builder.py | 293 +++++++++++++++++++++++++++++++ src/biocontext/graph/models.py | 68 +++++++ src/biocontext/graph/store.py | 134 ++++++++++++++ src/biocontext/server.py | 44 +++++ tests/test_graph.py | 137 +++++++++++++++ 10 files changed, 730 insertions(+), 3 deletions(-) create mode 100644 src/biocontext/graph/__init__.py create mode 100644 src/biocontext/graph/builder.py create mode 100644 src/biocontext/graph/models.py create mode 100644 src/biocontext/graph/store.py create mode 100644 tests/test_graph.py diff --git a/docs/API.md b/docs/API.md index 926d7f1..2c5d0fc 100644 --- a/docs/API.md +++ b/docs/API.md @@ -150,8 +150,9 @@ asyncio.run(main()) | `targets` | `biocontext targets [-l INT]` | Retrieve Open Targets evidence-backed disease associations. | | `literature` | `biocontext literature [-l INT]` | Search supporting peer-reviewed publications from Europe PMC / PubMed. | | `paper` | `biocontext paper ` | Fetch detailed paper metadata and abstract by PMID, PMCID, or DOI. | +| `graph` | `biocontext graph [--taxon INT]` | Construct connected biological Knowledge Subgraph (PRD-21). | | `mouse` | `biocontext mouse ` | Lookup mouse gene model via MGI. | -| `cache` | `biocontext cache [stats\|clear]` | Inspect or clear SQLite cache. | +| `cache` | `biocontext cache [stats|clear]` | Inspect or clear SQLite cache. | | `serve` | `biocontext serve` | Start stdio MCP server for AI clients. | diff --git a/pyproject.toml b/pyproject.toml index f292011..44509ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "biocontext-mcp" -version = "0.7.0" +version = "0.8.0" description = "Authoritative Biological Entity Resolution & Contextual Intelligence Framework" readme = "README.md" authors = [ diff --git a/src/biocontext/cli.py b/src/biocontext/cli.py index 3039ab6..3c1a9c2 100644 --- a/src/biocontext/cli.py +++ b/src/biocontext/cli.py @@ -235,6 +235,20 @@ async def run_cli_async(args: argparse.Namespace) -> int: print(pub.model_dump_json(indent=2)) return 0 + elif args.command == "graph": + from biocontext.graph.builder import KnowledgeGraphBuilder + builder = KnowledgeGraphBuilder(resolver=resolver) + subgraph = await builder.build_gene_subgraph( + query=args.query, + taxon_id=args.taxon, + include_pathways=args.include_pathways, + include_go=args.include_go, + include_diseases=args.include_diseases, + include_literature=args.include_literature, + ) + print(subgraph.model_dump_json(indent=2)) + return 0 + elif args.command == "cache": cache = resolver.cache if args.cache_action == "clear": @@ -270,7 +284,7 @@ def main(): elif args.command in ( "resolve", "batch", "protein", "transcripts", "ortholog", "mouse", "annotate", "go", "pathway", "pathway-info", "disease", "targets", - "literature", "paper", "cache" + "literature", "paper", "graph", "cache" ): sys.exit(asyncio.run(run_cli_async(args))) else: diff --git a/src/biocontext/config.py b/src/biocontext/config.py index 121f10e..9abf791 100644 --- a/src/biocontext/config.py +++ b/src/biocontext/config.py @@ -197,6 +197,19 @@ def get_headers(cls, email: Optional[str] = None) -> Dict[str, str]: {"flags": ["identifier"], "help": "Publication ID (e.g. 30514107, PMC6280721, or 10.1038/...)"} ] }, + { + "name": "graph", + "help": "Construct connected biological Knowledge Subgraph (PRD-21)", + "arguments": [ + {"flags": ["query"], "help": "Gene symbol (e.g. TP53, BRCA1) or alias"}, + {"flags": ["--taxon"], "type": int, "default": 9606, "help": "NCBI Taxonomy ID (default: 9606)"}, + {"flags": ["--no-pathways"], "dest": "include_pathways", "action": "store_false", "help": "Exclude Reactome pathways"}, + {"flags": ["--no-go"], "dest": "include_go", "action": "store_false", "help": "Exclude Gene Ontology terms"}, + {"flags": ["--no-diseases"], "dest": "include_diseases", "action": "store_false", "help": "Exclude disease associations"}, + {"flags": ["--no-literature"], "dest": "include_literature", "action": "store_false", "help": "Exclude literature citations"} + ] + }, + { diff --git a/src/biocontext/graph/__init__.py b/src/biocontext/graph/__init__.py new file mode 100644 index 0000000..f09c482 --- /dev/null +++ b/src/biocontext/graph/__init__.py @@ -0,0 +1,23 @@ +"""BioContext Knowledge Graph module.""" + +from biocontext.graph.builder import KnowledgeGraphBuilder +from biocontext.graph.models import ( + EdgeType, + GraphEdge, + GraphNode, + NodeType, + Provenance, + SubGraphContext, +) +from biocontext.graph.store import InMemoryGraphStore + +__all__ = [ + "NodeType", + "EdgeType", + "Provenance", + "GraphNode", + "GraphEdge", + "SubGraphContext", + "InMemoryGraphStore", + "KnowledgeGraphBuilder", +] diff --git a/src/biocontext/graph/builder.py b/src/biocontext/graph/builder.py new file mode 100644 index 0000000..a7a50a9 --- /dev/null +++ b/src/biocontext/graph/builder.py @@ -0,0 +1,293 @@ +"""Knowledge Graph Builder pipeline integrating BioContext adapters into structured subgraphs.""" + +from typing import List, Optional +from biocontext.graph.models import ( + EdgeType, + GraphEdge, + GraphNode, + NodeType, + Provenance, + SubGraphContext, +) +from biocontext.graph.store import InMemoryGraphStore +from biocontext.logging import get_logger +from biocontext.resolver import EntityResolver +from biocontext.schemas import ResolutionResult + +logger = get_logger("graph.builder") + + +class KnowledgeGraphBuilder: + """Orchestrates entity resolution and builds connected biological subgraphs.""" + + def __init__(self, resolver: Optional[EntityResolver] = None): + self.resolver = resolver or EntityResolver() + + async def build_gene_subgraph( + self, + query: str, + taxon_id: int = 9606, + include_pathways: bool = True, + include_go: bool = True, + include_diseases: bool = True, + include_literature: bool = True, + limit_diseases: int = 5, + limit_pathways: int = 5, + limit_literature: int = 5, + ) -> SubGraphContext: + """Resolve a gene and construct an interconnected biological subgraph across databases.""" + store = InMemoryGraphStore() + + # 1. Resolve Gene + res: ResolutionResult = await self.resolver.resolve(query=query, taxon_id=taxon_id) + if not res.resolved_entity: + # Unresolved entity fallback node + unresolved_node = GraphNode( + id=f"unresolved:{query}", + name=query, + node_type=NodeType.GENE, + properties={"status": "unresolved", "query": query}, + ) + store.add_node(unresolved_node) + return store.get_subgraph(unresolved_node.id, max_depth=0) + + gene = res.resolved_entity + gene_node_id = f"gene:{gene.symbol}" + gene_node = GraphNode( + id=gene_node_id, + name=gene.symbol, + node_type=NodeType.GENE, + properties={ + "hgnc_id": gene.hgnc_id, + "ncbi_gene_id": gene.ncbi_gene_id, + "ensembl_gene_id": gene.ensembl_gene_id, + "description": gene.name, + "species": gene.species, + "synonyms": gene.synonyms, + }, + provenance=Provenance( + source_database="HGNC", + confidence=res.confidence_score, + source_id=gene.hgnc_id, + ), + ) + store.add_node(gene_node) + + # 2. Proteins & Transcripts + primary_accession = gene.uniprot_ids[0] if gene.uniprot_ids else None + for acc in gene.uniprot_ids[:2]: # Top 2 accessions + protein_node_id = f"protein:{acc}" + protein_node = GraphNode( + id=protein_node_id, + name=f"{gene.symbol} ({acc})", + node_type=NodeType.PROTEIN, + properties={"accession": acc, "gene_symbol": gene.symbol}, + provenance=Provenance( + source_database="UniProt", + confidence=1.0, + source_id=acc, + ), + ) + store.add_node(protein_node) + store.add_edge( + GraphEdge( + source_id=gene_node_id, + target_id=protein_node_id, + edge_type=EdgeType.ENCODES, + weight=1.0, + provenance=Provenance(source_database="UniProt", confidence=1.0, source_id=acc), + ) + ) + + # 3. Transcripts (from Ensembl if available) + for transcript in gene.transcripts[:3]: # Top 3 transcripts + t_node_id = f"transcript:{transcript.transcript_id}" + t_node = GraphNode( + id=t_node_id, + name=transcript.name or transcript.transcript_id, + node_type=NodeType.TRANSCRIPT, + properties={ + "transcript_id": transcript.transcript_id, + "is_canonical": transcript.is_canonical, + "biotype": transcript.biotype, + }, + provenance=Provenance( + source_database="Ensembl", + confidence=1.0, + source_id=transcript.transcript_id, + ), + ) + store.add_node(t_node) + store.add_edge( + GraphEdge( + source_id=gene_node_id, + target_id=t_node_id, + edge_type=EdgeType.TRANSCRIBES, + weight=1.0, + provenance=Provenance(source_database="Ensembl", confidence=1.0), + ) + ) + + # 4. Pathways (Reactome) + if include_pathways and primary_accession: + try: + p_ctx = await self.resolver.get_pathways(primary_accession) + for pathway in p_ctx.pathways[:limit_pathways]: + p_node_id = f"pathway:{pathway.st_id}" + p_node = GraphNode( + id=p_node_id, + name=pathway.name, + node_type=NodeType.PATHWAY, + properties={ + "st_id": pathway.st_id, + "url": pathway.url, + "is_in_disease": pathway.is_in_disease, + }, + provenance=Provenance( + source_database="Reactome", + confidence=1.0, + source_id=pathway.st_id, + ), + ) + store.add_node(p_node) + # Link to protein + protein_node_id = f"protein:{primary_accession}" + if store.get_node(protein_node_id): + store.add_edge( + GraphEdge( + source_id=protein_node_id, + target_id=p_node_id, + edge_type=EdgeType.PARTICIPATES_IN, + weight=1.0, + provenance=Provenance(source_database="Reactome", confidence=1.0), + ) + ) + else: + # Fallback link to gene + store.add_edge( + GraphEdge( + source_id=gene_node_id, + target_id=p_node_id, + edge_type=EdgeType.PARTICIPATES_IN, + weight=1.0, + provenance=Provenance(source_database="Reactome", confidence=1.0), + ) + ) + except Exception as e: + logger.warning(f"Failed to fetch pathways for graph: {e}") + + # 5. Gene Ontology (GO Terms) + if include_go and primary_accession: + try: + go_ctx = await self.resolver.get_functional_annotations(primary_accession) + # Sample 2 biological processes and 2 molecular functions + sampled_go = (go_ctx.biological_processes[:2] + go_ctx.molecular_functions[:2]) + for go in sampled_go: + go_node_id = f"go:{go.go_id}" + go_node = GraphNode( + id=go_node_id, + name=go.name or go.go_id, + node_type=NodeType.GO_TERM, + properties={"go_id": go.go_id, "aspect": go.aspect}, + provenance=Provenance( + source_database="GeneOntology", + confidence=1.0, + source_id=go.go_id, + evidence_code=go.evidence_code, + ), + ) + store.add_node(go_node) + protein_node_id = f"protein:{primary_accession}" + src_id = protein_node_id if store.get_node(protein_node_id) else gene_node_id + edge_type = ( + EdgeType.INVOLVED_IN + if go.aspect == "biological_process" + else EdgeType.HAS_FUNCTION + ) + store.add_edge( + GraphEdge( + source_id=src_id, + target_id=go_node_id, + edge_type=edge_type, + weight=1.0, + provenance=Provenance(source_database="GeneOntology", confidence=1.0), + ) + ) + except Exception as e: + logger.warning(f"Failed to fetch GO annotations for graph: {e}") + + # 6. Target-Disease Associations (Open Targets) + if include_diseases and gene.ensembl_gene_id: + try: + target_diseases = await self.resolver.get_target_diseases(gene.ensembl_gene_id) + for assoc in target_diseases.associations[:limit_diseases]: + disease_node_id = f"disease:{assoc.disease_id}" + disease_node = GraphNode( + id=disease_node_id, + name=assoc.disease_name, + node_type=NodeType.DISEASE, + properties={ + "disease_id": assoc.disease_id, + "association_score": assoc.score, + }, + provenance=Provenance( + source_database="OpenTargets", + confidence=assoc.score, + source_id=assoc.disease_id, + ), + ) + store.add_node(disease_node) + store.add_edge( + GraphEdge( + source_id=gene_node_id, + target_id=disease_node_id, + edge_type=EdgeType.ASSOCIATED_WITH, + weight=assoc.score, + provenance=Provenance( + source_database="OpenTargets", + confidence=assoc.score, + source_id=assoc.disease_id, + ), + properties={"datatype_scores": assoc.datatype_scores or {}}, + ) + ) + except Exception as e: + logger.warning(f"Failed to fetch disease associations for graph: {e}") + + # 7. Supporting Literature (Europe PMC) + if include_literature: + try: + lit_ctx = await self.resolver.get_supporting_publications(gene.symbol, limit=limit_literature) + for pub in lit_ctx.publications[:limit_literature]: + pub_id_str = pub.pmid or pub.doi or pub.title[:30] + pub_node_id = f"pub:{pub_id_str}" + pub_node = GraphNode( + id=pub_node_id, + name=pub.title, + node_type=NodeType.PUBLICATION, + properties={ + "pmid": pub.pmid, + "doi": pub.doi, + "year": pub.pub_year, + "journal": pub.journal, + }, + provenance=Provenance( + source_database="EuropePMC", + confidence=1.0, + source_id=pub.pmid, + ), + ) + store.add_node(pub_node) + store.add_edge( + GraphEdge( + source_id=gene_node_id, + target_id=pub_node_id, + edge_type=EdgeType.CITED_IN, + weight=1.0, + provenance=Provenance(source_database="EuropePMC", confidence=1.0), + ) + ) + except Exception as e: + logger.warning(f"Failed to fetch literature for graph: {e}") + + return store.get_subgraph(gene_node_id, max_depth=2) diff --git a/src/biocontext/graph/models.py b/src/biocontext/graph/models.py new file mode 100644 index 0000000..32b4692 --- /dev/null +++ b/src/biocontext/graph/models.py @@ -0,0 +1,68 @@ +"""Biological Knowledge Graph models and schemas for BioContext (PRD-21).""" + +from enum import Enum +from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field + + +class NodeType(str, Enum): + """Categorical types of biological entities in the Knowledge Graph.""" + GENE = "gene" + PROTEIN = "protein" + TRANSCRIPT = "transcript" + PATHWAY = "pathway" + DISEASE = "disease" + PUBLICATION = "publication" + GO_TERM = "go_term" + + +class EdgeType(str, Enum): + """Authoritative semantic biological relationships between entities.""" + ENCODES = "encodes" + ENCODED_BY = "encoded_by" + TRANSCRIBES = "transcribes" + TRANSCRIBED_FROM = "transcribed_from" + INTERACTS_WITH = "interacts_with" + PARTICIPATES_IN = "participates_in" + HAS_FUNCTION = "has_function" + INVOLVED_IN = "involved_in" + ASSOCIATED_WITH = "associated_with" + CITED_IN = "cited_in" + ORTHOLOG_OF = "ortholog_of" + + +class Provenance(BaseModel): + """Audit trail and confidence metadata for graph nodes and edges.""" + source_database: str = Field(..., description="Authoritative primary source (e.g. HGNC, UniProt, Reactome, OpenTargets)") + confidence: float = Field(1.0, ge=0.0, le=1.0, description="Confidence or association score between 0.0 and 1.0") + source_id: Optional[str] = Field(None, description="Identifier in source database (e.g. HGNC:11998, R-HSA-5357801)") + evidence_code: Optional[str] = Field(None, description="Experimental evidence code if applicable (e.g. IDA, IEA)") + + +class GraphNode(BaseModel): + """Single entity node in the Knowledge Graph.""" + id: str = Field(..., description="Unique canonical node ID (e.g. 'gene:TP53', 'protein:P04637', 'disease:MONDO:0018875')") + name: str = Field(..., description="Human-readable primary display name or symbol") + node_type: NodeType = Field(..., description="Biological entity category") + properties: Dict[str, Any] = Field(default_factory=dict, description="Additional structured properties") + provenance: Optional[Provenance] = Field(None, description="Provenance metadata for the node") + + +class GraphEdge(BaseModel): + """Directed semantic edge connecting two entities in the Knowledge Graph.""" + source_id: str = Field(..., description="Source node ID") + target_id: str = Field(..., description="Target node ID") + edge_type: EdgeType = Field(..., description="Semantic relationship category") + weight: float = Field(1.0, ge=0.0, le=1.0, description="Edge weight / association strength") + provenance: Optional[Provenance] = Field(None, description="Source attribution and evidence") + properties: Dict[str, Any] = Field(default_factory=dict, description="Edge specific annotations (e.g. mechanism, datatype_scores)") + + +class SubGraphContext(BaseModel): + """A connected subgraph centered around an entity query.""" + query: str = Field(..., description="Original input query or seed entity ID") + root_node_id: str = Field(..., description="The focal entity node ID") + nodes: List[GraphNode] = Field(default_factory=list, description="Unique nodes present in the subgraph") + edges: List[GraphEdge] = Field(default_factory=list, description="Relationships connecting nodes in the subgraph") + total_nodes: int = Field(0, description="Total node count") + total_edges: int = Field(0, description="Total edge count") diff --git a/src/biocontext/graph/store.py b/src/biocontext/graph/store.py new file mode 100644 index 0000000..d4da046 --- /dev/null +++ b/src/biocontext/graph/store.py @@ -0,0 +1,134 @@ +"""In-memory Knowledge Graph storage engine for BioContext.""" + +from typing import Dict, List, Optional, Set +from biocontext.graph.models import EdgeType, GraphEdge, GraphNode, NodeType, SubGraphContext + + +class InMemoryGraphStore: + """Lightweight, deterministic in-memory graph structure for biological entity relationships.""" + + def __init__(self): + self._nodes: Dict[str, GraphNode] = {} + self._out_edges: Dict[str, List[GraphEdge]] = {} + self._in_edges: Dict[str, List[GraphEdge]] = {} + + def add_node(self, node: GraphNode) -> None: + """Add or update a node in the graph.""" + self._nodes[node.id] = node + if node.id not in self._out_edges: + self._out_edges[node.id] = [] + if node.id not in self._in_edges: + self._in_edges[node.id] = [] + + def get_node(self, node_id: str) -> Optional[GraphNode]: + """Retrieve a node by its unique ID.""" + return self._nodes.get(node_id) + + def add_edge(self, edge: GraphEdge) -> None: + """Add a directed edge between two existing nodes.""" + if edge.source_id not in self._nodes: + raise ValueError(f"Source node '{edge.source_id}' does not exist in graph.") + if edge.target_id not in self._nodes: + raise ValueError(f"Target node '{edge.target_id}' does not exist in graph.") + + # Avoid exact duplicate edges + existing = [ + e for e in self._out_edges[edge.source_id] + if e.target_id == edge.target_id and e.edge_type == edge.edge_type + ] + if not existing: + self._out_edges[edge.source_id].append(edge) + self._in_edges[edge.target_id].append(edge) + + def get_outgoing_edges(self, node_id: str, edge_type: Optional[EdgeType] = None) -> List[GraphEdge]: + """Return all outgoing edges from node_id, optionally filtered by edge_type.""" + edges = self._out_edges.get(node_id, []) + if edge_type: + return [e for e in edges if e.edge_type == edge_type] + return list(edges) + + def get_incoming_edges(self, node_id: str, edge_type: Optional[EdgeType] = None) -> List[GraphEdge]: + """Return all incoming edges to node_id, optionally filtered by edge_type.""" + edges = self._in_edges.get(node_id, []) + if edge_type: + return [e for e in edges if e.edge_type == edge_type] + return list(edges) + + def get_neighbors( + self, + node_id: str, + direction: str = "both", + edge_type: Optional[EdgeType] = None, + node_type: Optional[NodeType] = None, + ) -> List[GraphNode]: + """Retrieve connected neighboring nodes according to direction and filters.""" + neighbor_ids: Set[str] = set() + + if direction in ("out", "both"): + for e in self.get_outgoing_edges(node_id, edge_type=edge_type): + neighbor_ids.add(e.target_id) + + if direction in ("in", "both"): + for e in self.get_incoming_edges(node_id, edge_type=edge_type): + neighbor_ids.add(e.source_id) + + neighbors = [self._nodes[nid] for nid in neighbor_ids if nid in self._nodes] + if node_type: + neighbors = [n for n in neighbors if n.node_type == node_type] + return neighbors + + def get_subgraph(self, root_id: str, max_depth: int = 1) -> SubGraphContext: + """Extract a connected subgraph up to max_depth hops around root_id.""" + if root_id not in self._nodes: + raise KeyError(f"Root node '{root_id}' not found in graph.") + + visited_nodes: Set[str] = {root_id} + frontier: Set[str] = {root_id} + collected_edges: List[GraphEdge] = [] + seen_edges: Set[tuple] = set() + + for _ in range(max_depth): + next_frontier: Set[str] = set() + for current_id in frontier: + # Outgoing edges + for edge in self.get_outgoing_edges(current_id): + edge_key = (edge.source_id, edge.target_id, edge.edge_type) + if edge_key not in seen_edges: + seen_edges.add(edge_key) + collected_edges.append(edge) + if edge.target_id not in visited_nodes: + visited_nodes.add(edge.target_id) + next_frontier.add(edge.target_id) + + # Incoming edges + for edge in self.get_incoming_edges(current_id): + edge_key = (edge.source_id, edge.target_id, edge.edge_type) + if edge_key not in seen_edges: + seen_edges.add(edge_key) + collected_edges.append(edge) + if edge.source_id not in visited_nodes: + visited_nodes.add(edge.source_id) + next_frontier.add(edge.source_id) + + frontier = next_frontier + if not frontier: + break + + nodes_list = [self._nodes[nid] for nid in visited_nodes] + root_node = self._nodes[root_id] + return SubGraphContext( + query=root_node.name, + root_node_id=root_id, + nodes=nodes_list, + edges=collected_edges, + total_nodes=len(nodes_list), + total_edges=len(collected_edges), + ) + + @property + def node_count(self) -> int: + return len(self._nodes) + + @property + def edge_count(self) -> int: + return sum(len(edges) for edges in self._out_edges.values()) diff --git a/src/biocontext/server.py b/src/biocontext/server.py index ecb2aa2..e0f6b28 100644 --- a/src/biocontext/server.py +++ b/src/biocontext/server.py @@ -305,6 +305,49 @@ async def get_publication_details(identifier: str) -> str: return pub.model_dump_json(indent=2) +@mcp.tool() +async def build_entity_subgraph( + query: str, + taxon_id: int = 9606, + include_pathways: bool = True, + include_go: bool = True, + include_diseases: bool = True, + include_literature: bool = True, +) -> str: + """Construct an ontological Knowledge Subgraph centered on a biological entity (PRD-21). + + Connects the seed gene across multiple biological layers: + - Encoded proteins (UniProt) + - Transcripts (Ensembl) + - Reactome pathways + - Gene Ontology terms (processes & molecular functions) + - Disease associations with evidence scores (Open Targets) + - Key supporting peer-reviewed publications (Europe PMC) + + Args: + query: Gene symbol (e.g. 'TP53', 'BRCA1', 'EGFR') or alias. + taxon_id: NCBI Taxonomy ID (default: 9606 for human). + include_pathways: Include Reactome pathway nodes and edges (default: True). + include_go: Include Gene Ontology nodes and edges (default: True). + include_diseases: Include Open Targets disease association nodes and edges (default: True). + include_literature: Include Europe PMC publication nodes and edges (default: True). + + Returns: + JSON string representing the connected subgraph nodes, edges, weights, and provenance. + """ + from biocontext.graph.builder import KnowledgeGraphBuilder + builder = KnowledgeGraphBuilder(resolver=resolver) + subgraph = await builder.build_gene_subgraph( + query=query, + taxon_id=taxon_id, + include_pathways=include_pathways, + include_go=include_go, + include_diseases=include_diseases, + include_literature=include_literature, + ) + return subgraph.model_dump_json(indent=2) + + def main(): """Run MCP server over stdio.""" mcp.run(transport="stdio") @@ -312,3 +355,4 @@ def main(): if __name__ == "__main__": main() + diff --git a/tests/test_graph.py b/tests/test_graph.py new file mode 100644 index 0000000..b629f04 --- /dev/null +++ b/tests/test_graph.py @@ -0,0 +1,137 @@ +"""Unit and integration tests for BioContext Knowledge Graph (PRD-21).""" + +import pytest +from biocontext.graph.builder import KnowledgeGraphBuilder +from biocontext.graph.models import ( + EdgeType, + GraphEdge, + GraphNode, + NodeType, + Provenance, + SubGraphContext, +) +from biocontext.graph.store import InMemoryGraphStore +from biocontext.resolver import EntityResolver + + +def test_graph_store_add_and_query(): + store = InMemoryGraphStore() + + gene_node = GraphNode( + id="gene:TP53", + name="TP53", + node_type=NodeType.GENE, + properties={"hgnc_id": "HGNC:11998"}, + provenance=Provenance(source_database="HGNC", confidence=1.0), + ) + protein_node = GraphNode( + id="protein:P04637", + name="TP53 (P04637)", + node_type=NodeType.PROTEIN, + properties={"accession": "P04637"}, + provenance=Provenance(source_database="UniProt", confidence=1.0), + ) + disease_node = GraphNode( + id="disease:MONDO:0018875", + name="Li-Fraumeni syndrome", + node_type=NodeType.DISEASE, + properties={"mondo_id": "MONDO:0018875"}, + provenance=Provenance(source_database="MONDO", confidence=1.0), + ) + + store.add_node(gene_node) + store.add_node(protein_node) + store.add_node(disease_node) + + assert store.node_count == 3 + assert store.get_node("gene:TP53") is not None + + # Add edges + edge_encodes = GraphEdge( + source_id="gene:TP53", + target_id="protein:P04637", + edge_type=EdgeType.ENCODES, + weight=1.0, + ) + edge_assoc = GraphEdge( + source_id="gene:TP53", + target_id="disease:MONDO:0018875", + edge_type=EdgeType.ASSOCIATED_WITH, + weight=0.95, + ) + + store.add_edge(edge_encodes) + store.add_edge(edge_assoc) + + assert store.edge_count == 2 + + # Query outgoing + outgoing = store.get_outgoing_edges("gene:TP53") + assert len(outgoing) == 2 + + # Filtered by edge_type + encodes_edges = store.get_outgoing_edges("gene:TP53", edge_type=EdgeType.ENCODES) + assert len(encodes_edges) == 1 + assert encodes_edges[0].target_id == "protein:P04637" + + # Query neighbors + neighbors = store.get_neighbors("gene:TP53", direction="out") + assert len(neighbors) == 2 + + # Subgraph extraction + subgraph: SubGraphContext = store.get_subgraph("gene:TP53", max_depth=1) + assert subgraph.total_nodes == 3 + assert subgraph.total_edges == 2 + assert subgraph.root_node_id == "gene:TP53" + + +@pytest.mark.asyncio +async def test_knowledge_graph_builder_tp53(): + resolver = EntityResolver() + builder = KnowledgeGraphBuilder(resolver=resolver) + + subgraph = await builder.build_gene_subgraph( + query="TP53", + include_pathways=True, + include_go=True, + include_diseases=True, + include_literature=True, + limit_diseases=3, + limit_pathways=3, + limit_literature=3, + ) + + assert subgraph.root_node_id == "gene:TP53" + assert subgraph.total_nodes > 1 + assert subgraph.total_edges > 1 + + node_types = {n.node_type for n in subgraph.nodes} + assert NodeType.GENE in node_types + assert NodeType.PROTEIN in node_types + + edge_types = {e.edge_type for e in subgraph.edges} + assert EdgeType.ENCODES in edge_types + + # Provenance verification + for n in subgraph.nodes: + if n.provenance: + assert n.provenance.source_database is not None + assert 0.0 <= n.provenance.confidence <= 1.0 + + +@pytest.mark.asyncio +async def test_mcp_tool_build_entity_subgraph(): + from biocontext.server import build_entity_subgraph + import json + + res_json = await build_entity_subgraph( + query="BRCA1", + include_pathways=True, + include_go=False, + include_diseases=True, + include_literature=False, + ) + data = json.loads(res_json) + assert data["root_node_id"] == "gene:BRCA1" + assert len(data["nodes"]) >= 2 + assert len(data["edges"]) >= 1