fix: stabilize benchmark query expansion

Author: Alex-Hunterz

benchmarks/longmemeval/longmemeval_runner.py

@@ -32,7 +32,6 @@
 import asyncio
 import json
 import logging
-import os
 import time
 from dataclasses import dataclass
 from datetime import datetime
@@ -51,8 +50,14 @@
 logger = logging.getLogger(__name__)
 
 
+def _base_question_type(question_type: str) -> str:
+    """Strip LongMemEval abstention suffix while preserving the real task type."""
+    return question_type.removesuffix("_abs")
+
+
 # ── Config ────────────────────────────────────────────────────────────────────
 
+
 @dataclass
 class BenchmarkConfig:
     """Configuration for a LongMemEval benchmark run."""
@@ -68,7 +73,7 @@ class BenchmarkConfig:
     extraction_model: str = "gemini:gemini-2.5-flash-lite"
 
     # Retrieval
-    retrieval_depth: str = "l1"   # l0 | l1 | l2
+    retrieval_depth: str = "l1"  # l0 | l1 | l2
     top_k: int = 10
     context_window: int = 3
 
@@ -86,6 +91,7 @@ class BenchmarkConfig:
 
 # ── Runner ────────────────────────────────────────────────────────────────────
 
+
 class LongMemEvalBenchmark:
     """Main benchmark runner."""
 
@@ -128,7 +134,7 @@ async def _init_vektori(self) -> None:
             extraction_model=self.config.extraction_model,
             default_top_k=self.config.top_k,
             context_window=self.config.context_window,
-            async_extraction=False,   # benchmark needs sync so we can interleave cache ops
+            async_extraction=False,  # benchmark needs sync so we can interleave cache ops
         )
         await self.vektori_client._ensure_initialized()
         self.storage = self.vektori_client.db
@@ -218,10 +224,11 @@ async def _run_questions(self) -> None:
                 if done % 10 == 0 or done == total:
                     logger.info(
                         "Progress: %d/%d questions answered (cache: %d sessions)",
-                        done, total, await self._session_cache.count(),
+                        done,
+                        total,
+                        await self._session_cache.count(),
                     )
 
-
             except Exception as e:
                 logger.error("Question %s failed: %s", qid, e, exc_info=True)
                 self._checkpoint.mark_failed(qid, str(e))
@@ -251,7 +258,9 @@ async def _ingest_question(self, instance: dict[str, Any], user_id: str) -> None
 
             cached_facts = await self._session_cache.get(hsid)
             if cached_facts is not None:
