Spaces:
Running on Zero
Running on Zero
| from mcp.server.fastmcp import FastMCP | |
| from typing import Dict, List, Any, Optional | |
| import logging | |
| import json | |
| from pathlib import Path | |
| import os | |
| logger = logging.getLogger(__name__) | |
| # Create FastMCP server instance | |
| mcp = FastMCP("PDF Parser") | |
| # Global document storage | |
| documents: Dict[str, Dict] = {} | |
| # Storage file path - use absolute path to ensure it works from any directory | |
| STORAGE_FILE = Path(__file__).parent / "storage.json" | |
| def load_documents(): | |
| """Load documents from storage.json file""" | |
| global documents | |
| try: | |
| if STORAGE_FILE.exists(): | |
| with open(STORAGE_FILE, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| documents.clear() | |
| documents.update(data) | |
| logger.info(f"Loaded {len(documents)} documents from storage at {STORAGE_FILE}") | |
| else: | |
| logger.info(f"No storage file found at {STORAGE_FILE}, starting with empty documents") | |
| documents.clear() | |
| except Exception as e: | |
| logger.error(f"Error loading documents from storage: {str(e)}") | |
| documents.clear() | |
| def reload_documents(): | |
| """Reload documents from storage - called by tools to ensure fresh data""" | |
| load_documents() | |
| # Load documents on startup | |
| load_documents() | |
| def list_documents() -> str: | |
| """List all processed PDF documents""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| if not documents: | |
| return "No documents available." | |
| doc_list = [] | |
| for file_id, doc in documents.items(): | |
| doc_list.append(f"- {file_id}: {doc['filename']} (Status: {doc['status']})") | |
| return "Available documents:\n" + "\n".join(doc_list) | |
| def get_document_summary(file_id: str) -> str: | |
| """Get AI-generated summary of a specific document""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| if file_id not in documents: | |
| return f"Document {file_id} not found." | |
| doc = documents[file_id] | |
| if "summary" not in doc: | |
| return f"No summary available for document {file_id}." | |
| summary = doc["summary"] | |
| if summary["type"] == "single_chunk": | |
| return f"Document: {doc['filename']}\n\nSummary:\n{summary['summary']}" | |
| else: | |
| # Multi-chunk summary | |
| result = f"Document: {doc['filename']}\n\nOverall Summary:\n{summary['overall_summary']}\n\n" | |
| result += "Chunk Summaries:\n" | |
| for i, chunk_summary in enumerate(summary['chunk_summaries']): | |
| result += f"\nChunk {i+1}:\n{chunk_summary}\n" | |
| return result | |
| def get_document_content(file_id: str, chunk_id: Optional[int] = None) -> str: | |
| """Get full content of a specific document or a specific chunk""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| if file_id not in documents: | |
| return f"Document {file_id} not found." | |
| doc = documents[file_id] | |
| if chunk_id is not None: | |
| # Get specific chunk | |
| if "chunks" in doc and chunk_id < len(doc["chunks"]): | |
| chunk = doc["chunks"][chunk_id] | |
| return f"Document: {doc['filename']}\nChunk {chunk_id} (Pages {chunk['page_range']}):\n\n{chunk['text']}" | |
| else: | |
| return f"Chunk {chunk_id} not found in document {file_id}." | |
| else: | |
| # Get full content | |
| if "extracted_text" in doc: | |
| return f"Document: {doc['filename']}\n\nFull Content:\n{doc['extracted_text']}" | |
| else: | |
| return f"Content not available for document {file_id}." | |
| def search_documents(query: str, file_id: Optional[str] = None) -> str: | |
| """Search for specific content across all documents or within a specific document""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| results = [] | |
| documents_to_search = {file_id: documents[file_id]} if file_id else documents | |
| for doc_id, doc in documents_to_search.items(): | |
| if "extracted_text" not in doc: | |
| continue | |
| text = doc["extracted_text"].lower() | |
| query_lower = query.lower() | |
| if query_lower in text: | |
| # Find context around matches | |
| lines = doc["extracted_text"].split('\n') | |
| matching_lines = [] | |
| for i, line in enumerate(lines): | |
| if query_lower in line.lower(): | |
| # Include context (previous and next lines) | |
| start = max(0, i-2) | |
| end = min(len(lines), i+3) | |
| context = lines[start:end] | |
| matching_lines.append(f"Line {i+1}: " + "\n".join(context)) | |
| if matching_lines: | |
| results.append(f"Document: {doc['filename']} (ID: {doc_id})") | |
| results.extend(matching_lines[:5]) # Limit results | |
| results.append("") | |
| if not results: | |
| return f"No results found for query: {query}" | |
| return f"Search results for '{query}':\n\n" + "\n".join(results) | |
| def get_document_metadata(file_id: str) -> str: | |
| """Get metadata about a document (pages, size, processing time, etc.)""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| if file_id not in documents: | |
| return f"Document {file_id} not found." | |
| doc = documents[file_id] | |
| metadata = { | |
| "File ID": file_id, | |
| "Filename": doc.get("filename", "N/A"), | |
| "Status": doc.get("status", "N/A"), | |
| "Processed At": doc.get("processed_at", "N/A"), | |
| "Chunks": len(doc.get("chunks", [])), | |
| "Total Tokens": sum(chunk.get("tokens", 0) for chunk in doc.get("chunks", [])), | |
| "Has Summary": "summary" in doc | |
| } | |
| if "summary" in doc: | |
| metadata["Summary Type"] = doc["summary"].get("type", "N/A") | |
| text = f"Metadata for {doc.get('filename', file_id)}:\n\n" | |
| text += "\n".join([f"{key}: {value}" for key, value in metadata.items()]) | |
| return text | |
| def answer_question(question: str, file_id: str) -> str: | |
| """Answer a question about a specific document using its content and summary""" | |
| # Always reload documents to get latest data | |
| reload_documents() | |
| if file_id not in documents: | |
| return f"Document {file_id} not found." | |
| doc = documents[file_id] | |
| # Get document summary for context | |
| summary_text = "No summary available" | |
| if "summary" in doc: | |
| summary = doc["summary"] | |
| if summary["type"] == "single_chunk": | |
| summary_text = summary["summary"] | |
| else: | |
| summary_text = summary["overall_summary"] | |
| # Get document content (truncated for context) | |
| content = doc.get("extracted_text", "") | |
| content_preview = content[:5000] + "..." if len(content) > 5000 else content | |
| # Return context for the question | |
| return f"""Question: {question} | |
| Document: {doc['filename']} | |
| Summary: {summary_text} | |
| Content Preview: {content_preview} | |
| To get a detailed AI-powered answer, please use the Anthropic API integration in the main application.""" | |
| # Helper functions for document management | |
| def update_document_data(file_id: str, document_data: Dict[str, Any]): | |
| """Update document data in MCP server""" | |
| documents[file_id] = document_data | |
| logger.info(f"Updated document data for {file_id}") | |
| def remove_document(file_id: str): | |
| """Remove document from MCP server""" | |
| if file_id in documents: | |
| del documents[file_id] | |
| logger.info(f"Removed document {file_id} from MCP server") | |
| def get_server(): | |
| """Get the FastMCP server instance""" | |
| return mcp |