Spaces:
Running
Running
File size: 1,739 Bytes
0a40afa 54478a0 0a40afa 54478a0 0a40afa |
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 |
"""Extract tables from PDF using Azure Document Intelligence."""
from typing import Dict, Any
import logging
from .base_agent import BaseAgent
from services.azure_di_service import AzureDIService
class TableAgent(BaseAgent):
def __init__(self, settings):
self.service = AzureDIService(settings.AZURE_DI_ENDPOINT, settings.AZURE_DI_KEY)
self.logger = logging.getLogger(__name__)
def execute(self, ctx: Dict[str, Any]):
"""Extract tables from PDF."""
try:
pdf_file = ctx.get("pdf_file")
if not pdf_file:
self.logger.error("No PDF file found in context")
return {}
# Get the current position of the file pointer
current_pos = pdf_file.tell()
self.logger.info(f"Current file position: {current_pos}")
# Reset to beginning if not at start
if current_pos != 0:
self.logger.info("Resetting file pointer to beginning")
pdf_file.seek(0)
# Read the file
pdf_bytes = pdf_file.read()
self.logger.info(f"Read {len(pdf_bytes)} bytes from PDF")
# Extract content using Azure DI
result = self.service.extract_tables(pdf_bytes)
# Store both text and tables in context
ctx["text"] = result["text"]
#ctx["tables"] = result["tables"]
self.logger.info(f"Extracted {len(result['text'])} characters of text including tables")
return result
except Exception as e:
self.logger.error(f"Error in TableAgent: {str(e)}")
self.logger.exception("Full traceback:")
return {} |