From 6f47ab17a43a71ce7b09469ffbb5a9820fde8614 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sun, 14 May 2023 18:26:25 -0700
Subject: [PATCH] Harrison/param notion db (#4689)

Co-authored-by: Edward Park <ed.sh.park@gmail.com>
---
 .../document_loaders/examples/notiondb.ipynb        |  7 ++++++-
 langchain/document_loaders/notiondb.py              | 13 ++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/docs/modules/indexes/document_loaders/examples/notiondb.ipynb b/docs/modules/indexes/document_loaders/examples/notiondb.ipynb
index 28a6c093..eb18b918 100644
--- a/docs/modules/indexes/document_loaders/examples/notiondb.ipynb
+++ b/docs/modules/indexes/document_loaders/examples/notiondb.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "1dc7df1d",
    "metadata": {},
@@ -99,7 +100,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "loader = NotionDBLoader(integration_token=NOTION_TOKEN, database_id=DATABASE_ID)"
+    "loader = NotionDBLoader(\n",
+    "    integration_token=NOTION_TOKEN, \n",
+    "    database_id=DATABASE_ID,\n",
+    "    request_timeout_sec=30 # optional, defaults to 10\n",
+    ")"
    ]
   },
   {
diff --git a/langchain/document_loaders/notiondb.py b/langchain/document_loaders/notiondb.py
index 25c72959..f43fd5f4 100644
--- a/langchain/document_loaders/notiondb.py
+++ b/langchain/document_loaders/notiondb.py
@@ -1,6 +1,6 @@
 """Notion DB loader for langchain"""
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import requests
 
@@ -19,9 +19,15 @@ class NotionDBLoader(BaseLoader):
     Args:
         integration_token (str): Notion integration token.
         database_id (str): Notion database id.
+        request_timeout_sec (int): Timeout for Notion requests in seconds.
     """
 
-    def __init__(self, integration_token: str, database_id: str) -> None:
+    def __init__(
+        self,
+        integration_token: str,
+        database_id: str,
+        request_timeout_sec: Optional[int] = 10,
+    ) -> None:
         """Initialize with parameters."""
         if not integration_token:
             raise ValueError("integration_token must be provided")
@@ -35,6 +41,7 @@ class NotionDBLoader(BaseLoader):
             "Content-Type": "application/json",
             "Notion-Version": "2022-06-28",
         }
+        self.request_timeout_sec = request_timeout_sec
 
     def load(self) -> List[Document]:
         """Load documents from the Notion database.
@@ -148,7 +155,7 @@ class NotionDBLoader(BaseLoader):
             url,
             headers=self.headers,
             json=query_dict,
-            timeout=10,
+            timeout=self.request_timeout_sec,
         )
         res.raise_for_status()
         return res.json()