dhruv-punia-bits/memory-compaction-openenv
0
1from __future__ import annotations2 3from dataclasses import dataclass4 5from server.models import Difficulty, FutureQuery, MemoryItem, MemoryStatus, MemoryType, TaskDefinition, Turn6 7 8@dataclass(frozen=True)9class ScenarioConfig:10 seed: int11 token_budget: int12 recent_turn_window: int13 title: str14 summary: str15 profile_name: str16 travel_city: str17 travel_reason: str18 dietary_preference: str19 alternate_dietary_preference: str20 project_name: str21 project_context: str22 project_deadline: str23 corrected_deadline: str24 teammate_name: str25 favorite_tool: str26 mentor_name: str27 distraction_detail: str28 29 30def scenario(31 seed: int,32 difficulty: Difficulty,33 title: str,34 summary: str,35 profile_name: str,36 travel_city: str,37 travel_reason: str,38 dietary_preference: str,39 alternate_dietary_preference: str,40 project_name: str,41 project_context: str,42 project_deadline: str,43 corrected_deadline: str,44 teammate_name: str,45 favorite_tool: str,46 mentor_name: str,47 distraction_detail: str,48) -> ScenarioConfig:49 budgets = {50 Difficulty.EASY: (165, 5),51 Difficulty.MEDIUM: (155, 5),52 Difficulty.HARD: (145, 4),53 }54 token_budget, recent_turn_window = budgets[difficulty]55 return ScenarioConfig(56 seed=seed,57 token_budget=token_budget,58 recent_turn_window=recent_turn_window,59 title=title,60 summary=summary,61 profile_name=profile_name,62 travel_city=travel_city,63 travel_reason=travel_reason,64 dietary_preference=dietary_preference,65 alternate_dietary_preference=alternate_dietary_preference,66 project_name=project_name,67 project_context=project_context,68 project_deadline=project_deadline,69 corrected_deadline=corrected_deadline,70 teammate_name=teammate_name,71 favorite_tool=favorite_tool,72 mentor_name=mentor_name,73 distraction_detail=distraction_detail,74 )75 76 77SCENARIO_POOLS: dict[Difficulty, list[ScenarioConfig]] = {78 Difficulty.EASY: [79 scenario(1101, Difficulty.EASY, "Design Retreat Planner", "Travel planning plus product demo prep.", "Mira", "Kyoto", "a design retreat with her product team", "vegetarian lunches", "vegan lunches", "memory compaction demo", "a small showcase for the hackathon judging round", "Friday afternoon", "Friday afternoon", "Jules", "Notion", "Priya", "she finally finished Deep Work on the flight"),80 scenario(1102, Difficulty.EASY, "Founder Offsite", "Travel logistics and customer interview prep.", "Owen", "Berlin", "a founder offsite and customer interviews", "aisle seats on long flights", "window seats on long flights", "customer story workshop", "a prep sprint for next week's investor update", "Monday morning", "Monday morning", "Lena", "Linear", "Daniel", "he watched a documentary about Formula 1 over the weekend"),81 scenario(1103, Difficulty.EASY, "University Lab Visit", "Academic travel with coordination details.", "Anika", "Boston", "a university lab visit and recruiting day", "quiet hotel rooms", "hotels near transit", "campus partnership brief", "a collaboration deck for a faculty meeting", "May 6", "May 6", "Dev", "Asana", "Sara", "she spent Sunday reorganizing her bookshelf"),82 scenario(1104, Difficulty.EASY, "Sales Kickoff", "Conference travel and enablement planning.", "Noah", "Chicago", "a regional sales kickoff and partner dinner", "gluten-free breakfasts", "high-protein breakfasts", "partner enablement deck", "materials for a reseller kickoff", "June 10", "June 10", "Ivy", "Trello", "Priya", "he started learning espresso recipes at home"),83 scenario(1105, Difficulty.EASY, "Healthcare Workshop", "Client workshop with preference tracking.", "Riya", "Singapore", "a healthcare workshop and planning session", "dairy-free snacks", "fruit-based snacks", "clinical workflow pilot", "a narrow pilot for a hospital operations team", "July 9", "July 9", "Kabir", "ClickUp", "Daniel", "she has been testing a new yoga routine"),84 scenario(1106, Difficulty.EASY, "Product Filming Trip", "Creator trip with project coordination.", "Leo", "Los Angeles", "filming a product launch video and creator meetings", "early afternoon meetings", "late afternoon meetings", "launch cut review", "a review session for the public launch trailer", "August 14", "August 14", "Maya", "Airtable", "Sara", "he has been collecting vinyl records lately"),85 scenario(1107, Difficulty.EASY, "Civic Tech Demo", "Conference demo with operational details.", "Sana", "Delhi", "a civic tech summit and press briefing", "spicy vegetarian dinners", "mild vegetarian dinners", "demo booth setup", "a public sector technology booth for a summit", "September 2", "September 2", "Arjun", "Monday", "Priya", "she spent a whole evening tuning her bicycle gears"),86 scenario(1108, Difficulty.EASY, "Retail Ops Visit", "Store visits and field-notes workflow.", "Theo", "Toronto", "meeting retail operators and shadowing store managers", "decaf coffee after noon", "herbal tea after noon", "field visit recap", "a summary of operational issues from store visits", "October 1", "October 1", "Nina", "Coda", "Daniel", "he got absorbed in a long chess podcast"),87 scenario(1109, Difficulty.EASY, "Climate Grant Pitch", "Nonprofit trip with grant coordination.", "Asha", "Nairobi", "a climate grant pitch and stakeholder workshop", "plant-based catering", "vegetarian catering", "grant narrative revision", "the final revision round for a funding application", "November 11", "November 11", "Sam", "Basecamp", "Sara", "she baked sourdough for the first time"),88 scenario(1110, Difficulty.EASY, "Gaming Studio Sync", "Publisher trip with creative planning.", "Kai", "Tokyo", "a publisher sync and gameplay capture review", "low-sugar snacks", "protein bars", "capture review board", "a board that tracks bugs and footage notes", "December 3", "December 3", "Emi", "Jira", "Priya", "he spent all weekend fixing an old game console"),89 ],90 Difficulty.MEDIUM: [91 scenario(2201, Difficulty.MEDIUM, "Benchmark Launch", "Product launch planning with corrected deadlines.", "Sonia", "Lisbon", "a benchmark launch rehearsal and analyst briefing", "gluten-free snacks", "high-protein snacks", "benchmark launch", "the public release motion for a benchmarking product", "May 18", "May 21", "Arun", "Linear", "Daniel", "she listened to a three-hour history podcast"),92 scenario(2202, Difficulty.MEDIUM, "Fintech Rollout", "Enterprise rollout with changing requirements.", "Mason", "New York", "an enterprise rollout and customer steering review", "non-dairy cappuccinos", "black coffee", "merchant dashboard rollout", "a staged rollout for a merchant analytics dashboard", "June 4", "June 7", "Priyanka", "Jira", "Sara", "he spent the weekend comparing standing desks"),93 scenario(2203, Difficulty.MEDIUM, "Education Platform Pilot", "School pilot with corrected logistics.", "Leila", "Melbourne", "a school pilot kickoff and district workshop", "nut-free snacks", "fruit cups", "district pilot pack", "materials for a district-wide pilot of a learning tool", "July 12", "July 15", "Marcus", "Asana", "Priya", "she has been revisiting old watercolor sketches"),94 scenario(2204, Difficulty.MEDIUM, "Logistics Dashboard Review", "Ops dashboard with schedule changes.", "Carlos", "Mexico City", "a logistics dashboard review and warehouse visit", "quiet hotel rooms", "hotels near transit", "warehouse dashboard refresh", "a dashboard update for fulfillment managers", "August 8", "August 11", "Elena", "Monday", "Daniel", "he got sidetracked fixing a home speaker system"),95 scenario(2205, Difficulty.MEDIUM, "Streaming Analytics Demo", "Media analytics prep with preference change.", "Nora", "Seoul", "a streaming analytics demo and strategy offsite", "vegetarian dinners", "pescatarian dinners", "audience insights demo", "a demo for the executive strategy review", "September 17", "September 19", "Min", "Notion", "Sara", "she binge-read a novel over two evenings"),96 scenario(2206, Difficulty.MEDIUM, "Manufacturing QA Pilot", "Factory pilot with corrected task state.", "Harish", "Munich", "a manufacturing QA pilot and supplier meetings", "early check-ins", "late-morning check-ins", "quality signal pilot", "a computer vision pilot for the QA line", "October 6", "October 8", "Greta", "ClickUp", "Priya", "he spent Saturday pruning plants on his balcony"),97 scenario(2207, Difficulty.MEDIUM, "Travel Concierge App", "Consumer app planning with corrections.", "Zara", "Dubai", "a travel concierge app planning sprint and hotel partner meeting", "vegan breakfasts", "Mediterranean breakfasts", "concierge app beta", "a limited beta with hotel concierges", "November 5", "November 9", "Omar", "Linear", "Daniel", "she bought too many notebooks at an airport bookstore"),98 scenario(2208, Difficulty.MEDIUM, "Insurance Claims Flow", "Process redesign with updates.", "Ben", "Amsterdam", "an insurance claims flow workshop and process review", "low-noise hotel floors", "rooms away from elevators", "claims intake redesign", "a redesign of the claims intake workflow", "January 14", "January 16", "Sofia", "Airtable", "Sara", "he lost an afternoon building a custom keyboard"),99 scenario(2209, Difficulty.MEDIUM, "Food Delivery Ops", "Operations launch with revised preferences.", "Meera", "Bangkok", "a food delivery ops launch and restaurant partner roadshow", "spicy vegetarian lunches", "mild vegetarian lunches", "partner launch checklist", "a launch checklist for restaurant operations managers", "February 20", "February 24", "Tuan", "Coda", "Priya", "she became obsessed with documenting neighborhood cafes"),100 scenario(2210, Difficulty.MEDIUM, "HR Knowledge Assistant", "Internal tooling with corrected milestones.", "Ethan", "Sydney", "an internal tooling sprint and change-management workshop", "sparkling water instead of soda", "green tea instead of soda", "hr knowledge assistant", "an internal assistant for policy and onboarding questions", "March 3", "March 7", "Chloe", "Notion", "Daniel", "he spent Sunday restoring old family photos"),101 ],102 Difficulty.HARD: [103 scenario(3301, Difficulty.HARD, "Agent Memory Benchmark", "High-context benchmark planning with noise.", "Eli", "Seoul", "an agent benchmark working session and partner meeting", "quiet hotel rooms", "hotels near transit", "agent memory benchmark", "a benchmark suite for multi-step memory agents", "June 1", "June 4", "Nadia", "Obsidian", "Priya", "he went deep on a long essay about airport design"),104 scenario(3302, Difficulty.HARD, "Biotech Trial Ops", "Dense scientific coordination under budget pressure.", "Amina", "Basel", "a biotech trial operations review and sponsor meeting", "dairy-free breakfasts", "protein-heavy breakfasts", "trial ops dashboard", "an operations dashboard for a clinical trial program", "July 8", "July 11", "Jonas", "Jira", "Sara", "she spent the flight making a list of museums to visit later"),105 scenario(3303, Difficulty.HARD, "Cloud Cost Command Center", "Infrastructure planning with corrections and distractions.", "Victor", "Dublin", "a cloud cost command center review and CTO prep", "aisle seats on overnight flights", "bulkhead seats on overnight flights", "cost anomaly command center", "an operations center for cost anomaly triage", "August 13", "August 16", "Maeve", "Linear", "Daniel", "he read a long thread about espresso grinders"),106 scenario(3304, Difficulty.HARD, "Public Transit Redesign", "Civic operations planning with messy context.", "Farah", "Copenhagen", "a transit redesign workshop and policy roundtable", "quiet train-adjacent hotels", "hotels within walking distance of the metro", "ridership redesign brief", "a redesign brief for commuter information systems", "September 9", "September 12", "Jon", "Monday", "Priya", "she spent too much time comparing rain jackets online"),107 scenario(3305, Difficulty.HARD, "Fraud Review Console", "Trust-and-safety tooling with scoped updates.", "Gabriel", "Sao Paulo", "a fraud review console sprint and payments partner visit", "high-protein snacks", "fresh fruit snacks", "fraud review console", "a case-review tool for payment risk analysts", "October 21", "October 24", "Luisa", "Asana", "Sara", "he got pulled into an hour-long video on old camera lenses"),108 scenario(3306, Difficulty.HARD, "Carbon Reporting Pilot", "Policy-heavy project with long contextual turns.", "Isha", "Brussels", "a carbon reporting pilot and policy alignment meeting", "vegetarian dinners", "Mediterranean dinners", "carbon reporting pilot", "a reporting pilot for enterprise sustainability teams", "November 18", "November 22", "Pieter", "Airtable", "Daniel", "she filled a notes app with gift ideas during the flight"),109 scenario(3307, Difficulty.HARD, "Marketplace Integrity Launch", "Marketplace moderation planning with distractions.", "Rohan", "San Francisco", "a marketplace integrity launch and trust review", "rooms away from elevators", "rooms near the gym", "integrity launch review", "a launch review for marketplace abuse tooling", "December 5", "December 8", "Ava", "Notion", "Priya", "he spent a whole evening cleaning up duplicate photos"),110 scenario(3308, Difficulty.HARD, "Energy Grid Forecasting", "Operational forecasting with ambiguous future actions.", "Helena", "Oslo", "an energy grid forecasting review and regulator briefing", "early dinners", "late dinners", "grid forecasting refresh", "a forecasting refresh for an energy operations desk", "January 27", "January 30", "Mikkel", "ClickUp", "Sara", "she kept journaling unrelated ideas for a side project"),111 scenario(3309, Difficulty.HARD, "Customer Support Copilot", "Support tooling with many moving parts.", "Tariq", "Vancouver", "a support copilot rollout and escalation workshop", "decaf coffee after noon", "tea after noon", "support copilot rollout", "a rollout for an internal support assistant", "February 12", "February 15", "Elise", "Coda", "Daniel", "he got distracted reading about mountain rail systems"),112 scenario(3310, Difficulty.HARD, "Warehouse Robotics Demo", "Robotics demo with dense factual updates.", "Jun", "Taipei", "a warehouse robotics demo and investor walkthrough", "light vegetarian lunches", "rice-based lunches", "robotics demo review", "a demo review for a warehouse automation team", "March 19", "March 23", "Irene", "Obsidian", "Priya", "she spent a late night reorganizing recipe bookmarks"),113 ],114}115 116 117def _memory(118 memory_id: str,119 memory_type: MemoryType,120 subject: str,121 predicate: str,122 obj: str,123 source_turn_ids: list[int],124 status: MemoryStatus = MemoryStatus.ACTIVE,125 source_text: str = "",126 updated_from_memory_id: str | None = None,127 importance: float = 0.8,128 task_relevance: float = 0.8,129 requires_confirmation: bool = False,130) -> MemoryItem:131 return MemoryItem(132 memory_id=memory_id,133 type=memory_type,134 subject=subject,135 predicate=predicate,136 object=obj,137 source_turn_ids=source_turn_ids,138 source_text=source_text,139 status=status,140 updated_from_memory_id=updated_from_memory_id,141 importance=importance,142 task_relevance=task_relevance,143 requires_confirmation=requires_confirmation,144 )145 146 147def _scenario_for_seed(difficulty: Difficulty, seed: int | None) -> ScenarioConfig:148 pool = SCENARIO_POOLS[difficulty]149 if seed is None:150 return pool[0]151 for config in pool:152 if config.seed == seed:153 return config154 return pool[abs(seed) % len(pool)]155 156 157def build_task_definitions() -> list[TaskDefinition]:158 definitions: list[TaskDefinition] = []159 for difficulty in Difficulty:160 for config in SCENARIO_POOLS[difficulty]:161 payload = build_episode_payload(difficulty, config.seed)162 definitions.append(163 TaskDefinition(164 difficulty=difficulty,165 title=f"{difficulty.value.title()} - {config.title}",166 description=config.summary,167 seed=config.seed,168 token_budget=config.token_budget,169 turn_count=len(payload["turns"]),170 evaluation_query_count=len(payload["future_queries"]),171 )172 )173 return definitions174 175 176def build_episode_payload(difficulty: Difficulty, seed: int | None = None) -> dict[str, object]:177 config = _scenario_for_seed(difficulty, seed)178 name = config.profile_name179 city = config.travel_city180 trip_reason = config.travel_reason181 diet = config.dietary_preference182 alternate_diet = config.alternate_dietary_preference183 project = config.project_name184 project_context = config.project_context185 deadline = config.project_deadline186 corrected_deadline = config.corrected_deadline187 teammate = config.teammate_name188 tool = config.favorite_tool189 mentor = config.mentor_name190 distraction = config.distraction_detail191 192 turns = [193 Turn(194 turn_id=1,195 speaker="user",196 text=(197 f"Before we dive into planning, my name is {name}, and I am traveling to {city} next month for {trip_reason}. "198 f"I am juggling several threads at once, so I need you to remember the durable parts and ignore the fluff."199 ),200 ),201 Turn(202 turn_id=2,203 speaker="assistant",204 text=(205 f"Understood. I will keep track of the lasting details around your {city} trip and the work that depends on it, "206 "while keeping the running context compact."207 ),208 ),209 Turn(210 turn_id=3,211 speaker="user",212 text=(213 f"One standing preference that actually matters when you make suggestions: please remember that I prefer {diet}. "214 "That comes up often enough that I do not want to repeat it."215 ),216 ),217 Turn(218 turn_id=4,219 speaker="user",220 text=(221 f"Separately, I am building a {project} with {teammate}. The project is basically {project_context}, "222 "so if I ask for next steps later I need you to keep that frame in mind."223 ),224 ),225 Turn(226 turn_id=5,227 speaker="assistant",228 text=(229 f"Got it. The durable pieces so far are the {project} effort, your collaboration with {teammate}, "230 "and the constraints you want me to respect while planning around it."231 ),232 ),233 Turn(234 turn_id=6,235 speaker="user",236 text=(237 f"The first deadline is {deadline}, and I track work in {tool}. "238 "I know that sounds operational, but it matters because I will ask you to continue the plan later."239 ),240 ),241 ]242 243 gold_memories = [244 _memory("name", MemoryType.FACT, name, "is", "traveler", [1], source_text=turns[0].text, importance=0.7, task_relevance=0.45),245 _memory("trip", MemoryType.FACT, name, "travel_city", city, [1], source_text=turns[0].text, importance=0.95, task_relevance=0.95),246 _memory("diet", MemoryType.PREFERENCE, name, "prefers", diet, [3], source_text=turns[2].text, importance=0.9, task_relevance=0.95),247 _memory("project", MemoryType.PLAN, name, "building", project, [4], source_text=turns[3].text, importance=0.95, task_relevance=0.95),248 _memory("teammate", MemoryType.FACT, project, "teammate", teammate, [4], source_text=turns[3].text, importance=0.85, task_relevance=0.9),249 _memory("deadline", MemoryType.TASK, project, "deadline", deadline, [6], source_text=turns[5].text, importance=0.95, task_relevance=1.0),250 _memory("tool", MemoryType.PREFERENCE, name, "tracks_work_in", tool, [6], source_text=turns[5].text, importance=0.75, task_relevance=0.8),251 ]252 253 new_memory_ids_by_turn: dict[int, list[str]] = {254 1: ["name", "trip"],255 3: ["diet"],256 4: ["project", "teammate"],257 6: ["deadline", "tool"],258 }259 260 ambiguous_turn_ids: list[int] = []261 low_value_turn_ids: list[int] = []262 263 if difficulty in {Difficulty.MEDIUM, Difficulty.HARD}:264 turns.extend(265 [266 Turn(267 turn_id=7,268 speaker="user",269 text=(270 f"Quick correction so we do not anchor on the wrong thing: correction, the deadline moved to {corrected_deadline}. "271 "The original date is stale now, so please treat the newer one as the live commitment."272 ),273 ),274 Turn(275 turn_id=8,276 speaker="assistant",277 text=(278 f"Updated. I will treat {corrected_deadline} as the current deadline and avoid surfacing the older date unless you ask for history."279 ),280 ),281 Turn(282 turn_id=9,283 speaker="user",284 text=(285 f"Another detail that is useful but less central: my mentor for this work is {mentor}. "286 "That occasionally matters when I ask who should review a draft."287 ),288 ),289 Turn(290 turn_id=10,291 speaker="user",292 text=(293 "There is a chance I might present this work next week if a slot opens up, "294 "but that is not confirmed yet and I do not want you to lock it in as a hard plan."295 ),296 ),297 Turn(298 turn_id=11,299 speaker="user",300 text=(301 f"Also, a real preference change: actually I switched from {diet} to {alternate_diet}. "302 "The earlier preference is out of date, so future suggestions should use the newer one."303 ),304 ),305 Turn(306 turn_id=12,307 speaker="assistant",308 text="Understood. I will preserve the latest confirmed preference and keep the tentative presentation idea marked as uncertain.",309 ),310 ]311 )312 gold_memories.extend(313 [314 _memory(315 "deadline_v2",316 MemoryType.CORRECTION,317 project,318 "deadline",319 corrected_deadline,320 [7],321 source_text=turns[6].text,322 updated_from_memory_id="deadline",323 importance=0.95,324 task_relevance=1.0,325 ),326 _memory("mentor", MemoryType.FACT, project, "mentor", mentor, [9], source_text=turns[8].text, importance=0.7, task_relevance=0.75),327 _memory(328 "diet_v2",329 MemoryType.CORRECTION,330 name,331 "prefers",332 alternate_diet,333 [11],334 source_text=turns[10].text,335 updated_from_memory_id="diet",336 importance=0.95,337 task_relevance=1.0,338 ),339 ]340 )341 new_memory_ids_by_turn.update({7: ["deadline_v2"], 9: ["mentor"], 11: ["diet_v2"]})342 ambiguous_turn_ids = [10]343 344 if difficulty == Difficulty.HARD:345 turns.extend(346 [347 Turn(348 turn_id=13,349 speaker="user",350 text=(351 f"One thing that is probably noise for future planning: {distraction}. "352 "I am mentioning it because it came up in conversation, not because it should become durable memory."353 ),354 ),355 Turn(356 turn_id=14,357 speaker="user",358 text=(359 "More generally, if I include reflective or chatty details alongside operational ones, "360 "please keep the durable memory scoped to what would still matter in a later task."361 ),362 ),363 Turn(364 turn_id=15,365 speaker="assistant",366 text="Makes sense. I will keep only durable operational context, confirmed preferences, and the latest task state.",367 ),368 ]369 )370 low_value_turn_ids = [13, 14, 15]371 372 future_queries = [373 FutureQuery(query_id="q1", prompt="What travel city should the assistant remember for future planning?", expected_memory_keys=["trip"]),374 FutureQuery(query_id="q2", prompt="What project and teammate should be used when continuing the work plan?", expected_memory_keys=["project", "teammate"]),375 ]376 377 if difficulty == Difficulty.EASY:378 future_queries.append(379 FutureQuery(380 query_id="q3",381 prompt="Which standing preference and planning tool should the assistant honor later?",382 expected_memory_keys=["diet", "tool"],383 )384 )385 else:386 future_queries.extend(387 [388 FutureQuery(389 query_id="q3",390 prompt="Which deadline is the current one after the correction?",391 expected_memory_keys=["deadline_v2"],392 forbidden_memory_keys=["deadline"],393 ),394 FutureQuery(395 query_id="q4",396 prompt="Which latest preference should be respected after the user changed their mind?",397 expected_memory_keys=["diet_v2"],398 forbidden_memory_keys=["diet"],399 ),400 ]401 )402 403 if difficulty == Difficulty.HARD:404 future_queries.append(405 FutureQuery(406 query_id="q5",407 prompt="What mentor detail is still relevant, and what should have been treated as a distraction instead?",408 expected_memory_keys=["mentor"],409 )410 )411 412 return {413 "seed": config.seed if seed is None else seed,414 "scenario_title": config.title,415 "token_budget": config.token_budget,416 "recent_turn_window": config.recent_turn_window,417 "turns": turns,418 "future_queries": future_queries,419 "gold_memories": gold_memories,420 "new_memory_ids_by_turn": new_memory_ids_by_turn,421 "ambiguous_turn_ids": ambiguous_turn_ids,422 "low_value_turn_ids": low_value_turn_ids,423 }424