File size: 4,531 Bytes
be5a5d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6020272
 
 
 
 
 
 
be5a5d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6020272
 
 
 
 
 
 
be5a5d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6020272
 
 
 
 
 
 
be5a5d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6020272
 
 
 
 
 
be5a5d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6020272
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
from __future__ import annotations

from typing import Any

from .trajectory_store import trajectory_store


class ExampleExtractor:
    """Extract few-shot examples from successful trajectories."""

    def extract_classification_examples(
        self,
        task_id: str,
        persona: str,
        limit: int = 2,
    ) -> list[dict[str, Any]]:
        """Extract classification action examples from high-scoring runs."""
        trajectories = trajectory_store.get_similar_trajectories(task_id, persona, limit)
        examples = []

        for traj in trajectories:
            for step in traj.get("trajectory", []):
                action = step.get("action", {})
                if action.get("action_type") == "classify":
                    examples.append(
                        {
                            "email_id": action.get("email_id"),
                            "label": action.get("label"),
                            "reason": step.get("reason", ""),
                        }
                    )
                    if len(examples) >= limit:
                        return examples

        return examples

    def extract_reply_examples(
        self,
        task_id: str,
        persona: str,
        limit: int = 2,
    ) -> list[dict[str, Any]]:
        """Extract reply action examples from high-scoring runs."""
        trajectories = trajectory_store.get_similar_trajectories(task_id, persona, limit)
        examples = []

        for traj in trajectories:
            for step in traj.get("trajectory", []):
                action = step.get("action", {})
                if action.get("action_type") == "reply":
                    examples.append(
                        {
                            "email_id": action.get("email_id"),
                            "content": action.get("content", "")[:200],
                            "sender_role": step.get("sender_role", "unknown"),
                        }
                    )
                    if len(examples) >= limit:
                        return examples

        return examples

    def extract_escalation_examples(
        self,
        task_id: str,
        persona: str,
        limit: int = 2,
    ) -> list[dict[str, Any]]:
        """Extract escalation action examples from high-scoring runs."""
        trajectories = trajectory_store.get_similar_trajectories(task_id, persona, limit)
        examples = []

        for traj in trajectories:
            for step in traj.get("trajectory", []):
                action = step.get("action", {})
                if action.get("action_type") == "escalate":
                    examples.append(
                        {
                            "email_id": action.get("email_id"),
                            "escalate_to": action.get("escalate_to"),
                            "reason": step.get("reason", ""),
                        }
                    )
                    if len(examples) >= limit:
                        return examples

        return examples

    def extract_prioritization_examples(
        self,
        task_id: str,
        persona: str,
        limit: int = 2,
    ) -> list[dict[str, Any]]:
        """Extract prioritization examples from high-scoring runs."""
        trajectories = trajectory_store.get_similar_trajectories(task_id, persona, limit)
        examples = []

        for traj in trajectories:
            for step in traj.get("trajectory", []):
                action = step.get("action", {})
                if action.get("action_type") == "prioritize":
                    examples.append(
                        {
                            "priority_order": action.get("priority_order", []),
                            "reason": step.get("reason", ""),
                        }
                    )
                    if len(examples) >= limit:
                        return examples

        return examples

    def extract_all_examples(
        self,
        task_id: str,
        persona: str,
    ) -> dict[str, list[dict[str, Any]]]:
        """Extract all types of examples for a task/persona combination."""
        return {
            "classify": self.extract_classification_examples(task_id, persona),
            "reply": self.extract_reply_examples(task_id, persona),
            "escalate": self.extract_escalation_examples(task_id, persona),
            "prioritize": self.extract_prioritization_examples(task_id, persona),
        }


example_extractor = ExampleExtractor()