Spaces:
Sleeping
Sleeping
| 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() | |