mcp_server.py

import os
import uvicorn
from mcp.server.fastmcp import FastMCP
from sqlalchemy import create_engine, text

# 1. Database Connection (Optimized for 2 CPU / 16GB RAM)
DB_URI = os.getenv("DB_URI", "postgresql://postgres:password@localhost:5432/postgres")
# Pool size 5 prevents too many context switches on 2 vCPUs
engine = create_engine(DB_URI, pool_size=5, max_overflow=10)

# 2. Initialize FastMCP
mcp = FastMCP("Shamela4")

def fetch_data(query: str, params: dict = {}):
    try:
        with engine.connect() as conn:
            result = conn.execute(text(query), params)
            return [dict(row._mapping) for row in result]
    except Exception as e:
        return {"error": str(e)}

def fetch_one(query: str, params: dict = {}):
    try:
        with engine.connect() as conn:
            result = conn.execute(text(query), params).fetchone()
            if result:
                return dict(result._mapping)
            return None
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
def get_database_stats() -> dict:
    """Get database statistics: counts of books, authors, categories, and pages."""
    try:
        with engine.connect() as conn:
            books_count = conn.execute(text("SELECT COUNT(book_id) FROM books")).scalar() or 0
            authors_count = conn.execute(text("SELECT COUNT(author_id) FROM authors")).scalar() or 0
            categories_count = conn.execute(text("SELECT COUNT(category_id) FROM categories")).scalar() or 0
            pages_count = conn.execute(text("SELECT COUNT(page_id) FROM pages")).scalar() or 0
            return {
                "total_categories": categories_count,
                "total_authors": authors_count,
                "total_books": books_count,
                "total_pages": pages_count,
            }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
def search_categories(limit: int = 50) -> list:
    """Get list of categories in the library."""
    limit = min(limit, 100) # Hard limit
    return fetch_data("SELECT category_id, category_name_ar, name_en FROM categories ORDER BY sort_order LIMIT :limit", {"limit": limit})

@mcp.tool()
def search_authors(search_term: str, limit: int = 10) -> list:
    """Search for authors by their Arabic name."""
    limit = min(limit, 20)
    query = "SELECT author_id, name_ar, death_hijri FROM authors WHERE name_ar LIKE :search ORDER BY death_hijri LIMIT :limit"
    return fetch_data(query, {"search": f"%{search_term}%", "limit": limit})

@mcp.tool()
def search_books(title_or_author: str, limit: int = 10) -> list:
    """Search for books by title or author name."""
    limit = min(limit, 20)
    query = """
        SELECT book_id, title_ar, main_author_name_ar, main_author_death_hijri 
        FROM books 
        WHERE title_ar LIKE :search OR main_author_name_ar LIKE :search 
        ORDER BY book_id LIMIT :limit
    """
    return fetch_data(query, {"search": f"%{title_or_author}%", "limit": limit})

@mcp.tool()
def get_book_toc(book_id: int) -> list:
    """Get the Table of Contents (Daftar Isi) for a specific book_id."""
    query = "SELECT title_id, title_text, page_id, parent_id FROM toc WHERE book_id = :book_id ORDER BY title_id"
    return fetch_data(query, {"book_id": book_id})

@mcp.tool()
def search_text_in_books(query: str, limit: int = 5) -> list:
    """
    Full-text search inside book pages. 
    Use this to find specific Arabic text/words inside the books.
    Returns highly relevant text snippets (highlights).
    """
    limit = min(limit, 10) # Strict limit for efficiency on 16GB RAM
    
    # Use websearch_to_tsquery for efficient and safe natural language parsing
    sql = """
        SELECT p.book_id, p.page_num, p.sequence_num, b.title_ar as book_title,
               ts_headline('arabic', p.body, websearch_to_tsquery('arabic', :query)) as snippet
        FROM pages p
        JOIN books b ON p.book_id = b.book_id
        WHERE to_tsvector('arabic', p.body) @@ websearch_to_tsquery('arabic', :query)
        ORDER BY p.book_id, p.sequence_num
        LIMIT :limit
    """
    return fetch_data(sql, {"query": query, "limit": limit})

@mcp.tool()
def get_page_content(book_id: int, page_num: int) -> dict:
    """Retrieve the exact, full text of a specific page number in a book."""
    query = "SELECT page_num, body, footnotes FROM pages WHERE book_id = :book_id AND page_num = :page_num LIMIT 1"
    return fetch_one(query, {"book_id": book_id, "page_num": page_num})

if __name__ == "__main__":
    print("Starting Shamela4 FastMCP Server with SSE on port 8001...")
    mcp.run(transport="sse", host="0.0.0.0", port=8001)