fix: correct the input params of "load_data" in NotionPageReader (#17529

)

llama-index-integrations/tools/llama-index-tools-notion/llama_index/tools/notion/base.py

@@ -54,15 +54,17 @@ def get_fn_schema_from_fn_name(
             raise ValueError(f"Invalid function name: {fn_name}")
 
     def load_data(
-        self, page_ids: Optional[List[str]] = None, database_id: Optional[str] = None
+        self,
+        page_ids: Optional[List[str]] = None,
+        database_ids: Optional[List[str]] = None,
     ) -> str:
-        """Loads content from a set of page ids or a database id.
+        """Loads content from a set of page ids or database ids.
 
-        Don't use this endpoint if you don't know the page ids or database id.
+        Don't use this endpoint if you don't know the page ids or database ids.
 
         """
         page_ids = page_ids or []
-        docs = self.reader.load_data(page_ids=page_ids, database_id=database_id)
+        docs = self.reader.load_data(page_ids=page_ids, database_ids=database_ids)
         return "\n".join([doc.get_content() for doc in docs])
 
     def search_data(
@@ -73,11 +75,12 @@ def search_data(
         value: Optional[str] = None,
         property: Optional[str] = None,
         page_size: int = 100,
-    ) -> str:
+    ) -> List[Dict[str, Any]]:
         """Search a list of relevant pages.
 
         Contains metadata for each page (but not the page content).
-
+        params:
+           query: the title of the page or database to search for, which is fuzzy matched.
         """
         payload: Dict[str, Any] = {
             "query": query,

llama-index-integrations/tools/llama-index-tools-notion/pyproject.toml

@@ -28,7 +28,7 @@ license = "MIT"
 maintainers = ["jerryjliu"]
 name = "llama-index-tools-notion"
 readme = "README.md"
-version = "0.3.0"
+version = "0.3.1"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"

llama-index-integrations/tools/llama-index-tools-notion/tests/test_tools_notion.py

@@ -1,7 +1,56 @@
+import os
+
+import pytest
 from llama_index.core.tools.tool_spec.base import BaseToolSpec
 from llama_index.tools.notion import NotionToolSpec
 
+# Get yourself a page id and database id from your notion account
+# Refer to the page: https://developers.notion.com/docs/create-a-notion-integration#give-your-integration-page-permissions
+
+page_ids = ["17d66c19670f80c5aaddfb8a0a449179"]  # replace with your page id
+database_ids = ["16066c19-670f-801d-adb8-fa9d1cdaa053"]  # replace with your database id
+
 
 def test_class():
     names_of_base_classes = [b.__name__ for b in NotionToolSpec.__mro__]
     assert BaseToolSpec.__name__ in names_of_base_classes
+
+
+@pytest.mark.skipif(
+    "NOTION_INTEGRATION_TOKEN" not in os.environ,
+    reason="NOTION_INTEGRATION_TOKEN is not set",
+)
+def test_load_data_with_page_ids():
+    tool = NotionToolSpec()
+    content = tool.load_data(page_ids=page_ids)
+    assert content
+
+
+@pytest.mark.skipif(
+    "NOTION_INTEGRATION_TOKEN" not in os.environ,
+    reason="NOTION_INTEGRATION_TOKEN is not set",
+)
+def test_load_data_with_database_ids():
+    tool = NotionToolSpec()
+    content = tool.load_data(database_ids=database_ids)
+    assert content
+
+
+@pytest.mark.skipif(
+    "NOTION_INTEGRATION_TOKEN" not in os.environ,
+    reason="NOTION_INTEGRATION_TOKEN is not set",
+)
+def test_load_data_with_page_ids_and_database_ids():
+    tool = NotionToolSpec()
+    content = tool.load_data(page_ids=page_ids, database_ids=database_ids)
+    assert content
+
+
+@pytest.mark.skipif(
+    "NOTION_INTEGRATION_TOKEN" not in os.environ,
+    reason="NOTION_INTEGRATION_TOKEN is not set",
+)
+def test_search_data():
+    tool = NotionToolSpec()
+    result = tool.search_data(query="Website")  # replace with your search query
+    assert len(result) > 0