feat: Add edge attributes fetching and schema retrieval (#52)

Author: j6k4m8

motifstudio-web/src/app/GraphStats.tsx

@@ -58,6 +58,17 @@ export function GraphStats({
             host_id: graph?.id,
         })
     );
+    const {
+        data: edgeAttrData,
+        error: edgeAttrError,
+        isLoading: edgeAttrIsLoading,
+    } = useSWR<{
+        attributes: { [key: string]: string };
+    }>([`${BASE_URL}/queries/edges/attributes`, graph?.id], () =>
+        bodiedFetcher(`${BASE_URL}/queries/edges/attributes`, {
+            host_id: graph?.id,
+        })
+    );
 
     // To handle the fact that the attributes are loaded asynchronously, we
     // provide a callback function to the parent component to share the
@@ -70,8 +81,8 @@ export function GraphStats({
 
     // Use client-only check to avoid hydration mismatch
     if (!isClient) return <div>Loading...</div>;
-    if (vertIsLoading || edgeIsLoading) return <div>Loading...</div>;
-    if (vertError || edgeError) return <div>Error: {vertError}</div>;
+    if (vertIsLoading || edgeIsLoading || vertAttrIsLoading || edgeAttrIsLoading) return <div>Loading...</div>;
+    if (vertError || edgeError || vertAttrError || edgeAttrError) return <div>Error: {vertError || edgeError || vertAttrError || edgeAttrError}</div>;
     if (!vertData || !edgeData) return <div>No data</div>;
 
     /**
@@ -153,13 +164,30 @@ export function GraphStats({
                 <div className="flex gap-2">
                     {vertAttrData?.attributes
                         ? Object.entries(vertAttrData?.attributes).map(([key, value]) => (
-                              <span
-                                  key={key}
-                                  className="px-2 py-1 bg-blue-50 rounded-md shadow-sm text-sm font-medium text-blue-800"
-                              >
-                                  {key} <b className="font-mono">({value})</b>
-                              </span>
-                          ))
+                            <span
+                                key={key}
+                                className="px-2 py-1 bg-blue-50 rounded-md shadow-sm text-sm font-medium text-blue-800"
+                            >
+                                {key} <b className="font-mono">({value})</b>
+                            </span>
+                        ))
+                        : null}
+                </div>
+
+                <hr className="my-2 w-full" />
+
+                {/* Edge attributes list */}
+                <h3 className="text-lg font-mono w-full">Edge Attributes</h3>
+                <div className="flex gap-2">
+                    {edgeAttrData?.attributes
+                        ? Object.entries(edgeAttrData?.attributes).map(([key, value]) => (
+                            <span
+                                key={key}
+                                className="px-2 py-1 bg-green-50 rounded-md shadow-sm text-sm font-medium text-green-800"
+                            >
+                                {key} <b className="font-mono">({value})</b>
+                            </span>
+                        ))
                         : null}
                 </div>
 

server/src/host_provider/host_provider/host_provider.py

@@ -72,6 +72,18 @@ def get_edge_count(self, uri: str) -> int:
         """
         ...
 
+    def get_edge_attribute_schema(self, uri: str) -> AttributeSchema:
+        """Return the schema of the edge attributes in the graph.
+
+        Arguments:
+            uri (str): The URI of the host.
+
+        Returns:
+            dict[str, str]: The schema of the edge attributes in the graph.
+
+        """
+        ...
+
     def get_motif_count(self, uri: str, motif_string: str) -> int:
         """Get a count of motifs in the graph.
 
@@ -175,6 +187,31 @@ def get_edge_count(self, uri: str) -> int:
         """
         return len(self.get_networkx_graph(uri).edges)
 
+    def get_edge_attribute_schema(self, uri: str) -> dict[str, str]:
+        """Return the schema of the edge attributes in the graph.
+
+        Arguments:
+            uri (str): The URI of the host.
+
+        Returns:
+            dict[str, str]: The schema of the edge attributes in the graph.
+
+        """
+        g = self.get_networkx_graph(uri)
+        # TODO: This exhaustive search is no good for very large graphs.
+
+        # Detect types of attributes, which may be different for different
+        # edges.
+        attribute_types = {}
+        for edge in g.edges:
+            for attribute in g.edges[edge].keys():
+                if attribute not in attribute_types:
+                    attribute_types[attribute] = type(g.edges[edge][attribute]).__name__
+                else:
+                    if attribute_types[attribute] != type(g.edges[edge][attribute]).__name__:
+                        attribute_types[attribute] = "str"
+        return attribute_types
+
     def get_motif_count(self, uri: str, motif_string: str) -> int:
         """Count the number of instances of a motif in the graph.
 

server/src/models.py

@@ -98,6 +98,19 @@ class EdgeCountQueryResponse(_QueryResponseBase):
     edge_count: int
 
 
+class EdgeAttributeQueryRequest(_QueryRequestBase):
+    """A request to get the edge attributes for a host graph."""
+
+    ...
+
+
+class EdgeAttributeQueryResponse(_QueryResponseBase):
+    """A response with the edge attribute results for a host graph."""
+
+    # Attribute name to attribute schema:
+    attributes: AttributeSchema
+
+
 class MotifParseQueryRequest(_QueryRequestBase):
     """A request to parse a motif query."""
 

server/src/server/routers/queries.py

@@ -13,6 +13,8 @@
 from ...models import (
     EdgeCountQueryRequest,
     EdgeCountQueryResponse,
+    EdgeAttributeQueryRequest,
+    EdgeAttributeQueryResponse,
     MotifCountQueryRequest,
     MotifCountQueryResponse,
     MotifParseQueryRequest,
@@ -222,6 +224,36 @@ def query_count_edges(
     )
 
 
+@router.post("/edges/attributes")
+def query_edge_attributes(
+    edge_attribute_query_request: EdgeAttributeQueryRequest,
+    commons: Annotated[HostProviderRouterGlobalDep, Depends(provider_router)],
+) -> EdgeAttributeQueryResponse:
+    """Get the edge attributes for a given host."""
+    tic = time.time()
+    uri = commons.get_uri_from_id(edge_attribute_query_request.host_id)
+    if uri is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No host found with ID {edge_attribute_query_request.host_id}",
+        )
+
+    provider = commons.host_provider_router.provider_for(uri)
+    if provider is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No provider found for host {edge_attribute_query_request.host_id}",
+        )
+
+    attributes = provider.get_edge_attribute_schema(uri)
+    return EdgeAttributeQueryResponse(
+        attributes=attributes,
+        host_id=edge_attribute_query_request.host_id,
+        response_time=datetime.datetime.now().isoformat(),
+        response_duration_ms=(time.time() - tic) * 1000,
+    )
+
+
 @router.post("/motifs/count")
 def query_count_motifs(
     motif_count_query_request: MotifCountQueryRequest,