Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Test script for Enhanced Parser Agent | |
| Tests multimodal extraction: text + figures + tables + embeddings | |
| """ | |
| import asyncio | |
| import sys | |
| from pathlib import Path | |
| # Add parent directory to path | |
| sys.path.insert(0, str(Path(__file__).parent.parent)) | |
| from agents.enhanced_parser_agent import EnhancedParserAgent | |
| from agents.base_agent import AgentContext | |
| from services.neo4j_service import Neo4jService | |
| from services.mongo_service import MongoService | |
| from services.llm_service import LLMService | |
| from services.vector_store_service import VectorStoreService | |
| from services.lancedb_service import LanceDBService | |
| from memory.graph_memory import GraphMemory | |
| from utils.helpers import generate_id | |
| from utils.logger import get_logger | |
| logger = get_logger("test_parser") | |
| async def test_enhanced_parser(pdf_path: str) -> dict: | |
| """Test the enhanced parser on a PDF file.""" | |
| if not Path(pdf_path).exists(): | |
| logger.error(f"PDF file not found: {pdf_path}") | |
| return {"error": f"File not found: {pdf_path}"} | |
| logger.info(f"Testing enhanced parser on: {pdf_path}") | |
| # Initialize services | |
| try: | |
| neo4j = Neo4jService() | |
| mongo = MongoService() | |
| llm = LLMService() | |
| embedder = VectorStoreService() | |
| vector_db = LanceDBService() | |
| logger.info("Services initialized") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize services: {e}") | |
| return {"error": str(e)} | |
| # Create context | |
| job_id = "test_" + generate_id("job")[:8] | |
| doc_id = generate_id("doc") | |
| graph_memory = GraphMemory(neo4j=neo4j, embedder=embedder) | |
| ctx = AgentContext( | |
| job_id=job_id, | |
| doc_id=doc_id, | |
| pdf_path=pdf_path, | |
| neo4j=neo4j, | |
| mongo=mongo, | |
| llm=llm, | |
| embedder=embedder, | |
| graph_memory=graph_memory, | |
| ) | |
| ctx.vector_db = vector_db | |
| # Run enhanced parser | |
| parser = EnhancedParserAgent() | |
| try: | |
| result = await parser.execute(ctx) | |
| logger.info(f"Parser result: {result}") | |
| return result | |
| except Exception as e: | |
| logger.error(f"Parser execution failed: {e}") | |
| return {"error": str(e)} | |
| async def main(): | |
| """Main test function.""" | |
| # Check for PDF argument | |
| if len(sys.argv) < 2: | |
| logger.info("Usage: python test_docling_parser.py <pdf_path>") | |
| logger.info("Example: python test_docling_parser.py ./samples/paper.pdf") | |
| return | |
| pdf_path = sys.argv[1] | |
| logger.info("=" * 60) | |
| logger.info("ENHANCED PARSER AGENT TEST") | |
| logger.info("=" * 60) | |
| result = await test_enhanced_parser(pdf_path) | |
| logger.info("=" * 60) | |
| logger.info("TEST RESULTS:") | |
| logger.info("=" * 60) | |
| if "error" in result: | |
| logger.error(f"TEST FAILED: {result['error']}") | |
| else: | |
| logger.info(f"✓ Parser: {result.get('parser', 'unknown')}") | |
| logger.info(f"✓ Title: {result.get('title', 'N/A')}") | |
| logger.info(f"✓ Sections: {result.get('section_count', 0)}") | |
| logger.info(f"✓ Figures: {result.get('figure_count', 0)}") | |
| logger.info(f"✓ Tables: {result.get('table_count', 0)}") | |
| logger.info(f"✓ Vector indexed: {result.get('vector_indexed', False)}") | |
| logger.info("\nTEST PASSED ✓") | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |