basion-agent 0.4.0__py3-none-any.whl

basion_agent/tools/knowledge_graph.py

@@ -0,0 +1,306 @@
+"""Knowledge Graph Tool for querying biomedical knowledge graphs."""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import aiohttp
+
+
+@dataclass
+class Entity:
+    """Base entity from knowledge graph."""
+
+    id: str
+    name: str
+    entity_type: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Entity":
+        return cls(
+            id=str(data.get("id", "")),
+            name=data.get("name") or data.get("label") or data.get("symbol", ""),
+            entity_type=data.get("type", "unknown"),
+        )
+
+
+@dataclass
+class Edge:
+    """Relationship edge in knowledge graph."""
+
+    source_id: str
+    source_type: str
+    target_id: str
+    target_type: str
+    relation_type: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Edge":
+        return cls(
+            source_id=str(data.get("source_id", "")),
+            source_type=data.get("source_type", ""),
+            target_id=str(data.get("target_id", "")),
+            target_type=data.get("target_type", ""),
+            relation_type=data.get("relation_type", ""),
+        )
+
+
+@dataclass
+class SimilarDisease:
+    """Disease similarity result."""
+
+    disease_id: str
+    disease_name: str
+    similarity_score: float
+    shared_count: int
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SimilarDisease":
+        return cls(
+            disease_id=str(data.get("disease_id", "")),
+            disease_name=data.get("disease_name", ""),
+            similarity_score=float(data.get("similarity_score", 0)),
+            shared_count=int(data.get("shared_count", 0)),
+        )
+
+
+@dataclass
+class PathStep:
+    """Step in shortest path."""
+
+    node_id: str
+    node_name: str
+    node_type: str
+    relation: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "PathStep":
+        return cls(
+            node_id=str(data.get("node_id", "")),
+            node_name=data.get("node_name", ""),
+            node_type=data.get("node_type", ""),
+            relation=data.get("relation"),
+        )
+
+
+class KnowledgeGraphTool:
+    """
+    Tool for querying biomedical knowledge graphs.
+
+    Usage:
+        # Initialize with gateway URL
+        kg = KnowledgeGraphTool(base_url="http://gateway:8080/s/knowledge-graph")
+
+        # Or from gateway client
+        kg = KnowledgeGraphTool.from_gateway(gateway_client)
+
+        # Query diseases
+        diseases = await kg.search_diseases(name="Huntington")
+
+        # Find similar diseases
+        similar = await kg.find_similar_diseases("Huntington Disease", limit=10)
+
+        # Get entity connections
+        edges = await kg.get_entity_network("BRCA1", "protein")
+
+        # Find path between entities
+        path = await kg.find_shortest_path("BRCA1", "protein", "Breast Cancer", "disease")
+    """
+
+    def __init__(self, base_url: str):
+        self.base_url = base_url.rstrip("/")
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    @classmethod
+    def from_gateway(cls, gateway_client) -> "KnowledgeGraphTool":
+        """Create from GatewayClient instance."""
+        return cls(gateway_client.knowledge_graph_url)
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            timeout = aiohttp.ClientTimeout(total=30.0)
+            self._session = aiohttp.ClientSession(timeout=timeout)
+        return self._session
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def _get(self, path: str, params: Optional[Dict] = None) -> Any:
+        session = await self._get_session()
+        url = f"{self.base_url}{path}"
+        async with session.get(url, params=params) as response:
+            response.raise_for_status()
+            return await response.json()
+
+    # === Entity Lookup Methods ===
+
+    async def search_diseases(
+        self,
+        name: Optional[str] = None,
+        orphacode: Optional[str] = None,
+        omim: Optional[str] = None,
+        limit: int = 20,
+    ) -> List[Dict]:
+        """Search for diseases by name or codes."""
+        params: Dict[str, Any] = {"limit": limit}
+        if name:
+            params["label"] = name
+        if orphacode:
+            params["orphacode"] = orphacode
+        if omim:
+            params["omim"] = omim
+        return await self._get("/diseases/", params)
+
+    async def get_disease(self, disease_id: int) -> Dict:
+        """Get disease by ID."""
+        return await self._get(f"/diseases/{disease_id}")
+
+    async def search_proteins(
+        self,
+        symbol: Optional[str] = None,
+        ensembl_id: Optional[str] = None,
+        limit: int = 20,
+    ) -> List[Dict]:
+        """Search for proteins/genes."""
+        params: Dict[str, Any] = {"limit": limit}
+        if symbol:
+            params["symbol"] = symbol
+        if ensembl_id:
+            params["ensembl_id"] = ensembl_id
+        return await self._get("/proteins/", params)
+
+    async def search_phenotypes(
+        self,
+        name: Optional[str] = None,
+        hpo_id: Optional[str] = None,
+        limit: int = 20,
+    ) -> List[Dict]:
+        """Search for phenotypes (HPO terms)."""
+        params: Dict[str, Any] = {"limit": limit}
+        if name:
+            params["name"] = name
+        if hpo_id:
+            params["hpo_id"] = hpo_id
+        return await self._get("/phenotypes/", params)
+
+    async def search_drugs(
+        self,
+        name: Optional[str] = None,
+        limit: int = 20,
+    ) -> List[Dict]:
+        """Search for drugs."""
+        params: Dict[str, Any] = {"limit": limit}
+        if name:
+            params["name"] = name
+        return await self._get("/drugs/", params)
+
+    async def search_pathways(
+        self,
+        name: Optional[str] = None,
+        limit: int = 20,
+    ) -> List[Dict]:
+        """Search for biological pathways."""
+        params: Dict[str, Any] = {"limit": limit}
+        if name:
+            params["name"] = name
+        return await self._get("/pathways/", params)
+
+    # === Similarity Analysis ===
+
+    async def find_similar_diseases(
+        self,
+        disease_name: str,
+        limit: int = 10,
+    ) -> List[SimilarDisease]:
+        """Find diseases similar by shared phenotypes."""
+        data = await self._get(
+            "/diseases/similarity",
+            params={"disease_name": disease_name, "limit": limit},
+        )
+        return [SimilarDisease.from_dict(d) for d in data.get("similar_diseases", [])]
+
+    async def find_similar_diseases_by_genes(
+        self,
+        disease_name: str,
+        limit: int = 10,
+    ) -> List[SimilarDisease]:
+        """Find diseases similar by shared genes."""
+        data = await self._get(
+            "/diseases/similarity-by-genes",
+            params={"disease_name": disease_name, "limit": limit},
+        )
+        return [SimilarDisease.from_dict(d) for d in data.get("similar_diseases", [])]
+
+    # === Graph Traversal ===
+
+    async def get_entity_network(
+        self,
+        entity_name: str,
+        entity_type: str,
+    ) -> List[Edge]:
+        """
+        Get all connections for an entity.
+
+        Args:
+            entity_name: Name of the entity (e.g., "BRCA1", "Huntington Disease")
+            entity_type: One of: protein, phenotype, disease, pathway,
+                        molecular_function, cellular_component, biological_process
+        """
+        data = await self._get(
+            "/edges/network",
+            params={"entity_name": entity_name, "entity_type": entity_type},
+        )
+        return [Edge.from_dict(e) for e in data]
+
+    async def k_hop_traversal(
+        self,
+        entity_name: str,
+        entity_type: str,
+        k: int = 2,
+        limit_edges: int = 100,
+    ) -> Dict:
+        """
+        Perform k-hop BFS traversal from entity.
+
+        Returns subgraph with nodes and edges within k hops.
+        """
+        return await self._get(
+            "/graph/k-hop",
+            params={
+                "entity_name": entity_name,
+                "entity_type": entity_type,
+                "k": k,
+                "limit_edges": limit_edges,
+            },
+        )
+
+    async def find_shortest_path(
+        self,
+        start_name: str,
+        start_type: str,
+        end_name: str,
+        end_type: str,
+        max_hops: int = 5,
+    ) -> List[PathStep]:
+        """
+        Find shortest path between two entities (Mechanism of Action).
+
+        Args:
+            start_name: Starting entity name
+            start_type: Starting entity type
+            end_name: Target entity name
+            end_type: Target entity type
+            max_hops: Maximum path length
+        """
+        data = await self._get(
+            "/paths/shortest",
+            params={
+                "start_name": start_name,
+                "start_type": start_type,
+                "end_name": end_name,
+                "end_type": end_type,
+                "max_hops": max_hops,
+            },
+        )
+        return [PathStep.from_dict(s) for s in data.get("path", [])]