chore: add backoff and retry logic to paperstack wflow

Author: GangGreenTemperTatum

.DS_Store

@@ -4,35 +4,39 @@ on:
   workflow_call:
     inputs:
       search-arxiv:
-        description: 'Search Arxiv?'
+        description: "Search Arxiv?"
         required: true
-        type: 'boolean'
+        type: "boolean"
       search-scholar:
-        description: 'Search Semantic Scholar?'
+        description: "Search Semantic Scholar?"
         required: true
-        type: 'boolean'
+        type: "boolean"
+  push:
+    branches:
+      - "ads/eng-1935-paperstack-1-notion-api-timeouts-causing-paperstack"
 
 jobs:
   run-script:
     runs-on: ubuntu-latest
+    timeout-minutes: 120 # Adding 2-hour timeout for the entire job
     steps:
-    - name: Check out code
-      uses: actions/checkout@v2
+      - name: Check out code
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 #v4.2.2
 
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.x'
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 #v5.6.0
+        with:
+          python-version: "3.x"
 
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install -r requirements.txt
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
 
-    - name: Run paperstack
-      run: |
-        python paperstack.py ${{ inputs.search-arxiv && '--search-arxiv' || '' }} ${{ inputs.search-scholar && '--search-semantic-scholar' || '' }}
-      env:
-        NOTION_TOKEN: ${{ secrets.NOTION_TOKEN }}
-        NOTION_DATABASE_ID: ${{ secrets.NOTION_DATABASE_ID }}
-        OPENAI_API_TOKEN: ${{ secrets.OPENAI_API_TOKEN }}
+      - name: Run paperstack
+        run: |
+          python paperstack.py ${{ inputs.search-arxiv && '--search-arxiv' || '' }} ${{ inputs.search-scholar && '--search-semantic-scholar' || '' }}
+        env:
+          NOTION_TOKEN: ${{ secrets.NOTION_TOKEN }}
+          NOTION_DATABASE_ID: ${{ secrets.NOTION_DATABASE_ID }}
+          OPENAI_API_TOKEN: ${{ secrets.OPENAI_API_TOKEN }}

.github/workflows/core.yml

@@ -158,3 +158,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+notion_utils.py

.gitignore

@@ -1,28 +1,52 @@
 import asyncio
 import typing as t
 from datetime import datetime
+import time
 
 from notion_client import AsyncClient
+from notion_client.errors import RequestTimeoutError, APIResponseError
 from notion_client.helpers import async_collect_paginated_api
 from tqdm import tqdm  # type: ignore
 
 from _types import AttackType, Paper, Focus
 
+# Retry constants
+MAX_RETRIES = 5
+RETRY_DELAY = 5
+MAX_BATCH_SIZE = 5
+
 NotionClient = AsyncClient
 
 
 def get_notion_client(token: str) -> NotionClient:
-    return NotionClient(auth=token)
+    return NotionClient(auth=token, timeout_ms=60000)  # 60-second timeout
 
 
 async def get_papers_from_notion(client: NotionClient, database_id: str, *, max: int | None = None) -> list[Paper]:
-    if max:
-        results = await client.databases.query(database_id=database_id, page_size=max)
-        results = results['results']
-    else:
-        results = await async_collect_paginated_api(
-            client.databases.query, database_id=database_id
-        )
+    retries = 0
+    results = []
+
+    while retries < MAX_RETRIES:
+        try:
+            if max:
+                response = await client.databases.query(database_id=database_id, page_size=max)
+                results = response['results']
+            else:
+                results = await async_collect_paginated_api(
+                    client.databases.query, database_id=database_id
+                )
+            break
+        except (RequestTimeoutError, APIResponseError) as e:
+            retries += 1
+            if retries >= MAX_RETRIES:
+                print(f"Failed to get papers from Notion after {MAX_RETRIES} attempts: {str(e)}")
+                return []
+            else:
+                print(f"Notion API error when fetching papers, retrying ({retries}/{MAX_RETRIES}): {str(e)}")
+                # Exponential backoff with jitter
+                wait_time = RETRY_DELAY * (2 ** (retries - 1)) + (RETRY_DELAY * 0.1 * retries)
+                print(f"Waiting {wait_time:.1f} seconds before retry...")
+                await asyncio.sleep(wait_time)
 
     papers: list[Paper] = []
     for result in results:
