modules/entity_search.py

"""
Entity Search & Filter Module — NEW feature
Search, filter and rank extracted entities.
"""


def search_entities(entities: list[dict], query: str) -> list[dict]:
    """Case-insensitive keyword search within extracted entities."""
    if not query.strip():
        return entities
    q = query.lower()
    return [e for e in entities if q in e["word"].lower()]


def filter_by_label(entities: list[dict], labels: list[str]) -> list[dict]:
    """Filter entities to only include specified label types."""
    if not labels:
        return entities
    return [e for e in entities if e["label"] in labels]


def filter_by_confidence(entities: list[dict], min_score: float = 80.0) -> list[dict]:
    """Keep only entities above a confidence threshold."""
    return [e for e in entities if e["score"] >= min_score]


def get_context(text: str, entity: dict, window: int = 80) -> str:
    """
    Return surrounding context for an entity.
    Shows 'window' characters before and after the entity span.
    """
    start = max(0, entity["start"] - window)
    end   = min(len(text), entity["end"] + window)
    snippet = text[start:end]
    return f"...{snippet}..." if start > 0 or end < len(text) else snippet