agents/postprocess.py

"""
Shared post-processing for all agents.

After model.generate() produces text, this module:
1. Parses text into structured fields (AnswerParser)
2. Annotates the source frame with visual overlays (FrameAnnotator)
3. Returns an enriched Result with multimodal fields
"""

from agents.base import Result
from agents.answer_parser import AnswerParser
from agents.annotator import FrameAnnotator


def enrich_result(result: Result, image_path: str = None) -> Result:
    """Post-process a Result: parse text → structured data, annotate frame.

    Args:
        result: The base Result with answer text and expert_used
        image_path: Path to the source frame (for annotation)

    Returns:
        Same Result with multimodal fields populated
    """
    if not result.ok or not result.answer:
        return result

    # Clean up model output: strip list-string format artifacts
    answer = result.answer
    if answer.startswith("[") and "'" in answer:
        # Model generated Python list string like "['bear', 'bear', ...]"
        import re
        items = re.findall(r"'([^']+)'", answer)
        if items:
            # Deduplicate preserving order
            seen = set()
            unique = [x for x in items if not (x in seen or seen.add(x))]
            answer = ", ".join(unique)
    # Strip leading/trailing whitespace and quotes
    answer = answer.strip().strip("'\"")
    result.answer = answer

    # Parse text into structured fields
    parsed = AnswerParser.parse(result.answer, result.expert_used)

    result.output_type = parsed.get("output_type", "text")
    result.detections = parsed.get("detections", [])
    result.counts = parsed.get("counts", {})
    result.text_regions = parsed.get("text_regions", [])
    result.alert = parsed.get("alert", {})
    result.analysis = parsed.get("analysis", {})
    result.tracks = parsed.get("tracks", [])
    result.scene_attributes = parsed.get("scene_attributes", {})

    # Annotate frame if visual output is needed
    if result.output_type != "text" and image_path:
        try:
            import cv2
            frame = cv2.imread(image_path)
            if frame is not None:
                result.annotated_frame_base64 = FrameAnnotator.annotate_and_encode(
                    frame,
                    detections=result.detections,
                    agent_type=result.expert_used,
                    counts=result.counts if result.counts else None,
                    text_regions=result.text_regions if result.text_regions else None,
                    tracks=result.tracks if result.tracks else None,
                    alert=result.alert if result.alert else None,
                )
        except Exception:
            pass  # Annotation is best-effort — don't crash on failure

    return result