@@ -67,34 +91,64 @@ async def get_papers_from_notion(client: NotionClient, database_id: str, *, max:
 async def write_papers_to_notion(
     client: NotionClient, database_id: str, papers: list[Paper]
 ) -> None:
-    for paper in tqdm(papers):
-        properties: dict[str, t.Any] = {}
-        if paper.title and paper._original_state["title"] != paper.title:
-            properties["Title"] = {"title": [{"text": {"content": paper.title}}]}
-        if paper.url and paper._original_state["url"] != paper.url:
-            properties["URL"] = {"url": paper.url}
-        if paper.summary and paper._original_state["summary"] != paper.summary:
-            properties["Summary"] = {
-                "rich_text": [{"text": {"content": paper.summary}}]
-            }
-        if paper.authors and paper._original_state["authors"] != paper.authors:
-            properties["Authors"] = {
-                "multi_select": [{"name": author} for author in paper.authors[:5]] # Limit to 5 authors
-            }
-        if paper.published and paper._original_state["published"] != paper.published:
-            properties["Published"] = {"date": {"start": paper.published.isoformat()}}
-        if paper.focus and paper._original_state["focus"] != paper.focus:
-            properties["Focus"] = {"select": {"name": paper.focus.value}}
-        if paper.attack_type and paper._original_state["attack_type"] != paper.attack_type:
-            properties["Attack Type"] = {"select": {"name": paper.attack_type.value}}
-        if paper.explored and paper._original_state["explored"] != paper.explored:
-            properties["Explored"] = {"checkbox": paper.explored}
-
-        if paper.page_id:
-            await client.pages.update(paper.page_id, properties=properties)
-        else:
-            await client.pages.create(
-                parent={"database_id": database_id}, properties=properties
-            )
+    # Process papers in smaller batches with pauses between
+    for i in range(0, len(papers), MAX_BATCH_SIZE):
+        batch = papers[i:i+MAX_BATCH_SIZE]
+        print(f"Processing batch {i//MAX_BATCH_SIZE + 1}/{(len(papers) + MAX_BATCH_SIZE - 1)//MAX_BATCH_SIZE}")
+
+        for paper in tqdm(batch):
+            properties: dict[str, t.Any] = {}
+            if paper.title:
+                properties["Title"] = {"title": [{"text": {"content": paper.title}}]}
+            if paper.url:
+                properties["URL"] = {"url": paper.url}
+            if paper.summary:
+                properties["Summary"] = {
+                    "rich_text": [{"text": {"content": paper.summary}}]
+                }
+            if paper.authors:
+                properties["Authors"] = {
+                    "multi_select": [{"name": author} for author in paper.authors]
+                }
+            if paper.published:
+                properties["Published"] = {"date": {"start": paper.published.isoformat()}}
+            if paper.focus:
+                properties["Focus"] = {"select": {"name": paper.focus.value}}
+            if paper.attack_type:
+                properties["Attack Type"] = {"select": {"name": paper.attack_type.value}}
+            if paper.explored is not None:
+                properties["Explored"] = {"checkbox": paper.explored}
+
+            # Retry logic with progressive backoff
+            retries = 0
+            while retries < MAX_RETRIES:
+                try:
+                    if paper.page_id:
+                        await client.pages.update(paper.page_id, properties=properties)
+                    else:
+                        await client.pages.create(
+                            parent={"database_id": database_id}, properties=properties
+                        )
+                    # Success, break out of retry loop
+                    break
+                except (RequestTimeoutError, APIResponseError) as e:
+                    retries += 1
+                    if retries >= MAX_RETRIES:
+                        print(f"Failed to update/create paper after {MAX_RETRIES} attempts: {paper.title[:50]}...")
+                        # Don't raise - continue with other papers
+                        break
+                    else:
+                        print(f"Notion API error, retrying ({retries}/{MAX_RETRIES}): {str(e)}")
+                        # Exponential backoff with longer delays
+                        wait_time = RETRY_DELAY * (2 ** (retries - 1)) + (RETRY_DELAY * 0.1 * retries)
+                        print(f"Waiting {wait_time:.1f} seconds before retry...")
+                        await asyncio.sleep(wait_time)
+
+            # Add a small delay between papers regardless of success/failure
+            await asyncio.sleep(1)
+
+        if i + MAX_BATCH_SIZE < len(papers):
+            print(f"Pausing for 10 seconds between batches...")
+            await asyncio.sleep(10)
 
     return None