Spaces:
Running
Running
File size: 3,331 Bytes
0e38162 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/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())
|