src/nodes/entity_extractor.py

import re
from src.state import AgentState
from langchain_core.messages import HumanMessage, AIMessage

# Olist order IDs are 32-char hex strings
ORDER_ID_RE = re.compile(r'\b([a-f0-9]{32})\b', re.IGNORECASE)
SHORT_ID_RE = re.compile(r'\b(ORD|ORDER)[-_]?([A-Z0-9]{4,})\b', re.IGNORECASE)

PRONOUN_TRIGGERS = ["that", "it", "the same", "this", "that order",
                    "that one", "the order", "it again", "same one",
                    "the seller", "that seller", "the product", "that product"]

def extract_entities(state: AgentState) -> dict:
    """Extract entities and resolve pronouns using session context and conversation history"""
    print("[entity_extractor] Starting entity extraction")
    text = state.get("user_input", "").lower()
    
    result = {}
    session_context_updates = {}
    error_log_updates = []

    # 1. Try explicit order ID in current message
    match = ORDER_ID_RE.search(state.get("user_input", ""))
    if match:
        order_id = match.group(1).lower()
        result["last_order_id"] = order_id
        session_context_updates["order_id"] = order_id
        print(f"[entity_extractor] Extracted order_id: {order_id}")

    short_match = SHORT_ID_RE.search(state.get("user_input", ""))
    if short_match and not match:
        order_id = short_match.group(0)
        result["last_order_id"] = order_id
        print(f"[entity_extractor] Extracted short order_id: {order_id}")

    # 2. Check if this is a clarification response (user providing ID after being asked)
    prior_clarification = state.get("session_context", {}).get("last_clarification_type")
    if prior_clarification == "order_id" and not match and not short_match:
        # Current message might BE the order ID
        potential_id = state.get("user_input", "").strip()
        if len(potential_id) == 32 and ORDER_ID_RE.match(potential_id):
            order_id = potential_id.lower()
            result["last_order_id"] = order_id
            session_context_updates["order_id"] = order_id
            session_context_updates["last_clarification_type"] = None  # Clear flag
            print(f"[entity_extractor] Resolved clarification response to order_id: {order_id}")

    # 3. Pronoun resolution — use conversation history and session context
    if not match and not short_match:
        if any(trigger in text for trigger in PRONOUN_TRIGGERS):
            # First try to get from current state (most recent)
            carried_id = state.get("last_order_id")
            
            # Fall back to session context
            if not carried_id:
                carried_id = state.get("session_context", {}).get("order_id")
            
            # Last resort: extract from recent conversation history
            if not carried_id:
                recent_messages = state.get("messages", [])[-4:]  # Last 2 turns
                for msg in recent_messages:
                    if isinstance(msg, (HumanMessage, AIMessage)):
                        msg_match = ORDER_ID_RE.search(msg.content)
                        if msg_match:
                            carried_id = msg_match.group(1).lower()
                            break
            
            if carried_id:
                result["last_order_id"] = carried_id
                session_context_updates["order_id"] = carried_id
                print(f"[entity_extractor] Resolved pronoun to order_id={carried_id}")
                error_log_updates.append(f"[entity_extractor] Resolved pronoun to order_id={carried_id}")

    # 4. Carry forward seller_id and product_id from prior context if not in current message
    if not state.get("last_seller_id"):
        seller_id = state.get("session_context", {}).get("seller_id")
        if seller_id:
            result["last_seller_id"] = seller_id
    
    if not state.get("last_product_id"):
        product_id = state.get("session_context", {}).get("product_id")
        if product_id:
            result["last_product_id"] = product_id

    # 5. Extract category mentions (for product recommendations)
    KNOWN_CATEGORIES = ["computers", "electronics", "furniture", "toys",
                        "sports", "books", "health", "fashion", "phones",
                        "computer"]  # Singular forms last
    for cat in KNOWN_CATEGORIES:
        if cat in text:
            result["last_category"] = cat
            break
    
    # Merge session_context updates
    if session_context_updates:
        result["session_context"] = {**state.get("session_context", {}), **session_context_updates}
    
    # Merge error_log updates
    if error_log_updates:
        result["error_log"] = state.get("error_log", []) + error_log_updates

    return result