Angad2005/ChatRAG

1# mcp_main.py2import io3import os4from html import escape5from fastapi import FastAPI6from fastapi import HTTPException7from fastapi.responses import StreamingResponse8from pydantic import BaseModel, Field9from docx import Document10from reportlab.lib.pagesizes import letter11from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer12from reportlab.lib.styles import getSampleStyleSheet13 14# --- Model and App Initialization ---15 16app = FastAPI(17    title="Document Summarization MCP",18    description="A microservice to summarize texts and generate a document.",19    version="1.0.0"20)21 22generator = None23MODEL_NAME = os.getenv("SUMMARIZATION_MODEL", "sshleifer/distilbart-cnn-12-6")24 25def get_generator():26    """Load the summarizer only when the first summary request is received."""27    global generator28    if generator is None:29        from transformers import pipeline30        generator = pipeline("summarization", model=MODEL_NAME)31    return generator32 33# --- Pydantic Models for Request Body ---34 35class SummarizationRequest(BaseModel):36    documents: dict[str, str] = Field(..., description="A dictionary where keys are filenames and values are the document's text content.")37    doc_type: str = Field("pdf", description="The desired output document type. Either 'pdf' or 'docx'.")38 39# --- Helper Functions for Document Generation ---40 41def create_pdf_from_summaries(summaries: dict[str, str]) -> io.BytesIO:42    """Generates a PDF document from a dictionary of summaries."""43    buffer = io.BytesIO()44    doc = SimpleDocTemplate(buffer, pagesize=letter)45    styles = getSampleStyleSheet()46    story = []47 48    for filename, summary in summaries.items():49        story.append(Paragraph(escape(f"Summary for: {filename}"), styles['h2']))50        story.append(Spacer(1, 12))51        story.append(Paragraph(escape(summary).replace("\n", "<br/>"), styles['BodyText']))52        story.append(Spacer(1, 24))53 54    doc.build(story)55    buffer.seek(0)56    return buffer57 58def create_docx_from_summaries(summaries: dict[str, str]) -> io.BytesIO:59    """Generates a DOCX document from a dictionary of summaries."""60    buffer = io.BytesIO()61    doc = Document()62 63    for filename, summary in summaries.items():64        doc.add_heading(f"Summary for: {filename}", level=2)65        doc.add_paragraph(summary)66        doc.add_paragraph()  # Add a little space67 68    doc.save(buffer)69    buffer.seek(0)70    return buffer71 72# --- Helper Function for Text Generation ---73 74def generate_summary(text: str) -> str:75    """Generate a summary of the given text using the language model."""76    max_input_length = 50077    truncated_text = text[:max_input_length]78    result = get_generator()(truncated_text, max_length=150, truncation=True)79    return result[0]["summary_text"].strip()80 81# --- API Endpoint ---82 83@app.post("/summarize-and-create-document/")84async def summarize_and_create(request: SummarizationRequest):85    """86    Receives document texts, summarizes them, and returns a single87    consolidated PDF or DOCX file.88    """89    summaries = {}90    for filename, content in request.documents.items():91        try:92            summary = generate_summary(content)93        except Exception as exc:94            raise HTTPException(status_code=503, detail=f"Summarization model unavailable: {exc}") from exc95        summaries[filename] = summary96 97    if request.doc_type.lower() == 'pdf':98        buffer = create_pdf_from_summaries(summaries)99        media_type = "application/pdf"100        filename = "summaries.pdf"101    elif request.doc_type.lower() == 'docx':102        buffer = create_docx_from_summaries(summaries)103        media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"104        filename = "summaries.docx"105    else:106        raise HTTPException(status_code=400, detail="Invalid doc_type. Must be 'pdf' or 'docx'.")107 108    headers = {109        'Content-Disposition': f'attachment; filename="{filename}"'110    }111 112    return StreamingResponse(buffer, media_type=media_type, headers=headers)113 114@app.get("/health")115def health_check():116    return {"status": "ok" if generator is not None else "ready", "model": MODEL_NAME}117 118# --- FIX: this FastAPI app previously had no way to actually start ---119# `app = FastAPI(...)` only defines the app object; nothing ran a server120# process for it. If this file is meant to run as a standalone service121# (e.g. a separate Space, or a sidecar process next to the Gradio app),122# it needs an entrypoint like this:123if __name__ == "__main__":124    import uvicorn125    port = int(os.getenv("MCP_PORT", "8001"))126    uvicorn.run(app, host="0.0.0.0", port=port)127