notion_mcp_server.py

#!/usr/bin/env python3
import os
import json
import requests
from dotenv import load_dotenv
from fastmcp import FastMCP, Context
from typing import Dict, Any
from pydantic import BaseModel, Field

# Initialize FastMCP server
mcp = FastMCP(
    name="Notion Knowledge Base",
    description="MCP server for querying a Notion knowledge base",
    version="0.1.0",
    dependencies=["python-dotenv", "requests"]
)

# Load environment variables
load_dotenv()

# Get API key from environment
API_KEY = os.getenv('DIFY_API_BACKEND_KEY')
if not API_KEY:
    raise ValueError("DIFY_API_BACKEND_KEY environment variable not set")

class NotionResponse(BaseModel):
    """Structure for Notion API responses."""
    answer: str = Field(default="")
    notion_page_url: str = Field(default="")
    notion_page_id: str = Field(default="")

    @classmethod
    def from_api_response(cls, data: Dict[str, Any]) -> "NotionResponse":
        """Create a NotionResponse from API response data."""
        outputs = data.get('data', {}).get('outputs', {})
        return cls(
            answer=outputs.get('answer', ''),
            notion_page_url=outputs.get('notion_page_url', ''),
            notion_page_id=outputs.get('notion_page_id', '')
        )

@mcp.tool()
def ask_notion_question(question: str, ctx: Context) -> Dict[str, Any]:
    """Ask a question about the Notion knowledge base.
    
    Args:
        question: The question to ask about the Notion knowledge base
        ctx: MCP context for logging and progress tracking
        
    Returns:
        Dictionary containing the answer and related Notion page information
        
    Raises:
        ValueError: If the API key is not set or if the request fails
    """
    ctx.info(f"Processing question: {question}")
    
    try:
        url = "https://dify.rickydata.com/v1/workflows/run"
        headers = {
            'Authorization': f"Bearer {API_KEY}",
            'Content-Type': 'application/json'
        }
        payload = {
            'inputs': {'question': question},
            'response_mode': 'blocking',
            'user': 'curation_agent_python'
        }
        
        ctx.debug("Sending request to Notion API")
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        
        data = response.json()
        if not data.get('data', {}).get('outputs', {}):
            raise ValueError("Invalid response format from API")
        
        result = NotionResponse.from_api_response(data)
        ctx.debug("Successfully received response from Notion API")
        
        return result.model_dump()
        
    except requests.RequestException as e:
        error_msg = f"API request failed: {str(e)}"
        ctx.error(error_msg)
        raise ValueError(error_msg)
    except json.JSONDecodeError as e:
        error_msg = f"Invalid JSON response: {str(e)}"
        ctx.error(error_msg)
        raise ValueError(error_msg)
    except Exception as e:
        error_msg = f"Unexpected error: {str(e)}"
        ctx.error(error_msg)
        raise ValueError(error_msg)

if __name__ == "__main__":
    mcp.run()