server.py

import os
import httpx
from typing import Dict, Any
from mcp.server.fastmcp import FastMCP
from modules import vertex_ai, vector_store
from dotenv import load_dotenv

# Initialize FastMCP server
mcp = FastMCP("Ebook Knowledge Base")

# Load configuration
load_dotenv()

@mcp.tool()
async def ask_ebook_knowledge_base(query: str) -> str:
    """Asks a question grounded in the private ebook knowledge base."""
    try:
        embedding = vertex_ai.generate_embedding(query)
        context_chunks = vector_store.search_relevant_chunks(embedding)
        if not context_chunks: return "No relevant information found."
        return vertex_ai.synthesize_answer(query, context_chunks)
    except Exception as e: return f"Error: {str(e)}"

@mcp.tool()
async def generate_ebook_outline_from_kb(topic: str) -> str:
    """Drafts a new 10-chapter book outline inspired by your existing library."""
    try:
        embedding = vertex_ai.generate_embedding(topic)
        context_chunks = vector_store.search_relevant_chunks(embedding, limit=10)
        return vertex_ai.generate_outline_from_context(topic, context_chunks)
    except Exception as e: return f"Error: {str(e)}"

@mcp.tool()
async def ingest_document(bucket_name: str, blob_name: str, ebook_id: str) -> str:
    """Triggers the ingestion pipeline to index a new document."""
    service_url = os.getenv("INGESTION_SERVICE_URL")
    if not service_url: return "Error: INGESTION_SERVICE_URL not configured."
    async with httpx.AsyncClient() as client:
        try:
            res = await client.post(f"{service_url.rstrip('/')}/ingest", json={"bucket_name": bucket_name, "blob_name": blob_name, "ebook_id": ebook_id}, timeout=10.0)
            return f"Ingestion triggered: {res.text}" if res.status_code == 200 else f"Failed: {res.text}"
        except Exception as e: return f"Error: {str(e)}"

@mcp.tool()
def get_knowledge_base_status() -> Dict[str, Any]:
    """Returns statistics of the knowledge base."""
    return vector_store.get_stats()

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