request working

src/app.py

@@ -1,10 +1,26 @@
 import streamlit as st
 import asyncio
-import aiohttp
-from bs4 import BeautifulSoup
-from typing import Dict, Any, Optional, List, Set, cast
+import websockets
+import json
+import platform
+from typing import Dict, Any, Optional, List, Literal, TypedDict, cast
 
-CacheType = Dict[int, Dict[str, Any]]
+GGRAPHER_URI = "wss://ggrphr.davidalber.net"
+
+class StartNodeRequest(TypedDict):
+    recordId: int
+    getAdvisors: bool
+    getDescendants: bool
+
+class RequestPayload(TypedDict):
+    kind: Literal["build-graph"]
+    options: Dict[Literal["reportingCallback"], bool]
+    startNodes: List[StartNodeRequest]
+
+class ProgressCallback(TypedDict):
+    queued: int
+    fetching: int
+    done: int
 
 def get_id_from_input(val: str) -> Optional[int]:
     try:
@@ -12,186 +28,67 @@ def get_id_from_input(val: str) -> Optional[int]:
     except Exception:
         return None
 
-async def fetch_record(
-    record_id: int,
-    session: aiohttp.ClientSession,
-    cache: Optional[CacheType],
-    delay: float = 0.0,
-) -> Optional[Dict[str, Any]]:
-    if cache is not None and record_id in cache:
-        return cache[record_id]
-    url = f"https://www.mathgenealogy.org/id.php?id={record_id}"
-    async with session.get(url) as resp:
-        html = await resp.text()
-    if delay > 0:
-        await asyncio.sleep(delay)
-    soup = BeautifulSoup(html, "html.parser")
-    if soup.string == "Non-numeric id supplied. Aborting.":
-        return None
-    if (
-        soup.p is not None
-        and soup.p.string
-        == "You have specified an ID that does not exist in the database. Please back up and try again."
-    ):
-        return None
-    h2 = soup.find("h2")
-    name = h2.get_text(strip=True) if h2 and hasattr(h2, "get_text") else str(record_id)
-    institution = None
-    year = None
-    for inst in soup.find_all("div", style="line-height: 30px; text-align: center; margin-bottom: 1ex"):
-        if hasattr(inst, "find"):
-            span1 = inst.find("span")
-            if span1 and hasattr(span1, "find"):
-                span2 = span1.find("span")
-                if span2 and hasattr(span2, "text"):
-                    institution = span2.text
-            if span1 and hasattr(span1, "contents") and span1.contents:
-                y = span1.contents[-1]
-                if isinstance(y, str):
-                    y = y.strip()
-                    if y:
-                        y = y.split(",")[0].strip()
-                        if y.isdigit():
-                            year = int(y)
-    advisors = []
-    for a in soup.find_all(string=lambda s: s and ("Advisor" in cast(str, s) or "Promotor" in cast(str, s))):
-        # Ensure 'a' itself is treated as a string for the 'in' check
-        if "Advisor: Unknown" in str(a):
-            continue
-        next_tag = a.find_next() if hasattr(a, "find_next") else None
-        if next_tag and hasattr(next_tag, "attrs") and "href" in next_tag.attrs:
-            try:
-                advisors.append(int(next_tag.attrs["href"].split("=")[-1]))
-            except Exception:
-                pass
-    table = soup.find("table")
-    descendants = []
-    if table and hasattr(table, "find_all"):
-        for a in table.find_all("a"):
-            if hasattr(a, "attrs") and "href" in a.attrs:
-                try:
-                    descendants.append(int(a.attrs["href"].split("=")[-1]))
-                except Exception:
-                    pass
-    record = {
-        "id": record_id,
-        "name": name,
-        "institution": institution,
-        "year": year,
-        "advisors": advisors,
-        "descendants": descendants,
+def make_payload(record_id: int) -> RequestPayload:
+    return {
+        "kind": "build-graph",
+        "options": {"reportingCallback": True},
+        "startNodes": [{
+            "recordId": record_id,
+            "getAdvisors": True,
+            "getDescendants": False,
+        }],
     }
-    if cache is not None:
-        cache[record_id] = record
-    return record
 
-async def fetch_advisors_parallel(
-    rid: int,
-    session: aiohttp.ClientSession,
-    cache: Optional[CacheType],
-    semaphore: asyncio.Semaphore,
-    tree: Dict[int, Dict[str, Any]],
-    progress_cb,
-    visited: Set[int],
-    delay: float = 0.0,
-):
-    if rid in visited:
-        return
-    visited.add(rid)
-    async with semaphore:
-        rec = await fetch_record(rid, session, cache, delay)
-    if rec is not None:
-        tree[rid] = rec
-        await progress_cb(tree)
-        tasks = [
-            fetch_advisors_parallel(
-                adv, session, cache, semaphore, tree, progress_cb, visited, delay
-            )
-            for adv in rec["advisors"]
-        ]
-        if tasks:
-            await asyncio.gather(*tasks)
+async def get_graph(payload: RequestPayload, progress_cb=None) -> Dict[str, Any]:
+    def intify_record_keys(d: Dict[Any, Any]) -> Dict[Any, Any]:
+        if "nodes" in d:
+            ret = {k: v for k, v in d.items() if k != "nodes"}
+            ret["nodes"] = {int(k): v for k, v in d["nodes"].items()}
+            return ret
+        return d
 
-def tree_to_dot(tree: Dict[int, Dict[str, Any]]) -> str:
+    async with websockets.connect(  # type: ignore[attr-defined]
+        GGRAPHER_URI,
+    ) as ws:
+        await ws.send(json.dumps(payload))
+        while True:
+            response_json = await ws.recv()
+            response = json.loads(response_json, object_hook=intify_record_keys)
+            response_payload = response.get("payload")
+            if response["kind"] == "graph":
+                return cast(Dict[str, Any], response_payload)
+            elif response["kind"] == "progress" and progress_cb:
+                progress = cast(ProgressCallback, response_payload)
+                progress_cb(progress)
+            else:
+                continue
+
+def tree_to_dot(graph: Dict[str, Any]) -> str:
+    nodes = graph.get("nodes", {})
     lines = [
         "digraph G {",
-        "  rankdir=TB; // Top to Bottom layout, could also be LR (Left to Right)",
-        '  node [shape=box, style="rounded,filled", fillcolor=lightyellow]; // Optional: style nodes',
+        "  rankdir=TB;",
+        '  node [shape=box, style="rounded,filled", fillcolor=lightyellow];',
         '  edge [arrowhead=vee];'
     ]
-
-    # Group nodes by year
-    nodes_by_year: Dict[Optional[int], List[int]] = {}
-    for node_id, node_data in tree.items():
-        year = node_data.get("year") # Might be None
-        if year not in nodes_by_year:
-            nodes_by_year[year] = []
-        nodes_by_year[year].append(node_id)
-
     # Define nodes and their labels
-    for node_id, node in tree.items():
+    for node_id, node in nodes.items():
         name = node.get("name", str(node_id))
         year_str = f" ({node.get('year')})" if node.get('year') is not None else " (Year Unknown)"
         label = f"{name}{year_str}"
-        # Add a tooltip with more info if possible (Graphviz might not render HTML tooltips in all viewers)
-        tooltip = f"ID: {node_id}\nName: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}"
+        tooltip = f"ID: {node_id}\\nName: {name}\\nYear: {node.get('year', 'N/A')}\\nInstitution: {node.get('institution', 'N/A')}"
         lines.append(f'  "{node_id}" [label="{label}", tooltip="{tooltip}"];')
-
     # Define edges
-    for node_id, node in tree.items():
+    for node_id, node in nodes.items():
         for adv_id in node.get("advisors", []):
-            if adv_id in tree:
+            if adv_id in nodes:
                 lines.append(f'  "{adv_id}" -> "{node_id}";')
-
-    # Add rank constraints for years
-    # Sort years, placing None (unknown year) perhaps at the top or bottom.
-    # For this example, None years will not be part of explicit rank=same groups.
-    # Dot will place them based on connections.
-    # Consider creating a specific rank for "Unknown Year" if desired.
-    
-    # Sort known years.
-    # Years are typically displayed with earlier years at the top (if rankdir=TB)
-    # or left (if rankdir=LR). Dot usually handles this naturally with directed edges.
-    # The rank=same constraint is more about aligning nodes *within* the same year.
-    sorted_known_years = sorted([y for y in nodes_by_year.keys() if y is not None])
-
-    for year in sorted_known_years:
-        nodes_in_year = nodes_by_year[year]
-        if len(nodes_in_year) > 0: # Create a rank group even for single nodes if you want to ensure year separation
-            # Invisible subgraph for ranking
-            lines.append(f"  subgraph cluster_year_{year} {{") # Naming cluster helps visually if you add style
-            lines.append("    label=\"\"; // No visible label for the cluster itself")
-            lines.append("    style=invis;")
-            lines.append("    {")
-            lines.append("      rank=same;")
-            # Add nodes to the rank group
-            for node_id in nodes_in_year:
-                lines.append(f'      "{node_id}";')
-            lines.append("    }")
-            lines.append("  }")
-            
-    # Handle nodes with unknown years (year is None)
-    # Option 1: Let them float (current behavior if not explicitly ranked)
-    # Option 2: Group them in their own rank (e.g., at the top or bottom)
-    if None in nodes_by_year and nodes_by_year[None]:
-        lines.append("  subgraph cluster_year_unknown {")
-        lines.append("    label=\"\"; style=invis;")
-        lines.append("    {")
-        lines.append("      rank=same; // Or rank=min/max if you want them at very top/bottom")
-        for node_id in nodes_by_year[None]:
-            lines.append(f'      "{node_id}";')
-        lines.append("    }")
-        lines.append("  }")
-
     lines.append("}")
     return "\n".join(lines)
 
 def main():
-    st.title("Math Genealogy Ancestor Tree (Async Parallel Scrape)")
+    st.title("Math Genealogy Ancestor Tree (WebSocket API)")
     mgp_id_str = st.text_input("Enter MGP ID (integer):")
-    use_cache = st.checkbox("Use in-memory cache", value=True)
-    concurrency = st.slider("Max parallel requests", min_value=1, max_value=10, value=5)
-    delay = st.slider("Delay between requests (seconds)", min_value=0.0, max_value=2.0, value=0.2, step=0.05)
     progress_placeholder = st.empty()
     graph_placeholder = st.empty()
     run_btn = st.button("Show Ancestor Tree")
@@ -200,28 +97,22 @@ def main():
         if mgp_id is None:
             st.error("Please enter a valid integer MGP ID.")
             return
-        cache: Optional[CacheType] = {} if use_cache else None
+        payload = make_payload(mgp_id)
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-        progress_nodes: List[int] = []
-        async def progress_cb(tree):
-            progress_nodes.clear()
-            progress_nodes.extend(tree.keys())
-            progress_placeholder.info(f"Ancestors found: {len(tree)}")
-            dot = tree_to_dot(tree)
-            graph_placeholder.graphviz_chart(dot)
+        def progress_cb(progress):
+            progress_placeholder.info(
+                f"Queued: {progress['queued']} | Fetching: {progress['fetching']} | Done: {progress['done']}"
+            )
         async def runner():
-            async with aiohttp.ClientSession() as session:
-                semaphore = asyncio.Semaphore(concurrency)
-                tree: Dict[int, Dict[str, Any]] = {}
-                visited: Set[int] = set()
-                await fetch_advisors_parallel(
-                    mgp_id, session, cache, semaphore, tree, progress_cb, visited, delay
-                )
+            graph = await get_graph(payload, progress_cb)
+            dot = tree_to_dot(graph)
+            graph_placeholder.graphviz_chart(dot)
         try:
             loop.run_until_complete(runner())
             progress_placeholder.success("Done!")
         except Exception as e:
+            print(f"Error: {e}")
             progress_placeholder.error(f"Error: {e}")
 
 if __name__ == "__main__":