Update app.py

app.py

@@ -133,7 +133,7 @@ STRICT IDENTITY RULES:
 
 3. NATURAL CONVERSATION & TOOL RESPONSE STYLE:
    - Do NOT introduce your company, founder, or model name in general responses unless requested.
-   - Always summarize and format Google Connector execution results into clean, polite, and natural language automatically matching the language used by the user.
+   - Always summarize and format Google and Notion Connector execution results into clean, polite, and natural language automatically matching the language used by the user.
    - Never output raw technical logs, IDs, or developer headers directly to the user."""
 
 TOOLS = [
@@ -358,6 +358,65 @@ TOOLS = [
                 "required": ["file_id", "content"]
             }
         }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "notion_search",
+            "description": "Search Notion pages, databases, or documents by query",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "Search keyword or empty to list recent pages"},
+                    "max_results": {"type": "integer", "description": "Number of results to fetch, default 5"}
+                }
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "notion_get_page_content",
+            "description": "Read text content and blocks of a specific Notion page by page_id",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "page_id": {"type": "string", "description": "The ID of the Notion page"}
+                },
+                "required": ["page_id"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "notion_create_page",
+            "description": "Create a new page in Notion inside a parent page or database",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "parent_id": {"type": "string", "description": "Optional parent page or database ID"},
+                    "title": {"type": "string", "description": "Title of the new Notion page"},
+                    "content": {"type": "string", "description": "Initial text content for the page"}
+                },
+                "required": ["title"]
+            }
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "notion_append_block_children",
+            "description": "Append text or paragraph blocks to an existing Notion page by page_id",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "page_id": {"type": "string", "description": "The ID of the Notion page"},
+                    "content": {"type": "string", "description": "Text content to append to the page"}
+                },
+                "required": ["page_id", "content"]
+            }
+        }
     }
 ]
 
@@ -444,6 +503,157 @@ async def get_google_creds(user_id: str, service_name: str = "google"):
             pass
     return creds
 
+async def notion_get_headers(user_id: str):
+    token_data = await get_token_from_supabase(user_id, "notion")
+    if not token_data:
+        return None
+    access_token = token_data.get("access_token") or token_data.get("token")
+    if not access_token:
+        return None
+    return {
+        "Authorization": f"Bearer {access_token}",
+        "Notion-Version": "2022-06-28",
+        "Content-Type": "application/json"
+    }
+
+async def notion_search(user_id: str, query: str = "", max_results: int = 5):
+    headers = await notion_get_headers(user_id)
+    if not headers:
+        return "Notion is not connected. Please connect Notion account."
+    try:
+        url = "https://api.notion.com/v1/search"
+        payload = {"page_size": max_results}
+        if query:
+            payload["query"] = query
+        async with httpx.AsyncClient() as c:
+            res = await c.post(url, headers=headers, json=payload)
+            if res.status_code != 200:
+                return f"Error searching Notion: {res.text}"
+            data = res.json()
+            results = data.get("results", [])
+            if not results:
+                return "No pages or databases found in Notion."
+            items = []
+            for item in results:
+                item_id = item.get("id", "")
+                obj_type = item.get("object", "page")
+                title = "Untitled"
+                if obj_type == "page":
+                    props = item.get("properties", {})
+                    for k, v in props.items():
+                        if v.get("type") == "title":
+                            t_arr = v.get("title", [])
+                            if t_arr:
+                                title = "".join([t.get("plain_text", "") for t in t_arr])
+                            break
+                elif obj_type == "database":
+                    t_arr = item.get("title", [])
+                    if t_arr:
+                        title = "".join([t.get("plain_text", "") for t in t_arr])
+                url_val = item.get("url", "")
+                items.append(f"Page/DB ID: {item_id} | Type: {obj_type} | Title: {title} | URL: {url_val}")
+            return "\n".join(items)
+    except Exception as e:
+        return f"Error searching Notion: {str(e)}"
+
+async def notion_get_page_content(user_id: str, page_id: str):
+    headers = await notion_get_headers(user_id)
+    if not headers:
+        return "Notion is not connected."
+    try:
+        clean_id = page_id.replace("-", "").strip()
+        url = f"https://api.notion.com/v1/blocks/{clean_id}/children"
+        async with httpx.AsyncClient() as c:
+            res = await c.get(url, headers=headers)
+            if res.status_code != 200:
+                return f"Error reading Notion page content: {res.text}"
+            data = res.json()
+            results = data.get("results", [])
+            if not results:
+                return "Notion page is empty."
+            texts = []
+            for b in results:
+                b_type = b.get("type")
+                if b_type and b_type in b and "rich_text" in b[b_type]:
+                    rt = b[b_type].get("rich_text", [])
+                    t = "".join([x.get("plain_text", "") for x in rt])
+                    if t:
+                        texts.append(t)
+            return "\n".join(texts) if texts else "Page contains non-text blocks."
+    except Exception as e:
+        return f"Error reading Notion page: {str(e)}"
+
+async def notion_create_page(user_id: str, parent_id: str = None, title: str = "Untitled", content: str = ""):
+    headers = await notion_get_headers(user_id)
+    if not headers:
+        return "Notion is not connected."
+    try:
+        url = "https://api.notion.com/v1/pages"
+        payload = {
+            "properties": {
+                "title": {
+                    "title": [{"text": {"content": title}}]
+                }
+            }
+        }
+        if parent_id:
+            clean_parent = parent_id.replace("-", "").strip()
+            payload["parent"] = {"page_id": clean_parent}
+        else:
+            search_res = await notion_search(user_id, max_results=1)
+            if "Page/DB ID: " in search_res:
+                p_id = search_res.split("Page/DB ID: ")[1].split(" ")[0].strip()
+                payload["parent"] = {"page_id": p_id.replace("-", "")}
+            else:
+                return "Cannot find parent page. Please provide a parent_id to create a new page in Notion."
+
+        if content:
+            payload["children"] = [
+                {
+                    "object": "block",
+                    "type": "paragraph",
+                    "paragraph": {
+                        "rich_text": [{"type": "text", "text": {"content": content}}]
+                    }
+                }
+            ]
+
+        async with httpx.AsyncClient() as c:
+            res = await c.post(url, headers=headers, json=payload)
+            if res.status_code != 200:
+                return f"Error creating Notion page: {res.text}"
+            data = res.json()
+            new_id = data.get("id", "")
+            return f"Successfully created Notion page '{title}' (Page ID: {new_id})."
+    except Exception as e:
+        return f"Error creating Notion page: {str(e)}"
+
+async def notion_append_block_children(user_id: str, page_id: str, content: str):
+    headers = await notion_get_headers(user_id)
+    if not headers:
+        return "Notion is not connected."
+    try:
+        clean_id = page_id.replace("-", "").strip()
+        url = f"https://api.notion.com/v1/blocks/{clean_id}/children"
+        payload = {
+            "children": [
+                {
+                    "object": "block",
+                    "type": "paragraph",
+                    "paragraph": {
+                        "rich_text": [{"type": "text", "text": {"content": content}}]
+                    }
+                }
+            ]
+        }
+        async with httpx.AsyncClient() as c:
+            res = await c.patch(url, headers=headers, json=payload)
+            if res.status_code != 200:
+                return f"Error appending content to Notion: {res.text}"
+            return f"Successfully appended content to Notion page ID: {page_id}."
+    except Exception as e:
+        return f"Error appending content to Notion: {str(e)}"
+
 async def gmail_search_emails(user_id: str, query: str = "", max_results: int = 5):
     creds = await get_google_creds(user_id, "gmail")
     if not creds:
@@ -724,6 +934,14 @@ async def execute_tool_call(user_id: str, func_name: str, args: dict) -> str:
         return await drive_delete_file(user_id, args.get("file_id", ""))
     elif func_name == "drive_update_file":
         return await drive_update_file(user_id, args.get("file_id", ""), args.get("content", ""))
+    elif func_name == "notion_search":
+        return await notion_search(user_id, args.get("query", ""), args.get("max_results", 5))
+    elif func_name == "notion_get_page_content":
+        return await notion_get_page_content(user_id, args.get("page_id", ""))
+    elif func_name == "notion_create_page":
+        return await notion_create_page(user_id, args.get("parent_id"), args.get("title", "Untitled"), args.get("content", ""))
+    elif func_name == "notion_append_block_children":
+        return await notion_append_block_children(user_id, args.get("page_id", ""), args.get("content", ""))
     return "Unknown tool function."
 
 def count_tokens(messages):
@@ -932,7 +1150,7 @@ async def generate(request: Request):
 
                                         if content and not has_tool_call:
                                             yield json.dumps({"answer": content}) + "\n"
-                                me:
+                                except json.JSONDecodeError:
                                     continue
 
                     if has_tool_call and tool_calls_acc: