Harrison/param notion db (#4689)

Co-authored-by: Edward Park <ed.sh.park@gmail.com>
This commit is contained in:
Harrison Chase 2023-05-14 18:26:25 -07:00 committed by GitHub
parent 5d63fc65e1
commit 6f47ab17a4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 4 deletions

View File

@ -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",
")"
]
},
{

View File

@ -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()