Harrison/param notion db (#4689)

Co-authored-by: Edward Park <ed.sh.park@gmail.com>
dynamic_agent_tools
Harrison Chase 1 year ago committed by GitHub
parent 5d63fc65e1
commit 6f47ab17a4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -1,6 +1,7 @@
{ {
"cells": [ "cells": [
{ {
"attachments": {},
"cell_type": "markdown", "cell_type": "markdown",
"id": "1dc7df1d", "id": "1dc7df1d",
"metadata": {}, "metadata": {},
@ -99,7 +100,11 @@
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "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",
")"
] ]
}, },
{ {

@ -1,6 +1,6 @@
"""Notion DB loader for langchain""" """Notion DB loader for langchain"""
from typing import Any, Dict, List from typing import Any, Dict, List, Optional
import requests import requests
@ -19,9 +19,15 @@ class NotionDBLoader(BaseLoader):
Args: Args:
integration_token (str): Notion integration token. integration_token (str): Notion integration token.
database_id (str): Notion database id. 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.""" """Initialize with parameters."""
if not integration_token: if not integration_token:
raise ValueError("integration_token must be provided") raise ValueError("integration_token must be provided")
@ -35,6 +41,7 @@ class NotionDBLoader(BaseLoader):
"Content-Type": "application/json", "Content-Type": "application/json",
"Notion-Version": "2022-06-28", "Notion-Version": "2022-06-28",
} }
self.request_timeout_sec = request_timeout_sec
def load(self) -> List[Document]: def load(self) -> List[Document]:
"""Load documents from the Notion database. """Load documents from the Notion database.
@ -148,7 +155,7 @@ class NotionDBLoader(BaseLoader):
url, url,
headers=self.headers, headers=self.headers,
json=query_dict, json=query_dict,
timeout=10, timeout=self.request_timeout_sec,
) )
res.raise_for_status() res.raise_for_status()
return res.json() return res.json()

Loading…
Cancel
Save