""" GROBID extractor plugin. Registers the GrobidTrainingExtractor with the extraction registry. Provides a download endpoint for reviewers to download complete training packages. """ import logging from pathlib import Path from typing import Any, Callable from fastapi_app.lib.plugins.plugin_base import Plugin, PluginContext from fastapi_app.lib.plugins.plugin_tools import get_plugin_config from fastapi_app.plugins.tei_wizard.plugin import TeiWizardPlugin from fastapi_app.lib.extraction import ExtractorRegistry from fastapi_app.lib.sse.event_bus import get_event_bus from .extractor import GrobidTrainingExtractor logger = logging.getLogger(__name__) class GrobidPlugin(Plugin): """Plugin that provides GROBID-based extraction.""" def __init__(self) -> None: get_plugin_config("plugin.grobid.server.url", "GROBID_SERVER_URL", default="") get_plugin_config("plugin.grobid.server.timeout", "GROBID_SERVER_TIMEOUT", default=10) get_plugin_config("plugin.grobid.extraction.timeout", "GROBID_EXTRACTION_TIMEOUT", default=300) get_plugin_config("plugin.grobid.cache.disabled", "GROBID_DISABLE_CACHE", default=False) @property def metadata(self) -> dict[str, Any]: """Return plugin metadata.""" return { "id": "grobid", "name": "GROBID Extractor", "description": "Extract training data using GROBID server", "category": "extractor", "version": "1.0.0", "required_roles": ["user"], "endpoints": [ { "name": "download_training", "label": "Download GROBID Training Data", "description": "Download complete GROBID training package for a collection", "category": "collection", "state_params": ["collection"], "required_roles": ["reviewer"], }, ], "dependencies": ["tei-wizard"], } def get_endpoints(self) -> dict[str, Callable]: """Return available endpoints.""" return { "download_training": self.download_training, } @classmethod def is_available(cls) -> bool: """Check if GROBID server URL is configured.""" from fastapi_app.lib.utils.config_utils import get_config return bool(get_config().get("plugin.grobid.server.url")) async def initialize(self, context: PluginContext) -> None: """Register the GROBID extractor and event handlers.""" from fastapi_app.lib.plugins.frontend_extension_registry import FrontendExtensionRegistry ext_registry = FrontendExtensionRegistry.get_instance() extension_file = Path(__file__).parent / "extensions" / "grobid-sync.js" logger.debug("DEBUG grobid-sync extension file path: %s (exists=%s)", extension_file, extension_file.exists()) if extension_file.exists(): ext_registry.register_extension(extension_file, self.metadata["id"]) logger.info("Registered grobid-sync frontend extension") else: logger.warning("grobid-sync extension file not found: %s", extension_file) registry = ExtractorRegistry.get_instance() registry.register(GrobidTrainingExtractor) # Register event handler for file deletion cache cleanup event_bus = get_event_bus() event_bus.on("file.deleted", self._on_file_deleted) """Register the TEI header enrichment enhancement with tei-wizard.""" # todo: create auto-discover reusable utility func tei_wizard = context.get_dependency("tei-wizard") if isinstance(tei_wizard, TeiWizardPlugin): for enhancement_filename in ["split-bibl.js", "segment-footnotes.js", "desegment-footnotes.js"]: enhancement_file = Path(__file__).parent / "enhancements" / enhancement_filename if enhancement_file.exists(): tei_wizard.register_enhancement(enhancement_file, self.metadata["id"]) else: logger.warning(f"Enhancement file not found: {enhancement_file}") else: logger.debug("tei-wizard dependency not available") logger.info("GROBID extractor plugin initialized") async def cleanup(self) -> None: """Unregister the GROBID extractor and event handlers.""" registry = ExtractorRegistry.get_instance() registry.unregister("grobid") # Unregister event handler event_bus = get_event_bus() event_bus.off("file.deleted", self._on_file_deleted) logger.info("GROBID extractor plugin cleaned up") async def _on_file_deleted(self, stable_id: str, **kwargs) -> None: """ Clean up cached GROBID training data when a file is deleted. This handler is called when any file is deleted. It checks if the file was a PDF and removes any cached training data for that document. Args: stable_id: The stable_id of the deleted file """ from fastapi_app.lib.core.dependencies import get_db from fastapi_app.lib.repository.file_repository import FileRepository try: # Get file info to check if it was a PDF db = get_db() file_repo = FileRepository(db) file_info = file_repo.get_file_by_stable_id(stable_id) if not file_info or file_info.file_type != "pdf": return doc_id = file_info.doc_id if not doc_id: return # Check if any other TEI files still exist for this doc_id doc_files = file_repo.get_files_by_doc_id(doc_id) other_teis = [f for f in doc_files if f.file_type == "tei" and f.stable_id != stable_id and not f.deleted] if other_teis: # Other PDFs exist, don't delete cache return # Delete cached training data for this document from fastapi_app.plugins.grobid.cache import delete_cache_for_doc if delete_cache_for_doc(doc_id): logger.info(f"Deleted cached GROBID training data for {doc_id}") except Exception as e: logger.warning(f"Failed to clean up GROBID cache for {stable_id}: {e}") async def download_training( self, context: PluginContext, params: dict[str, Any] ) -> dict[str, Any]: """ Generate download URL for GROBID training package. Args: context: Plugin context params: Parameters including 'collection' (collection ID) Returns: downloadUrl pointing to the download route """ collection = params.get("collection") if not collection: return { "error": "No collection selected", "message": "Please select a collection first.", } # Build download URL with optional parameters # no_progress=false enables SSE progress events for UI usage download_url = f"/api/plugins/grobid/download?collection={collection}&no_progress=false" # Add optional parameters if provided if params.get("force_refresh"): download_url += "&force_refresh=true" if params.get("flavor"): download_url += f"&flavor={params['flavor']}" return { "downloadUrl": download_url, "collection": collection, }