Spaces:
Paused
Paused
File size: 4,657 Bytes
70a33a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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)
|