Spaces:
Build error
Build error
| """ | |
| 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) | |
| 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, | |
| } | |
| 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, | |
| } | |