File size: 4,706 Bytes
cf796c5
 
d2c5868
cf796c5
 
 
 
 
 
d2c5868
 
cf796c5
f9e8eed
d2c5868
 
cf796c5
f9e8eed
 
 
 
cf796c5
 
 
 
f9e8eed
 
 
 
cf796c5
 
 
f9e8eed
 
 
cf796c5
d2c5868
 
 
 
 
 
f9e8eed
 
 
 
 
d2c5868
 
cf796c5
 
d2c5868
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf796c5
f9e8eed
 
d2c5868
f9e8eed
cf796c5
d2c5868
 
f9e8eed
 
 
d2c5868
 
f9e8eed
 
 
d2c5868
 
cf796c5
8d73fdc
 
cf796c5
 
f9e8eed
cf796c5
f9e8eed
 
 
 
 
 
 
 
cf796c5
f9e8eed
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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