Add server-side BFS entity graph endpoint

ontokit/api/routes/classes.py

@@ -3,8 +3,9 @@
 from typing import Annotated
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, Query, status
 
+from ontokit.schemas.graph import EntityGraphResponse
 from ontokit.schemas.owl_class import (
     OWLClassCreate,
     OWLClassListResponse,
@@ -57,6 +58,39 @@ async def create_class(
     return await service.create_class(ontology_id, owl_class)
 
 
+@router.get(
+    "/ontologies/{ontology_id}/classes/graph",
+    response_model=EntityGraphResponse,
+)
+async def get_class_graph(
+    ontology_id: UUID,
+    service: Annotated[OntologyService, Depends(get_ontology_service)],
+    class_iri: str = Query(description="IRI of the class to build the graph around"),
+    branch: str = "main",
+    ancestors_depth: int = Query(default=5, ge=0, le=10),
+    descendants_depth: int = Query(default=2, ge=0, le=10),
+    max_nodes: int = Query(default=200, ge=1, le=500),
+    include_see_also: bool = True,
+) -> EntityGraphResponse:
+    """Build a multi-hop entity graph around a class via BFS.
+
+    Returns nodes and edges for visualization, with lineage-based node types
+    for ontology-agnostic coloring (root, ancestor, focus, descendant, etc.).
+    """
+    result = await service.build_entity_graph(
+        ontology_id,
+        class_iri,
+        branch=branch,
+        ancestors_depth=ancestors_depth,
+        descendants_depth=descendants_depth,
+        max_nodes=max_nodes,
+        include_see_also=include_see_also,
+    )
+    if result is None:
+        raise HTTPException(status_code=404, detail="Class not found")
+    return result
+
+
 @router.get("/ontologies/{ontology_id}/classes/{class_iri:path}", response_model=OWLClassResponse)
 async def get_class(
     ontology_id: UUID,
@@ -97,26 +131,3 @@ async def delete_class(
     deleted = await service.delete_class(ontology_id, class_iri)
     if not deleted:
         raise HTTPException(status_code=404, detail="Class not found")
-
-
-@router.get("/ontologies/{ontology_id}/classes/{class_iri:path}/hierarchy")
-async def get_class_hierarchy(
-    ontology_id: UUID,
-    class_iri: str,
-    service: Annotated[OntologyService, Depends(get_ontology_service)],
-    direction: str = "both",
-    depth: int = 3,
-) -> dict[str, object]:
-    """
-    Get the class hierarchy around a specific class.
-
-    Args:
-        direction: 'ancestors', 'descendants', or 'both'
-        depth: Maximum depth to traverse
-    """
-    return await service.get_class_hierarchy(
-        ontology_id,
-        class_iri,
-        direction=direction,
-        depth=depth,
-    )

ontokit/api/routes/projects.py

@@ -27,6 +27,7 @@
 from ontokit.models.branch_metadata import BranchMetadata
 from ontokit.models.pull_request import GitHubIntegration, PRStatus, PullRequest
 from ontokit.models.user_github_token import UserGitHubToken
+from ontokit.schemas.graph import EntityGraphResponse
 from ontokit.schemas.owl_class import EntitySearchResponse, OWLClassResponse, OWLClassTreeResponse
 from ontokit.schemas.project import (
     BranchCreate,
@@ -629,6 +630,47 @@ async def get_ontology_tree_children(
     return OWLClassTreeResponse(nodes=nodes, total_classes=total_classes)
 
 
+@router.get(
+    "/{project_id}/ontology/classes/graph",
+    response_model=EntityGraphResponse,
+)
+async def get_ontology_class_graph(
+    project_id: UUID,
+    service: Annotated[ProjectService, Depends(get_service)],
+    ontology: Annotated[OntologyService, Depends(get_ontology)],
+    git: Annotated[GitRepositoryService, Depends(get_git)],
+    user: OptionalUser,
+    class_iri: str = Query(description="IRI of the class to build the graph around"),
+    branch: str | None = Query(default=None, description="Branch to read from"),
+    ancestors_depth: int = Query(default=5, ge=0, le=10),
+    descendants_depth: int = Query(default=2, ge=0, le=10),
+    max_nodes: int = Query(default=200, ge=1, le=500),
+    include_see_also: bool = Query(default=True),
+) -> EntityGraphResponse:
+    """Build a multi-hop entity graph around a class via BFS.
+
+    Returns nodes and edges for visualization, with lineage-based node types.
+    """
+    resolved_branch = branch or git.get_default_branch(project_id)
+    await _ensure_ontology_loaded(project_id, service, ontology, user, resolved_branch, git)
+
+    result = await ontology.build_entity_graph(
+        project_id,
+        class_iri,
+        branch=resolved_branch,
+        ancestors_depth=ancestors_depth,
+        descendants_depth=descendants_depth,
+        max_nodes=max_nodes,
+        include_see_also=include_see_also,
+    )
+    if result is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Class not found: {class_iri}",
+        )
+    return result
+
+
 @router.get("/{project_id}/ontology/classes/{class_iri:path}", response_model=OWLClassResponse)
 async def get_ontology_class(
     project_id: UUID,

ontokit/schemas/graph.py

@@ -0,0 +1,40 @@
+"""Pydantic models for the Entity Graph API."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class GraphNode(BaseModel):
+    """A node in the entity graph."""
+
+    id: str
+    label: str
+    iri: str
+    definition: str | None = None
+    is_focus: bool = False
+    is_root: bool = False
+    depth: int = 0
+    node_type: str = "class"
+    child_count: int | None = None
+
+
+class GraphEdge(BaseModel):
+    """An edge in the entity graph."""
+
+    id: str
+    source: str
+    target: str
+    edge_type: str
+    label: str | None = None
+
+
+class EntityGraphResponse(BaseModel):
+    """Complete graph response."""
+
+    focus_iri: str
+    focus_label: str
+    nodes: list[GraphNode]
+    edges: list[GraphEdge]
+    truncated: bool = False
+    total_concept_count: int = 0