-                await self._replay_session(session, hsid, user_id, session_time, session_date, cached_facts)
+                await self._replay_session(
+                    session, hsid, user_id, session_time, session_date, cached_facts
+                )
             else:
                 new_facts = await self._full_ingest_session(
                     session, hsid, user_id, session_time, session_date
@@ -293,9 +302,7 @@ async def _replay_session(
         )
 
         if inserted_facts:
-            conversation = "\n".join(
-                f"{msg['role'].upper()}: {msg['content']}" for msg in session
-            )
+            conversation = "\n".join(f"{msg['role'].upper()}: {msg['content']}" for msg in session)
             try:
                 await extractor._extract_episodes(
                     inserted_facts=inserted_facts,
@@ -346,16 +353,15 @@ async def _full_ingest_session(
 
     # ── Retrieval + QA ────────────────────────────────────────────────────────
 
-    async def _answer_question(
-        self, instance: dict[str, Any], user_id: str
-    ) -> dict[str, Any]:
+    async def _answer_question(self, instance: dict[str, Any], user_id: str) -> dict[str, Any]:
         qid = instance["question_id"]
         question = instance["question"]
         question_type = instance["question_type"]
+        base_question_type = _base_question_type(question_type)
         question_date = instance.get("question_date") or ""
 
         retrieval_t0 = time.perf_counter()
-        use_expansion = question_type in ("multi-session", "temporal-reasoning")
+        use_expansion = base_question_type in ("multi-session", "temporal-reasoning")
         search_results = await self.vektori_client.search(
             query=question,
             user_id=user_id,
@@ -450,7 +456,7 @@ async def _generate_answer(
 
         llm = create_llm(self.config.eval_model)
         prompt = self._build_qa_prompt(question, context, question_type, question_date)
-        max_tokens = 800 if question_type == "temporal-reasoning" else 500
+        max_tokens = 800 if _base_question_type(question_type) == "temporal-reasoning" else 500
         try:
             return (await llm.generate(prompt, max_tokens=max_tokens)).strip()
         except Exception as e:
@@ -461,6 +467,7 @@ def _build_qa_prompt(
         self, question: str, context: str, question_type: str, question_date: str = ""
     ) -> str:
         date_line = f"TODAY'S DATE: {question_date}\n\n" if question_date else ""
+        base_question_type = _base_question_type(question_type)
 
         abs_hint = ""
         if question_type.endswith("_abs"):
@@ -469,7 +476,7 @@ def _build_qa_prompt(
                 "recognise that the information was never mentioned"
             )
 
-        if question_type == "temporal-reasoning":
+        if base_question_type == "temporal-reasoning":
             return (
                 "You are an AI assistant answering questions based on provided context "
                 "from chat history.\n\n"
@@ -481,7 +488,8 @@ def _build_qa_prompt(
                 "- First, list the relevant dated events from the context\n"
                 "- Then compute the answer (count days/weeks/months between dates)\n"
                 "- If the context does not contain enough information to answer, say "
-                "\"I don't have that information\"\n\n"
+                '"I don\'t have that information"'
+                f"{abs_hint}\n\n"
                 "REASONING:\n"
             )
 
@@ -496,7 +504,7 @@ def _build_qa_prompt(
             "- Be concise and direct — a short phrase or sentence is preferred over a long explanation\n"
             "- For questions about time elapsed, use TODAY'S DATE above as your reference point\n"
             "- If the context does not contain enough information to answer, say "
-            "\"I don't have that information\" — do not guess or infer beyond what is stated"
+            '"I don\'t have that information" — do not guess or infer beyond what is stated'
             f"{abs_hint}\n\n"
             "ANSWER:"
         )
@@ -515,7 +523,10 @@ async def evaluate(self) -> None:
         jsonl_path = self.output_dir / f"{run_name}_qa_results.jsonl"
         with open(jsonl_path, "w", encoding="utf-8") as f:
             for r in qa_results:
-                f.write(json.dumps({"question_id": r["question_id"], "hypothesis": r["hypothesis"]}) + "\n")
+                f.write(
+                    json.dumps({"question_id": r["question_id"], "hypothesis": r["hypothesis"]})
+                    + "\n"
+                )
         logger.info("QA pairs saved to %s", jsonl_path)
 
         self._metrics = self._compute_metrics(qa_results)
@@ -634,22 +645,20 @@ def _print_summary(self) -> None:
             if lat:
                 print("\nLatency (ms):")
                 print(
-                    "  Ingestion      avg={0}  p95={1}".format(
+                    "  Ingestion      avg={}  p95={}".format(
                         lat.get("ingestion_avg"), lat.get("ingestion_p95")
                     )
                 )
                 print(
-                    "  Retrieval      avg={0}  p95={1}".format(
+                    "  Retrieval      avg={}  p95={}".format(
                         lat.get("retrieval_avg"), lat.get("retrieval_p95")
                     )
                 )
                 print(
-                    "  QA generation  avg={0}  p95={1}".format(
-                        lat.get("qa_avg"), lat.get("qa_p95")
-                    )
+                    "  QA generation  avg={}  p95={}".format(lat.get("qa_avg"), lat.get("qa_p95"))
                 )
                 print(
-                    "  Total/question avg={0}  p95={1}".format(
+                    "  Total/question avg={}  p95={}".format(
                         lat.get("total_question_avg"), lat.get("total_question_p95")
                     )
                 )
@@ -667,6 +676,7 @@ def _print_summary(self) -> None:
 
 # ── Helpers ───────────────────────────────────────────────────────────────────
 
+
 def _parse_date(date_str: str) -> datetime | None:
     """Parse a LongMemEval date string like '2023/05/30 (Tue) 23:40'."""
     if not date_str:
@@ -683,6 +693,7 @@ def _parse_date(date_str: str) -> datetime | None:
 
 # ── CLI entry point ───────────────────────────────────────────────────────────
 
+
 async def main() -> None:
     import argparse