src/agents/table_detection_agent.py

from PIL import Image
from typing import Dict, Any
from .base_agent import LLMAgent


class TableDetectionAgent(LLMAgent):
    name = "table_detection_agent"
    description = "Identifies and analyzes tables in PDF documents - detects structure, headers, rows, and columns"

    def build_prompt(self, pil_image: Image.Image) -> str:
        width, height = pil_image.size
        return f"""Analyze this PDF page ({width}x{height}) for tables.

Look for:
1. Are there any tables on this page?
2. What is the table structure (rows, columns, merged cells)?
3. Are headers clearly defined?
4. Is the table bordered or borderless?
5. What data does the table contain?

If no tables found, state clearly. If found, describe in detail."""

    def execute(self, pil_image: Image.Image, reasoning: str) -> Dict[str, Any]:
        detected = "table" in reasoning.lower() and (
            "row" in reasoning.lower() or "column" in reasoning.lower()
        )
        return {
            "table_count": 1 if detected else 0,
            "tables": [self._parse_table(reasoning)] if detected else [],
            "detected": detected,
            "reasoning": reasoning,
        }

    def _parse_table(self, reasoning: str) -> Dict[str, Any]:
        return {"headers": [], "rows": [], "style": "unknown"}