Support category-specific prompts

- Add APIConnectionError retry for OpenAI client (server disconnects)
- Add recommendation/preference question guidance to structured prompt
- Instruct model to build on user's existing tools/experiences

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

.gitignore

@@ -33,4 +33,6 @@ logs/
 hindsight-dev/benchmarks/locomo/results/
 hindsight-dev/benchmarks/longmemeval/results/
 hindsight-cli/target
-hindsight-clients/rust/target
+hindsight-clients/rust/target
+
+results/

category_prompts.json

@@ -0,0 +1,3 @@
+{
+  "single-session-preference": "You are a helpful assistant that must answer user questions based on the previous conversations.\n\n{context_instructions}**Answer Guidelines:**\n1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.\n2. Reason about all the memories and find the right answer, considering the most recent memory as an update of the current facts.\n3. If you have 2 possible answers, just say both.\n\nIn general the answer must be comprehensive and plenty of details from the retrieved context.\n\nFor quantitative/counting questions (\"how many...\"): First list each unique item in your reasoning (1. X, 2. Y, 3. Z...), scanning ALL facts, then count them for your answer.\nIf questions asks a location (where...?) make sure to include the location name.\nFor recommendation questions (\"can you recommend...\", \"suggest...\", \"any tips...\"): DO NOT give actual recommendations. Instead, describe what KIND the user would prefer based on their context. Example answer format: \"The user would prefer recommendations for [category] that focus on [their interest]. They would not prefer [what to avoid based on context].\"\nFor questions asking for help or instructions, consider the users' recent memories and previous interactions with the assistant to understand their current situation better (recent purchases, specific product models used..)\nFor specific number/value questions, use the context to understand what is the most up-to-date number based on recency, but also include the reasoning (in the answer) on previous possible values and why you think are less relevant.\nFor open questions, include as much details as possible from different sources that are relevant.\nFor questions where a specific entity/role is mentioned and it's different from your memory, just say the truth, don't make up anything just to fulfill the question. For example, if the question is about a specific sport, you should consider if the memories and the question are about the same sport. (e.g. american football vs soccer, shows vs podcasts)\nFor comparative questions, say you don't know the answer if you don't have information about both sides. (or more sides)\nFor questions related to time/date, carefully review the question date and the memories date to correctly answer the question.\nFor questions related to time/date calculation (e.g. How many days passed between X and Y?), carefully review the memories date to correctly answer the question and only provide an answer if you have information about both X and Y, otherwise say it's not possible to calculate and why.\n\nConsider assistant's previous actions (e.g., bookings, reminders) as impactful to the user experiences.\n\nQuestion: {question}\nQuestion Date: {formatted_question_date}\n\nRetrieved Context:\n{context}\n\nAnswer:"
+}

hindsight-api/hindsight_api/engine/llm_wrapper.py

@@ -5,12 +5,15 @@ import os
 import time
 import asyncio
 from typing import Optional, Any, Dict, List
-from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, LengthFinishReasonError
+from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, APIConnectionError, LengthFinishReasonError
 from google import genai
 from google.genai import types as genai_types
 from google.genai import errors as genai_errors
 import logging
 
+# Seed applied to every Groq request for deterministic behavior.
+DEFAULT_LLM_SEED = 4242
+
 logger = logging.getLogger(__name__)
 
 # Disable httpx logging
@@ -40,6 +43,7 @@ class LLMConfig:
         api_key: str,
         base_url: str,
         model: str,
+            reasoning_effort: str = "low",
     ):
         """
         Initialize LLM configuration.
@@ -54,6 +58,7 @@ class LLMConfig:
         self.api_key = api_key
         self.base_url = base_url
         self.model = model
+        self.reasoning_effort = reasoning_effort
 
         # Validate provider
         if self.provider not in ["openai", "groq", "ollama", "gemini"]:
@@ -136,10 +141,14 @@ class LLMConfig:
                 "messages": messages,
                 **kwargs
             }
+
+            if self.provider == "groq":
+                call_params["seed"] = DEFAULT_LLM_SEED
+            
             if self.provider == "groq":
                 call_params["extra_body"] = {
                     "service_tier": "auto",
-                    "reasoning_effort": "low",  # Reduce reasoning overhead
+                    "reasoning_effort": self.reasoning_effort,
                     "include_reasoning": False,  # Disable hidden reasoning tokens
                 }
 
@@ -202,6 +211,18 @@ class LLMConfig:
                         f"LLM output exceeded token limits. Input may need to be split into smaller chunks."
                     ) from e
 
+                except APIConnectionError as e:
+                    # Handle connection errors (server disconnected, network issues) with retry
+                    last_exception = e
+                    if attempt < max_retries:
+                        logger.warning(f"Connection error, retrying... (attempt {attempt + 1}/{max_retries + 1})")
+                        backoff = min(initial_backoff * (2 ** attempt), max_backoff)
+                        await asyncio.sleep(backoff)
+                        continue
+                    else:
+                        logger.error(f"Connection error after {max_retries + 1} attempts: {str(e)}")
+                        raise
+
                 except APIStatusError as e:
                     last_exception = e
                     if attempt < max_retries:
@@ -238,7 +259,7 @@ class LLMConfig:
         skip_validation: bool,
         start_time: float,
         **kwargs
-    ) -> Any:
+) -> Any:
         """Handle Gemini-specific API calls using google-genai SDK."""
         import json
 
@@ -287,6 +308,8 @@ class LLMConfig:
             config_kwargs['max_output_tokens'] = kwargs['max_tokens']
         if response_format is not None:
             config_kwargs['response_mime_type'] = 'application/json'
+            # Pass the Pydantic model directly as response_schema for structured output
+            config_kwargs['response_schema'] = response_format
 
         generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
 
@@ -302,6 +325,23 @@ class LLMConfig:
 
                 content = response.text
 
+                # Handle empty/None response (can happen with content filtering or timeouts)
+                if content is None:
+                    # Check if there's a block reason
+                    block_reason = None
+                    if hasattr(response, 'candidates') and response.candidates:
+                        candidate = response.candidates[0]
+                        if hasattr(candidate, 'finish_reason'):
+                            block_reason = candidate.finish_reason
+
+                    if attempt < max_retries:
+                        logger.warning(f"Gemini returned empty response (reason: {block_reason}), retrying... (attempt {attempt + 1}/{max_retries + 1})")
+                        backoff = min(initial_backoff * (2 ** attempt), max_backoff)
+                        await asyncio.sleep(backoff)
+                        continue
+                    else:
+                        raise RuntimeError(f"Gemini returned empty response after {max_retries + 1} attempts (reason: {block_reason})")
+
                 if response_format is not None:
                     # Parse the JSON response
                     json_data = json.loads(content)
@@ -326,6 +366,18 @@ class LLMConfig:
 
                 return result
 
+            except json.JSONDecodeError as e:
+                # Handle truncated JSON responses (often from MAX_TOKENS) with retry
+                last_exception = e
+                if attempt < max_retries:
+                    logger.warning(f"Gemini returned invalid JSON (truncated response?), retrying... (attempt {attempt + 1}/{max_retries + 1})")
+                    backoff = min(initial_backoff * (2 ** attempt), max_backoff)
+                    await asyncio.sleep(backoff)
+                    continue
+                else:
+                    logger.error(f"Gemini returned invalid JSON after {max_retries + 1} attempts: {str(e)}")
+                    raise
+
             except genai_errors.APIError as e:
                 # Handle rate limits and server errors with retry
                 if e.code in (429, 503, 500):
@@ -372,6 +424,37 @@ class LLMConfig:
             api_key=api_key,
             base_url=base_url,
             model=model,
+            reasoning_effort="low"
+        )
+
+    @classmethod
+    def for_answer_generation(cls) -> "LLMConfig":
+        """
+        Create configuration for answer generation operations from environment variables.
+
+        Falls back to memory LLM config if answer-specific config not set.
+        """
+        # Check if answer-specific config exists, otherwise fall back to memory config
+        provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
+        api_key = os.getenv("HINDSIGHT_API_ANSWER_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY"))
+        base_url = os.getenv("HINDSIGHT_API_ANSWER_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL"))
+        model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
+
+        # Set default base URL if not provided
+        if not base_url:
+            if provider == "groq":
+                base_url = "https://api.groq.com/openai/v1"
+            elif provider == "ollama":
+                base_url = "http://localhost:11434/v1"
+            else:
+                base_url = ""
+
+        return cls(
+            provider=provider,
+            api_key=api_key,
+            base_url=base_url,
+            model=model,
+            reasoning_effort="high"
         )
 
     @classmethod
@@ -401,4 +484,5 @@ class LLMConfig:
             api_key=api_key,
             base_url=base_url,
             model=model,
+            reasoning_effort="high"
         )

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -2965,8 +2965,6 @@ Guidelines:
                         uuid.UUID(obs_id), uuid.UUID(entity_id)
                     )
 
-        # Single consolidated log line
-        logger.info(f"[OBSERVATIONS] {entity_name}: {len(facts)} facts -> {len(created_ids)} observations")
         return created_ids
 
     async def _regenerate_observations_sync(

hindsight-api/hindsight_api/engine/retain/orchestrator.py

@@ -313,15 +313,6 @@ async def retain_batch(
                 contents, extracted_facts, is_duplicate_flags, unit_ids
             )
 
-            total_time = time.time() - start_time
-            log_buffer.append(f"{'='*60}")
-            log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(unit_ids)} units in {total_time:.3f}s")
-            if document_ids_added:
-                log_buffer.append(f"Documents: {', '.join(document_ids_added)}")
-            log_buffer.append(f"{'='*60}")
-
-            logger.info("\n" + "\n".join(log_buffer) + "\n")
-
         # Trigger background tasks AFTER transaction commits
         await _trigger_background_tasks(
             task_backend,
@@ -329,9 +320,20 @@ async def retain_batch(
             bank_id,
             unit_ids,
             non_duplicate_facts,
-            entity_links
+            entity_links,
+            log_buffer
         )
 
+        # Log final summary
+        total_time = time.time() - start_time
+        log_buffer.append(f"{'='*60}")
+        log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(unit_ids)} units in {total_time:.3f}s")
+        if document_ids_added:
+            log_buffer.append(f"Documents: {', '.join(document_ids_added)}")
+        log_buffer.append(f"{'='*60}")
+
+        logger.info("\n" + "\n".join(log_buffer) + "\n")
+
         return result_unit_ids
 
 
@@ -371,7 +373,8 @@ async def _trigger_background_tasks(
     bank_id: str,
     unit_ids: List[str],
     facts: List[ProcessedFact],
-    entity_links: List
+    entity_links: List,
+    log_buffer: List[str] = None
 ) -> None:
     """Trigger opinion reinforcement and observation regeneration (sync)."""
     # Trigger opinion reinforcement if there are entities
@@ -392,14 +395,19 @@ async def _trigger_background_tasks(
     if entity_links and regenerate_observations_fn:
         unique_entity_ids = set()
         for link in entity_links:
-            # links are tuples: (unit_id, entity_id, confidence)
-            if len(link) >= 2 and link[1]:
-                unique_entity_ids.add(str(link[1]))
+            # links are tuples: (from_unit_id, to_unit_id, link_type, weight, entity_id)
+            if len(link) >= 5 and link[4]:
+                unique_entity_ids.add(str(link[4]))
 
         if unique_entity_ids:
+            entities_to_process = list(unique_entity_ids)[:TOP_N_ENTITIES]
+            obs_start = time.time()
             # Run observation regeneration synchronously
             await regenerate_observations_fn(
                 bank_id=bank_id,
-                entity_ids=list(unique_entity_ids)[:TOP_N_ENTITIES],
+                entity_ids=entities_to_process,
                 min_facts=MIN_FACTS_THRESHOLD
             )
+            obs_time = time.time() - obs_start
+            if log_buffer is not None:
+                log_buffer.append(f"[11] Observations: {len(entities_to_process)} entities in {obs_time:.3f}s")

hindsight-control-plane/src/app/api/entities/[entityId]/route.ts

@@ -28,6 +28,13 @@ export async function GET(
       }
     });
 
+    if (response.error) {
+      return NextResponse.json(
+        { error: response.error },
+        { status: 500 }
+      );
+    }
+
     return NextResponse.json(response.data, { status: 200 });
   } catch (error) {
     console.error('Error getting entity:', error);

hindsight-control-plane/src/app/api/entities/route.ts

@@ -21,6 +21,13 @@ export async function GET(request: NextRequest) {
       query: { limit }
     });
 
+    if (response.error) {
+      return NextResponse.json(
+        { error: response.error },
+        { status: 500 }
+      );
+    }
+
     return NextResponse.json(response.data, { status: 200 });
   } catch (error) {
     console.error('Error listing entities:', error);

hindsight-control-plane/src/components/entities-view.tsx

@@ -164,6 +164,7 @@ export function EntitiesView() {
           </div>
 
           <div className="text-sm text-muted-foreground mb-4">
+            <div className="font-mono text-xs mb-1" title={selectedEntity.id}>ID: {selectedEntity.id}</div>
             <div>Mentions: {selectedEntity.mention_count}</div>
             <div>First seen: {formatDate(selectedEntity.first_seen)}</div>
             <div>Last seen: {formatDate(selectedEntity.last_seen)}</div>

hindsight-dev/benchmarks/common/benchmark_runner.py

@@ -38,6 +38,55 @@ import os
 console = Console()
 
 
+def get_model_config() -> Dict[str, Dict[str, str]]:
+    """
+    Get the model configuration for all three LLM roles.
+
+    Reads directly from environment variables without instantiating LLM clients.
+
+    Returns:
+        Dict with 'hindsight', 'answer_generation', and 'judge' keys,
+        each containing 'provider' and 'model' info.
+    """
+    # Memory/Hindsight config (base config)
+    memory_provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
+    memory_model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
+
+    # Answer generation config (falls back to memory config)
+    answer_provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", memory_provider)
+    answer_model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", memory_model)
+
+    # Judge config (falls back to memory config)
+    judge_provider = os.getenv("HINDSIGHT_API_JUDGE_LLM_PROVIDER", memory_provider)
+    judge_model = os.getenv("HINDSIGHT_API_JUDGE_LLM_MODEL", memory_model)
+
+    return {
+        'hindsight': {
+            'provider': memory_provider,
+            'model': memory_model,
+        },
+        'answer_generation': {
+            'provider': answer_provider,
+            'model': answer_model,
+        },
+        'judge': {
+            'provider': judge_provider,
+            'model': judge_model,
+        }
+    }
+
+
+def print_model_config():
+    """Print the model configuration to console."""
+    config = get_model_config()
+
+    console.print("\n[bold cyan]Model Configuration:[/bold cyan]")
+    console.print(f"  Hindsight:         {config['hindsight']['provider']}/{config['hindsight']['model']}")
+    console.print(f"  Answer Generation: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+    console.print(f"  LLM Judge:         {config['judge']['provider']}/{config['judge']['model']}")
+    console.print()
+
+
 async def create_memory_engine() -> MemoryEngine:
     """
     Create and initialize a MemoryEngine instance from environment variables.
@@ -411,16 +460,26 @@ class BenchmarkRunner:
             # Use MemoryEngine directly
             # Map thinking_budget to budget level
             budget = Budget.LOW if thinking_budget <= 30 else Budget.MID if thinking_budget <= 70 else Budget.HIGH
+
+            import time
+            recall_start_time = time.time()
             search_result = await self.memory.recall_async(
                 bank_id=agent_id,
                 query=question,
                 budget=budget,
                 max_tokens=max_tokens,
-                fact_type=["world", "bank"],
+                fact_type=["world", "experience"],
                 question_date=question_date,
                 include_entities=True,
                 include_chunks=True
             )
+            recall_time = time.time() - recall_start_time
+
+            # Log recall stats
+            num_results = len(search_result.results) if search_result.results else 0
+            num_chunks = len(search_result.chunks) if search_result.chunks else 0
+            num_entities = len(search_result.entities) if search_result.entities else 0
+            logging.info(f"Recall stats: {num_results} facts, {num_chunks} chunks, {num_entities} entities in {recall_time:.2f}s")
 
             # Convert entire RecallResult to dictionary for answer generation
             recall_result_dict = search_result.model_dump()
@@ -522,6 +581,7 @@ class BenchmarkRunner:
                             'reasoning': reasoning,
                             'category': category,
                             'retrieved_memories': memories_without_embeddings,
+                            'chunks': chunks,
                             'is_invalid': False,
                             'error': None
                         }
@@ -536,6 +596,7 @@ class BenchmarkRunner:
                             'reasoning': f'Error: {str(e)}',
                             'category': category,
                             'retrieved_memories': [],
+                            'chunks': {},
                             'is_invalid': True,
                             'error': str(e)
                         }
@@ -780,6 +841,9 @@ class BenchmarkRunner:
         console.print(f"\n[bold cyan]Benchmark Evaluation[/bold cyan]")
         console.print("=" * 80)
 
+        # Print model configuration
+        print_model_config()
+
         # Load dataset
         console.print(f"\n[1] Loading dataset from {dataset_path}...")
         items = self.dataset.load(dataset_path, max_items)
@@ -869,6 +933,7 @@ class BenchmarkRunner:
             'total_invalid': total_invalid,
             'total_valid': total_valid,
             'num_items': len(items),
+            'model_config': get_model_config(),
             'item_results': all_results
         }
 
@@ -1130,6 +1195,15 @@ class BenchmarkRunner:
         """Display benchmark results in a formatted table."""
         console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")
 
+        # Display model configuration
+        if 'model_config' in results:
+            config = results['model_config']
+            console.print("[bold cyan]Model Configuration:[/bold cyan]")
+            console.print(f"  Hindsight:         {config['hindsight']['provider']}/{config['hindsight']['model']}")
+            console.print(f"  Answer Generation: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+            console.print(f"  LLM Judge:         {config['judge']['provider']}/{config['judge']['model']}")
+            console.print()
+
         # Display results table
         table = Table(title="Benchmark Results", box=box.ROUNDED)
         table.add_column("Item ID", style="cyan")
@@ -1244,6 +1318,7 @@ class BenchmarkRunner:
             'total_invalid': total_invalid,
             'total_valid': total_valid,
             'num_items': len(all_results),
+            'model_config': get_model_config(),
             'item_results': all_results
         }
 

hindsight-dev/benchmarks/locomo/locomo_benchmark.py

@@ -103,8 +103,8 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
     """LoComo-specific answer generator using configurable LLM provider."""
 
     def __init__(self):
-        """Initialize with LLM configuration for memory operations."""
-        self.llm_config = LLMConfig.for_memory()
+        """Initialize with LLM configuration for answer generation."""
+        self.llm_config = LLMConfig.for_answer_generation()
         self.client = self.llm_config._client
         self.model = self.llm_config.model
 
@@ -444,6 +444,17 @@ def generate_markdown_table(results: dict, use_think: bool = False):
     mode_str = " (Think Mode)" if use_think else ""
     lines.append(f"# LoComo Benchmark Results{mode_str}")
     lines.append("")
+
+    # Add model configuration
+    if 'model_config' in results:
+        config = results['model_config']
+        lines.append("## Model Configuration")
+        lines.append("")
+        lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
+        lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+        lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
+        lines.append("")
+
     lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
     lines.append("")
     lines.append("| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |")

hindsight-dev/benchmarks/longmemeval/docs/judge_prompt_comparison.md

@@ -0,0 +1,81 @@
+# LongMemEval Judge Prompt Comparison: Original Paper vs Hindsight
+
+## 1. `single-session-user`, `single-session-assistant`, `multi-session`
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. If the response only contains a subset of the information required by the answer, answer no. | Evaluate if the model response contains the correct answer to the question. |
+| | I will give you a question, a correct answer, and a response from a model. Please set correct=true if the response contains the correct answer. Otherwise, set correct=no. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also set correct=true. If the response only contains a subset of the information required by the answer, set correct=false |
+| Question: {question} | Question: {question} |
+| Correct Answer: {answer} | Correct Answer: {correct_answer} |
+| Model Response: {response} | Model Response: {predicted_answer} |
+| Is the model response correct? Answer yes or no only. | Evaluation criteria: |
+| | - Set correct=true if the response contains the correct answer |
+| | - Set correct=true if the response is equivalent to the correct answer or contains intermediate steps |
+| | - Set correct=false if the response is incorrect or missing key information |
+| | Provide your evaluation as JSON with: |
+| | - reasoning: One sentence explanation |
+| | - correct: true or false |
+
+---
+
+## 2. `temporal-reasoning`
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. If the response only contains a subset of the information required by the answer, answer no. In addition, do not penalize off-by-one errors for the number of days. If the question asks for the number of days/weeks/months, etc., and the model makes off-by-one errors (e.g., predicting 19 days when the answer is 18), the model's response is still correct. | I will give you a question, a correct answer, and a response from a model. Please set correct=true if the response contains the correct answer. Otherwise, set correct=false. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also set correct=true. If the response only contains a subset of the information required by the answer, answer correct=false. In addition, do not penalize off-by-one errors for the number of days. If the question asks for the number of days/weeks/months, etc., and the model makes off-by-one errors (e.g., predicting 19 days when the answer is 18), the model's response is still correct. |
+| Question: {question} | Question: {question} |
+| Correct Answer: {answer} | Gold answer: {correct_answer} |
+| Model Response: {response} | Generated answer: {predicted_answer} |
+| Is the model response correct? Answer yes or no only. | First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred. If it's correct, set correct=true. |
+
+---
+
+## 3. `knowledge-update`
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response contains some previous information along with an updated answer, the response should be considered as correct as long as the updated answer is the required answer. | I will give you a question, a correct answer, and a response from a model. Please set correct=true if the response contains the correct answer. Otherwise, set correct=false. If the response contains some previous information along with an updated answer, the response should be considered as correct as long as the updated answer is the required answer. |
+| Question: {question} | Question: {question} |
+| Correct Answer: {answer} | Gold answer: {correct_answer} |
+| Model Response: {response} | Generated answer: {predicted_answer} |
+| Is the model response correct? Answer yes or no only. | First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred. If it's correct, set correct=true. |
+
+---
+
+## 4. `single-session-preference`
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| I will give you a question, a rubric for desired personalized response, and a response from a model. Please answer yes if the response satisfies the desired response. Otherwise, answer no. The model does not need to reflect all the points in the rubric. The response is correct as long as it recalls and utilizes the user's personal information correctly. | I will give you a question, a answer for desired personalized response, and a response from a model. Please set correct=true if the response satisfies the desired response. Otherwise, set correct=false. The model does not need to reflect all the points in the desired response. The response is correct as long as it recalls and utilizes the user's personal information correctly. |
+| Question: {question} | Question: {question} |
+| Rubric: {rubric} | Gold answer: {correct_answer} |
+| Model Response: {response} | Generated answer: {predicted_answer} |
+| Is the model response correct? Answer yes or no only. | First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred. If it's correct, set correct=true. |
+
+---
+
+## 5. `unanswerable` (abstention)
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| I will give you an unanswerable question, an explanation, and a response from a model. Please answer yes if the model correctly identifies the question as unanswerable. The model could say that the information is incomplete, or some other information is given but the asked information is not. | *Not implemented* |
+| Question: {question} | |
+| Explanation: {explanation} | |
+| Model Response: {response} | |
+| Does the model correctly identify the question as unanswerable? Answer yes or no only. | |
+
+---
+
+## 6. Default (fallback for unknown categories)
+
+| Original Paper | Hindsight |
+|----------------|-----------|
+| *No default - all categories have specific prompts* | Your task is to label an answer to a question as 'CORRECT' or 'WRONG'. You will be given the following data: (1) a question (posed by one user to another user), (2) a 'gold' (ground truth) answer, (3) a generated answer which you will score as CORRECT/WRONG. |
+| | The point of the question is to ask about something one user should know about the other user based on their prior conversations. The gold answer will usually be a concise and short answer that includes the referenced topic, for example: Question: Do you remember what I got the last time I went to Hawaii? Gold answer: A shell necklace The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT. |
+| | For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date. |
+| | There's an edge case where the actual answer can't be found in the data and in that case the gold answer will say so (e.g. 'You did not mention this information.'); if the generated answer says that it cannot be answered or it doesn't know all the details, it should be counted as CORRECT. |
+| | Question: {question} |
+| | Gold answer: {correct_answer} |
+| | Generated answer: {predicted_answer} |
+| | First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred. If it's correct, set correct=true. |

hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py

@@ -128,16 +128,239 @@ class LongMemEvalDataset(BenchmarkDataset):
 
 class QuestionAnswer(pydantic.BaseModel):
     answer: str
-    reasoning: str
+    reasoning: Optional[str] = None
 
 class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
     """LongMemEval-specific answer generator using configurable LLM provider."""
 
-    def __init__(self):
-        """Initialize with LLM configuration for memory operations."""
-        self.llm_config = LLMConfig.for_judge()
+    def __init__(self, context_format: str = "json", category_prompts: Optional[Dict[str, str]] = None):
+        """Initialize with LLM configuration for answer generation.
+
+        Args:
+            context_format: How to format the retrieved context. Options:
+                - "json": Raw JSON dump of recall_result (original behavior)
+                - "structured": Human-readable format with facts grouped with source chunks
+            category_prompts: Optional dict mapping category names to custom prompt templates.
+                Keys should be category names (e.g., 'single-session-preference').
+                Values should be prompt templates with placeholders: {context_instructions}, {question}, {formatted_question_date}, {context}
+                If None or category not found, uses default prompt.
+        """
+        self.llm_config = LLMConfig.for_answer_generation()
         self.client = self.llm_config._client
         self.model = self.llm_config.model
+        self.context_format = context_format
+        self.category_prompts = category_prompts or {}
+
+    def _format_context_json(self, recall_result: Dict[str, Any]) -> str:
+        """Original JSON dump format."""
+        return json.dumps(recall_result)
+
+    def _format_context_structured(self, recall_result: Dict[str, Any]) -> str:
+        """Human-readable format with facts grouped with their source chunks.
+
+        Format:
+            Fact 1: [fact text]
+            When: [date]
+            Source:
+              "[chunk text]"
+
+            ---
+
+            Fact 2: ...
+
+            === Entity Observations ===
+            Entity: [name]
+            - [observation 1]
+            - [observation 2]
+        """
+        results = recall_result.get("results", [])
+        chunks = recall_result.get("chunks", {})
+        entities = recall_result.get("entities", {})
+
+        if not results and not entities:
+            return "No memories found."
+
+        formatted_parts = []
+
+        for i, fact in enumerate(results, 1):
+            fact_text = fact.get("text", "")
+            fact_type = fact.get("fact_type", "unknown")
+
+            # Extract temporal information
+            occurred_start = fact.get("occurred_start")
+            occurred_end = fact.get("occurred_end")
+            mentioned_at = fact.get("mentioned_at")
+
+            # Build temporal string
+            when_parts = []
+            if occurred_start:
+                when_parts.append(f"occurred: {occurred_start}")
+            if mentioned_at:
+                when_parts.append(f"mentioned: {mentioned_at}")
+            when_str = " | ".join(when_parts) if when_parts else "unknown"
+
+            # Get the source chunk if available
+            chunk_id = fact.get("chunk_id")
+            chunk_text = None
+            if chunk_id and chunk_id in chunks:
+                chunk_info = chunks[chunk_id]
+                chunk_text = chunk_info.get("chunk_text", "")
+
+            # Build the formatted fact entry
+            entry_parts = [
+                f"Fact {i} ({fact_type}): {fact_text}",
+                f"When: {when_str}"
+            ]
+
+            # Add context field if present
+            context = fact.get("context")
+            if context:
+                entry_parts.append(f"Context: {context}")
+
+            # Add source chunk
+            if chunk_text:
+                # Truncate very long chunks
+                if len(chunk_text) > 1000:
+                    chunk_text = chunk_text[:1000] + "..."
+                entry_parts.append(f"Source chunk:\n  \"{chunk_text}\"")
+
+            formatted_parts.append("\n".join(entry_parts))
+
+        # Add entity observations section if present
+        if entities:
+            entity_parts = ["=== Entity Observations ==="]
+            for entity_name, entity_state in entities.items():
+                observations = entity_state.get("observations", [])
+                if observations:
+                    entity_parts.append(f"\nEntity: {entity_name}")
+                    for obs in observations:
+                        obs_text = obs.get("text", "")
+                        entity_parts.append(f"  - {obs_text}")
+            if len(entity_parts) > 1:  # More than just the header
+                formatted_parts.append("\n".join(entity_parts))
+
+        return "\n\n---\n\n".join(formatted_parts)
+
+    def _get_default_prompt_template(self) -> str:
+        """Get the default prompt template for answer generation."""
+        return """You are a helpful assistant that must answer user questions based on the previous conversations.
+
+{context_instructions}**Answer Guidelines:**
+1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
+2. Reason about all the memories and find the right answer, considering the most recent memory as an update of the current facts.
+3. If you have 2 possible answers, just say both.
+
+In general the answer must be comprehensive and plenty of details from the retrieved context.
+
+For quantitative/counting questions ("how many..."): First list each unique item in your reasoning (1. X, 2. Y, 3. Z...), scanning ALL facts, then count them for your answer.
+If questions asks a location (where...?) make sure to include the location name.
+For recommendation questions ("can you recommend...", "suggest...", "any tips..."): DO NOT give actual recommendations. Instead, describe what KIND the user would prefer based on their context. Example answer format: "The user would prefer recommendations for [category] that focus on [their interest]. They would not prefer [what to avoid based on context]."
+For questions asking for help or instructions, consider the users' recent memories and previous interactions with the assistant to understand their current situation better (recent purchases, specific product models used..)
+For specific number/value questions, use the context to understand what is the most up-to-date number based on recency, but also include the reasoning (in the answer) on previous possible values and why you think are less relevant.
+For open questions, include as much details as possible from different sources that are relevant.
+For questions where a specific entity/role is mentioned and it's different from your memory, just say the truth, don't make up anything just to fulfill the question. For example, if the question is about a specific sport, you should consider if the memories and the question are about the same sport. (e.g. american football vs soccer, shows vs podcasts)
+For comparative questions, say you don't know the answer if you don't have information about both sides. (or more sides)
+For questions related to time/date, carefully review the question date and the memories date to correctly answer the question.
+For questions related to time/date calculation (e.g. How many days passed between X and Y?), carefully review the memories date to correctly answer the question and only provide an answer if you have information about both X and Y, otherwise say it's not possible to calculate and why.
+
+Consider assistant's previous actions (e.g., bookings, reminders) as impactful to the user experiences.
+
+Question: {question}
+Question Date: {formatted_question_date}
+
+Retrieved Context:
+{context}
+
+Answer:"""
+
+    def _get_context_instructions(self) -> str:
+        """Get instructions for interpreting the context based on format."""
+        if self.context_format == "structured":
+            return """**Understanding the Retrieved Context:**
+The context contains memory facts extracted from previous conversations, each with its source chunk.
+
+1. **Fact**: A high-level summary/atomic fact (e.g., "User loves hiking in mountains")
+   - This is the searchable summary of what was stored
+
+2. **Source Chunk**: The actual raw conversation where the fact was extracted from
+   - **This is your primary source for detailed information**
+   - Look here for specifics, context, quotes, and evidence
+   - Prioritize information from chunks when facts seem ambiguous
+
+3. **Temporal Information**:
+   - "occurred": When the event actually happened
+   - "mentioned": When it was discussed in conversation
+   - Use this to understand the timeline and resolve conflicts (prefer more recent info)
+
+4. **Context**: Additional metadata about the conversation session
+
+**Date Calculations (CRITICAL - read carefully):**
+- When calculating days between two dates: count the days from Date A to Date B as (B - A)
+- Example: Jan 1 to Jan 8 = 7 days (not 8)
+- "X days ago" from Question Date means: Question Date minus X days
+- When a fact says "three weeks ago" on a certain mentioned date, that refers to 3 weeks before THAT mentioned date, NOT the question date
+- Always convert relative times ("last Friday", "two weeks ago") to absolute dates BEFORE comparing
+- Double-check your arithmetic - off-by-one errors are very common
+- **Important**: Read questions carefully for time anchors. "How many days ago did X happen when Y happened?" asks for the time between X and Y, NOT between X and the question date
+
+**Handling Relative Times in Facts:**
+- If a fact says "last Friday" or "two weeks ago", anchor it to the fact's "mentioned" date, NOT the question date
+- First convert ALL relative references to absolute dates, then answer the question
+- Show your date conversion work in your reasoning
+
+**Counting Questions (CRITICAL for "how many" questions):**
+- **Scan ALL facts first** - go through every single fact before counting, don't stop early
+- **List each item explicitly in your reasoning** before giving the count: "1. X, 2. Y, 3. Z = 3 total"
+- **Check all facts and chunks** before giving your final count
+- **Watch for duplicates**: The same item may appear in multiple facts. Deduplicate by checking if two facts refer to the same underlying item/event
+- **Watch for different descriptions of same thing**: "Dr. Patel (ENT specialist)" and "the ENT specialist" might be the same doctor
+- **Don't over-interpret**: A project you "completed" is different from a project you're "leading"
+- **Don't double-count**: If the same charity event is mentioned in two conversations, it's still one event
+
+**Disambiguation Guidance (CRITICAL - many errors come from over-counting):**
+- **Assume overlap by default**: If two facts describe similar events (same type, similar timeframe, similar details), assume they are the SAME event unless there's clear evidence they are different
+- If a person has a name AND a role mentioned, check if they're the same person before counting separately
+- If an amount is mentioned multiple times on different dates, check if it's the same event or different events
+- When facts reference the same underlying event from different sessions, count it once
+- **Check for aliases**: "my college roommate's wedding" and "Emily's wedding" might be the same event
+- **Check for time period overlap**: Two "week-long breaks" mentioned in overlapping time periods are likely the same break
+- **When in doubt, undercount**: It's better to miss a duplicate than to count the same thing twice
+
+**Question Interpretation (read carefully):**
+- "How many X before Y?" - count only X that happened BEFORE Y, not Y itself
+- "How many properties viewed before making an offer on Z?" - count OTHER properties, not Z
+- "How many X in the last week/month?" - calculate the exact date range from the question date, then filter
+- Pay attention to qualifiers like "before", "after", "initially", "currently", "in total"
+
+**When to Say "I Don't Know":**
+- If the question asks about something not in the retrieved context, say "I don't have information about X"
+- If comparing two things (e.g., "which happened first, X or Y?") but only one is mentioned, explicitly say the other is missing
+- Don't guess or infer dates that aren't explicitly stated in the facts or chunks
+- If you cannot find a specific piece of information after checking all facts and chunks, admit it
+- **Partial knowledge is OK**: If asked about two things and you only have info on one, provide what you know and note what's missing (don't just say "I don't know")
+
+**For Recommendation/Preference Questions (tips, suggestions, advice):**
+- **DO NOT invent specific recommendations** (no made-up product names, course names, paper titles, channel names, etc.)
+- **DO mention specific brands/products the user ALREADY uses** from the context
+- Describe WHAT KIND of recommendation the user would prefer, referencing their existing tools/brands
+- Keep answers concise - focus on key preferences (brand, quality level, specific interests) not exhaustive category lists
+- First scan ALL facts for user's existing tools, brands, stated preferences
+
+**How to Answer:**
+1. Scan ALL facts to find relevant memories - don't stop after finding a few
+2. **Read the source chunks carefully** - they contain the actual details you need
+3. Convert all relative times to absolute dates
+4. Use temporal information to understand when things happened
+5. Synthesize information from multiple facts if needed
+6. If facts conflict, prefer more recent information
+7. Double-check any date calculations before answering
+8. **For counting questions ("how many")**: First list each unique item in your reasoning (1. X, 2. Y, 3. Z...), then count them
+9. **For recommendations**: Reference the user's existing tools, experiences, or preferences explicitly
+
+"""
+        else:
+            # Original JSON format - minimal instructions
+            return ""
 
     async def generate_answer(
                 self,
@@ -159,56 +382,49 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
                 Tuple of (answer, reasoning, None)
                 - None indicates to use the memories from recall_result
             """
-            context = json.dumps(recall_result)
+            # Format context based on selected mode
+            if self.context_format == "structured":
+                context = self._format_context_structured(recall_result)
+            else:
+                context = self._format_context_json(recall_result)
+
+            context_instructions = self._get_context_instructions()
 
             # Format question date if provided
             formatted_question_date = question_date.strftime('%Y-%m-%d %H:%M:%S UTC') if question_date else "Not specified"
 
+            # Select prompt template based on category
+            if question_type and question_type in self.category_prompts:
+                prompt_template = self.category_prompts[question_type]
+            else:
+                prompt_template = self._get_default_prompt_template()
+
+            # Format the prompt with context
+            prompt_content = prompt_template.format(
+                context_instructions=context_instructions,
+                question=question,
+                formatted_question_date=formatted_question_date,
+                context=context
+            )
+
             # Use LLM to generate answer
             try:
                 answer_obj = await self.llm_config.call(
                     messages=[
                         {
                             "role": "user",
-                            "content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
-
-**How to Answer:**
-1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
-2. Reason about all the memories and find the right answer, considering the most recent memory as an update of the current facts.
-3. If you have 2 possible answers, just say both.
-
-In general the answer must be comprehensive and plenty of details from the retrieved context.
-
-For quantitative questions, use numbers and units. Example: 'How many..', just answer the number and which ones. Consider EACH item even if it's not the most recent one. Reason and do calculation for complex questions.
-If questions asks a location (where...?) make sure to include the location name.
-For recommendations/suggestions, use the retrieved context to understand the user's preferences and user's personal experiences, and provide a possible answer based on those. Include the reasoning and explicitly say what the user prefers, before making suggestions (user previous experiences or specific requests FROM the user). Consider as much user preferences as possible in your answer.
-For questions asking for help or instructions, consider the users' recent memories and previous interactions with the assistant to understand their current situation better (recent purchases, specific product models used..) 
-For specific number/value questions, use the context to understand what is the most up-to-date number based on recency, but also include the reasoning (in the answer) on previous possible values and why you think are less relevant.
-For open questions, include as much details as possible from different sources that are relevant.
-For questions where a specific entity/role is mentioned and it's different from your memory, just say the truth, don't make up anything just to fulfill the question. For example, if the question is about a specific sport, you should consider if the memories and the question are about the same sport. (e.g. american football vs soccer, shows vs podcasts)
-For comparative questions, say you don't know the answer if you don't have information about both sides. (or more sides) 
-For questions related to time/date, carefully review the question date and the memories date to correctly answer the question.
-For questions related to time/date calculation (e.g. How many days passed between X and Y?), carefully review the memories date to correctly answer the question and only provide an answer if you have information about both X and Y, otherwise say it's not possible to calculate and why.
-
-Consider assistant's previous actions (e.g., bookings, reminders) as impactful to the user experiences.
-
-
-Question: {question}
-Question Date: {formatted_question_date}
-
-Retrieved Context:
-{context}
-
-
-Answer:
-"""
+                            "content": prompt_content
                         }
                     ],
                     response_format=QuestionAnswer,
                     scope="memory",
-                    max_tokens=8192,
+                    max_tokens=32768,
                 )
-                return answer_obj.answer, answer_obj.reasoning + " (question date: " + formatted_question_date + ")", None
+                reasoning_text = answer_obj.reasoning or ""
+                if reasoning_text:
+                    reasoning_text = reasoning_text + " "
+                reasoning_text += f"(question date: {formatted_question_date})"
+                return answer_obj.answer, reasoning_text, None
             except Exception as e:
                 return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
 
@@ -227,7 +443,10 @@ async def run_benchmark(
     only_ingested: bool = False,
     category: str = None,
     max_concurrent_items: int = 1,
-    results_filename: str = "benchmark_results.json"
+    results_filename: str = "benchmark_results.json",
+    context_format: str = "json",
+    source_results: str = None,
+    category_prompts: Optional[Dict[str, str]] = None
 ):
     """
     Run the LongMemEval benchmark.
@@ -247,14 +466,18 @@ async def run_benchmark(
         category: Optional category to filter questions (e.g., 'single-session-user', 'multi-session', 'temporal-reasoning'). Mutually exclusive with max_instances and max_instances_per_category.
         max_concurrent_items: Maximum number of instances to process in parallel (default: 1 for sequential)
         results_filename: Filename for results (default: benchmark_results.json). Directory is fixed to results/.
+        context_format: How to format context for answer generation. "json" (raw JSON) or "structured" (human-readable with facts+chunks).
+        source_results: Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json.
+        category_prompts: Optional dict mapping category names to custom prompt templates. If provided, questions in matching categories will use the custom prompt instead of the default.
     """
     from rich.console import Console
     console = Console()
 
     # Validate mutually exclusive arguments
-    exclusive_args = [max_instances is not None, max_instances_per_category is not None, category is not None]
-    if sum(exclusive_args) > 1:
-        console.print("[red]Error: --max-instances, --max-questions-per-category, and --category are mutually exclusive[/red]")
+    # --max-instances-per-category can't be combined with --max-instances or --category
+    # But --category CAN be combined with --max-instances (to limit questions within a category)
+    if max_instances_per_category is not None and (max_instances is not None or category is not None):
+        console.print("[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]")
         return
 
     # Validate --only-ingested can't be combined with other dataset filters
@@ -316,12 +539,15 @@ async def run_benchmark(
     failed_question_ids = set()
     invalid_question_ids = set()
     if only_failed or only_invalid:
-        results_path = Path(__file__).parent / 'results' / 'benchmark_results.json'
+        # Use source_results if specified, otherwise default to benchmark_results.json
+        source_file = source_results if source_results else 'benchmark_results.json'
+        results_path = Path(__file__).parent / 'results' / source_file
         if not results_path.exists():
             console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
             console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
             return
 
+        console.print(f"[cyan]Reading failed/invalid questions from: {source_file}[/cyan]")
         with open(results_path, 'r') as f:
             previous_results = json.load(f)
 
@@ -388,9 +614,12 @@ async def run_benchmark(
         else:
             console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate[/green]")
 
-    answer_generator = LongMemEvalAnswerGenerator()
+    answer_generator = LongMemEvalAnswerGenerator(context_format=context_format, category_prompts=category_prompts)
     answer_evaluator = LLMAnswerEvaluator()
 
+    # Log context format being used
+    console.print(f"[blue]Context format: {context_format}[/blue]")
+
     # Create local memory engine
     from benchmarks.common.benchmark_runner import create_memory_engine
     memory = await create_memory_engine()
@@ -481,6 +710,9 @@ async def run_benchmark(
     # Generate detailed report by question type
     generate_type_report(results)
 
+    # Generate markdown results table
+    generate_markdown_table(results, output_path)
+
     return results
 
 
@@ -567,6 +799,67 @@ def generate_type_report(results: dict):
     console.print(table)
 
 
+def generate_markdown_table(results: dict, json_output_path: Path):
+    """Generate a markdown results table with model configuration."""
+    from rich.console import Console
+    console = Console()
+
+    # Aggregate stats by question type
+    type_stats = {}
+
+    for item_result in results['item_results']:
+        metrics = item_result['metrics']
+        by_category = metrics.get('category_stats', {})
+
+        for qtype, stats in by_category.items():
+            if qtype not in type_stats:
+                type_stats[qtype] = {'total': 0, 'correct': 0, 'invalid': 0}
+            type_stats[qtype]['total'] += stats['total']
+            type_stats[qtype]['correct'] += stats['correct']
+            type_stats[qtype]['invalid'] += stats.get('invalid', 0)
+
+    # Build markdown content
+    lines = []
+    lines.append("# LongMemEval Benchmark Results")
+    lines.append("")
+
+    # Add model configuration
+    if 'model_config' in results:
+        config = results['model_config']
+        lines.append("## Model Configuration")
+        lines.append("")
+        lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
+        lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+        lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
+        lines.append("")
+
+    lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
+    lines.append("")
+
+    # Results by question type
+    lines.append("## Results by Question Type")
+    lines.append("")
+    lines.append("| Question Type | Total | Correct | Invalid | Accuracy |")
+    lines.append("|---------------|-------|---------|---------|----------|")
+
+    for qtype in sorted(type_stats.keys()):
+        stats = type_stats[qtype]
+        valid_total = stats['total'] - stats['invalid']
+        acc = (stats['correct'] / valid_total * 100) if valid_total > 0 else 0
+        invalid_str = str(stats['invalid']) if stats['invalid'] > 0 else "-"
+        lines.append(f"| {qtype} | {stats['total']} | {stats['correct']} | {invalid_str} | {acc:.1f}% |")
+
+    # Add overall row
+    total_invalid = results.get('total_invalid', 0)
+    invalid_str = str(total_invalid) if total_invalid > 0 else "-"
+    lines.append(f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |")
+
+    # Write to file (same directory as JSON, but .md extension)
+    md_output_path = json_output_path.with_suffix('.md')
+    md_output_path.write_text('\n'.join(lines))
+    console.print(f"\n[green]✓[/green] Results table saved to {md_output_path}")
+
+
 if __name__ == "__main__":
     import logging
     import argparse
@@ -586,7 +879,7 @@ if __name__ == "__main__":
         type=int,
         default=None,
         dest="max_instances_per_category",
-        help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Mutually exclusive with --max-instances and --category."
+        help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category."
     )
     parser.add_argument(
         "--max-questions",
@@ -641,7 +934,7 @@ if __name__ == "__main__":
         "--category",
         type=str,
         default=None,
-        help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Mutually exclusive with --max-instances and --max-instances-per-category."
+        help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category."
     )
     parser.add_argument(
         "--parallel",
@@ -655,6 +948,25 @@ if __name__ == "__main__":
         default="benchmark_results.json",
         help="Filename for results output (default: benchmark_results.json). Saved in results/ directory."
     )
+    parser.add_argument(
+        "--context-format",
+        type=str,
+        choices=["json", "structured"],
+        default="json",
+        help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json."
+    )
+    parser.add_argument(
+        "--source-results",
+        type=str,
+        default=None,
+        help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified."
+    )
+    parser.add_argument(
+        "--category-prompts",
+        type=str,
+        default=None,
+        help="Path to JSON file containing category-specific prompt templates. Format: {\"category_name\": \"prompt_template_with_{placeholders}\"}"
+    )
 
     args = parser.parse_args()
 
@@ -663,13 +975,19 @@ if __name__ == "__main__":
         parser.error("Cannot use both --only-failed and --only-invalid at the same time")
 
     # Validate mutually exclusive arguments
-    exclusive_count = sum([
-        args.max_instances is not None,
-        args.max_instances_per_category is not None,
-        args.category is not None
-    ])
-    if exclusive_count > 1:
-        parser.error("--max-instances, --max-questions-per-category, and --category are mutually exclusive")
+    # --max-instances-per-category can't be combined with --max-instances or --category
+    if args.max_instances_per_category is not None and (args.max_instances is not None or args.category is not None):
+        parser.error("--max-questions-per-category cannot be combined with --max-instances or --category")
+
+    # Load category prompts if provided
+    category_prompts = None
+    if args.category_prompts:
+        try:
+            with open(args.category_prompts, 'r') as f:
+                category_prompts = json.load(f)
+            print(f"📝 Loaded category prompts for: {', '.join(category_prompts.keys())}")
+        except Exception as e:
+            parser.error(f"Failed to load category prompts from {args.category_prompts}: {str(e)}")
 
     results = asyncio.run(run_benchmark(
         max_instances=args.max_instances,
@@ -685,5 +1003,8 @@ if __name__ == "__main__":
         only_ingested=args.only_ingested,
         category=args.category,
         max_concurrent_items=args.parallel,
-        results_filename=args.results_filename
+        results_filename=args.results_filename,
+        context_format=args.context_format,
+        source_results=args.source_results,
+        category_prompts=category_prompts
     ))

scripts/analyze_benchmark_failures.py

@@ -0,0 +1,834 @@
+#!/usr/bin/env python3
+"""
+Analyze LongMemEval benchmark failures and generate prompt improvements.
+
+This script:
+1. Parses benchmark results to find incorrect questions by category
+2. Uses Groq to generate what correct answers should have been
+3. Analyzes the original benchmark prompt
+4. Suggests prompt improvements to fix the failure cases
+
+Usage:
+    python scripts/analyze_benchmark_failures.py --category single-session-preference
+    python scripts/analyze_benchmark_failures.py --category multi-session --max-examples 5
+"""
+
+import json
+import argparse
+import asyncio
+import os
+import time
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass
+import sys
+from dotenv import load_dotenv
+import pydantic
+
+# Load environment variables from .env file
+script_dir = Path(__file__).parent.parent
+env_file = script_dir / ".env"
+if env_file.exists():
+    load_dotenv(env_file)
+
+# Add the hindsight-api directory to Python path
+sys.path.append(str(script_dir / "hindsight-api"))
+
+from hindsight_api.engine.llm_wrapper import LLMConfig
+
+
+async def make_api_call_with_retry(client, model, messages, max_tokens=1000, temperature=0.1, seed=4242, response_format=None, max_retries=3):
+    """Make API call with retry logic for rate limits."""
+    for attempt in range(max_retries):
+        try:
+            kwargs = {
+                "model": model,
+                "messages": messages,
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "seed": seed
+            }
+            if response_format:
+                kwargs["response_format"] = response_format
+
+            response = await client.chat.completions.create(**kwargs)
+            return response
+        except Exception as e:
+            if "rate_limit_exceeded" in str(e) and attempt < max_retries - 1:
+                # Extract wait time from error message or use exponential backoff
+                if "Please try again in" in str(e):
+                    import re
+                    match = re.search(r'Please try again in (\d+\.?\d*)s', str(e))
+                    wait_time = float(match.group(1)) if match else 2 ** attempt
+                else:
+                    wait_time = 2 ** attempt
+
+                print(f"    Rate limit hit, waiting {wait_time:.1f}s before retry {attempt + 1}/{max_retries}...")
+                await asyncio.sleep(wait_time)
+            else:
+                raise e
+    raise Exception("Max retries exceeded")
+
+
+class JudgeResponse(pydantic.BaseModel):
+    """Judge response format."""
+    correct: bool
+    reasoning: str
+
+
+@dataclass
+class FailureCase:
+    """Represents a single failure case from the benchmark."""
+    question: str
+    correct_answer: str
+    predicted_answer: str
+    retrieved_memories: List[Dict[str, Any]]
+    chunks: Dict[str, Any]
+    correctness_reasoning: str
+    category: str
+
+
+class BenchmarkAnalyzer:
+    """Analyzes benchmark failures and generates prompt improvements."""
+
+    def __init__(self):
+        """Initialize the analyzer with Groq LLM configuration."""
+        self.llm_config = LLMConfig.for_answer_generation()
+        self.client = self.llm_config._client  # Use raw client to bypass LLMConfig issues
+
+        # Initialize judge LLM config
+        self.judge_config = LLMConfig.for_judge()
+        self.judge_client = self.judge_config._client
+
+    def load_benchmark_results(self, results_path: Path) -> Dict[str, Any]:
+        """Load benchmark results from JSON file."""
+        with open(results_path, 'r') as f:
+            return json.load(f)
+
+    def extract_failures_by_category(self, results: Dict[str, Any], category: str) -> List[FailureCase]:
+        """Extract all failure cases for a specific category."""
+        failures = []
+
+        for item_result in results.get('item_results', []):
+            for detail in item_result['metrics'].get('detailed_results', []):
+                # Check if this is an incorrect answer in the target category
+                if (detail.get('category') == category and
+                    detail.get('is_correct') == False and
+                    not detail.get('is_invalid', False)):
+
+                    failure = FailureCase(
+                        question=detail['question'],
+                        correct_answer=detail['correct_answer'],
+                        predicted_answer=detail['predicted_answer'],
+                        retrieved_memories=detail.get('retrieved_memories', []),
+                        chunks=detail.get('chunks', {}),
+                        correctness_reasoning=detail.get('correctness_reasoning', ''),
+                        category=category
+                    )
+                    failures.append(failure)
+
+        return failures
+
+    async def judge_answer(self, question: str, correct_answer: str, predicted_answer: str, category: str) -> Tuple[bool, str]:
+        """Use the same judge logic as the benchmark to validate an answer."""
+
+        # Use the exact same category-specific prompts as the benchmark
+        if category in ['single-session-user', 'single-session-assistant', 'multi-session']:
+            prompt_content = f"""Evaluate if the model response contains the correct answer to the question.
+
+I will give you a question, a correct answer, and a response from a model.
+Please set correct=true if the response contains the correct answer. Otherwise, set correct=no.
+If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also set correct=true.
+If the response only contains a subset of the information required by the answer, set correct=false
+
+Question: {question}
+
+Correct Answer: {correct_answer}
+
+Model Response: {predicted_answer}
+
+Evaluation criteria:
+- Set correct=true if the response contains the correct answer
+- Set correct=true if the response is equivalent to the correct answer or contains intermediate steps
+- Set correct=false if the response is incorrect or missing key information
+
+Provide your evaluation as JSON with:
+- reasoning: One sentence explanation
+- correct: true or false"""
+
+        elif category == 'temporal-reasoning':
+            prompt_content = f"""
+I will give you a question, a correct answer, and a response from a model.
+Please set correct=true if the response contains the correct answer. Otherwise, set correct=false.
+If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also set correct=true.
+If the response only contains a subset of the information required by the answer, answer correct=false.
+In addition, do not penalize off-by-one errors for the number of days. If the question asks for the number of days/weeks/months, etc., and the model makes off-by-one errors (e.g., predicting 19 days when the answer is 18), the model's response is still correct.
+
+Question: {question}
+Gold answer: {correct_answer}
+Generated answer: {predicted_answer}
+First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred.
+If it's correct, set correct=true."""
+
+        elif category == 'knowledge-update':
+            prompt_content = f"""
+I will give you a question, a correct answer, and a response from a model.
+Please set correct=true if the response contains the correct answer. Otherwise, set correct=false.
+If the response contains some previous information along with an updated answer, the response should be considered as correct as long as the updated answer is the required answer.
+
+Question: {question}
+Gold answer: {correct_answer}
+Generated answer: {predicted_answer}
+First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred.
+If it's correct, set correct=true."""
+
+        elif category == 'single-session-preference':
+            prompt_content = f"""
+I will give you a question, a answer for desired personalized response, and a response from a model.
+Please set correct=true if the response satisfies the desired response. Otherwise, set correct=false.
+The model does not need to reflect all the points in the desired response. The response is correct as long as it recalls and utilizes the user's personal information correctly.
+
+Question: {question}
+Gold answer: {correct_answer}
+Generated answer: {predicted_answer}
+First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred.
+If it's correct, set correct=true."""
+
+        else:
+            # Default LoComo-style evaluation
+            prompt_content = f"""Your task is to label an answer to a question as 'CORRECT' or 'WRONG'. You will be given the following data:
+(1) a question (posed by one user to another user),
+(2) a 'gold' (ground truth) answer,
+(3) a generated answer
+which you will score as CORRECT/WRONG.
+
+The point of the question is to ask about something one user should know about the other user based on their prior conversations.
+The gold answer will usually be a concise and short answer that includes the referenced topic, for example:
+Question: Do you remember what I got the last time I went to Hawaii?
+Gold answer: A shell necklace
+The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
+
+For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date.
+There's an edge case where the actual answer can't be found in the data and in that case the gold answer will say so (e.g. 'You did not mention this information.'); if the generated answer says that it cannot be answered or it doesn't know all the details, it should be counted as CORRECT.
+
+Question: {question}
+Gold answer: {correct_answer}
+Generated answer: {predicted_answer}
+First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred.
+If it's correct, set correct=true."""
+
+        try:
+            # Use retry function for judge call
+            messages = [{
+                "role": "user",
+                "content": f"{prompt_content}\n\nRespond in JSON format with 'correct' (boolean) and 'reasoning' (string) fields."
+            }]
+
+            response = await make_api_call_with_retry(
+                client=self.judge_client,
+                model=self.judge_config.model,
+                messages=messages,
+                max_tokens=4096,
+                temperature=0,
+                seed=4242,
+                response_format={"type": "json_object"}
+            )
+
+            content = response.choices[0].message.content
+            result = json.loads(content)
+
+            return result.get('correct', False), result.get('reasoning', 'No reasoning provided')
+
+        except Exception as e:
+            print(f"    Judge error: {str(e)}")
+            return False, f"Judge error: {str(e)}"
+
+    async def generate_correct_answer(self, failure: FailureCase, max_attempts: int = 10) -> str:
+        """Use Groq to generate what the correct answer should have been, with judge validation."""
+
+        # Recreate the exact same recall_result structure that the benchmark uses
+        recall_result = {
+            "results": failure.retrieved_memories,
+            "chunks": failure.chunks,
+            "entities": {}  # We don't have entities in the saved results, but benchmark expects this key
+        }
+
+        # Use the exact same context formatting as the benchmark (JSON format)
+        context = json.dumps(recall_result)
+
+        for attempt in range(max_attempts):
+            print(f"    Attempt {attempt + 1}/{max_attempts}...")
+
+            # Use the exact same prompt template as the benchmark
+            prompt_template = self.get_original_prompt_template()
+
+            # Format the prompt exactly like the benchmark does
+            prompt_content = prompt_template.format(
+                context_instructions="",  # Empty for default benchmark
+                question=failure.question,
+                formatted_question_date="Not specified",  # We don't have question dates in our failure data
+                context=context
+            )
+
+            try:
+                # Generate candidate answer using exact benchmark prompt with retry
+                response = await make_api_call_with_retry(
+                    client=self.client,
+                    model=self.llm_config.model,
+                    messages=[{"role": "user", "content": prompt_content}],
+                    max_tokens=1000,  # Allow much longer answers
+                    temperature=0.1 + (attempt * 0.05),  # Increase temperature slightly with each attempt
+                    seed=4242 + attempt  # Different seed for each attempt
+                )
+
+                candidate_answer = response.choices[0].message.content
+                if not candidate_answer:
+                    print(f"      Empty response, trying again...")
+                    continue
+
+                candidate_answer = candidate_answer.strip()
+                print(f"      Generated: {candidate_answer}")
+
+                # Validate with judge
+                is_correct, judge_reasoning = await self.judge_answer(
+                    failure.question,
+                    failure.correct_answer,
+                    candidate_answer,
+                    failure.category
+                )
+
+                print(f"      Judge says: {'✓ CORRECT' if is_correct else '✗ INCORRECT'} - {judge_reasoning}")
+
+                if is_correct:
+                    return candidate_answer
+
+            except Exception as e:
+                print(f"      Error in attempt {attempt + 1}: {str(e)}")
+
+        # If we get here, we failed all attempts
+        return f"Failed to generate judge-validated answer after {max_attempts} attempts"
+
+    def get_original_prompt_template(self) -> str:
+        """Extract the original prompt template from the longmemeval benchmark code."""
+        # This is the prompt template from longmemeval_benchmark.py
+        return """You are a helpful assistant that must answer user questions based on the previous conversations.
+
+{context_instructions}**Answer Guidelines:**
+1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
+2. Reason about all the memories and find the right answer, considering the most recent memory as an update of the current facts.
+3. If you have 2 possible answers, just say both.
+
+In general the answer must be comprehensive and plenty of details from the retrieved context.
+
+For quantitative/counting questions ("how many..."): First list each unique item in your reasoning (1. X, 2. Y, 3. Z...), scanning ALL facts, then count them for your answer.
+If questions asks a location (where...?) make sure to include the location name.
+For recommendation questions ("can you recommend...", "suggest...", "any tips..."): DO NOT give actual recommendations. Instead, describe what KIND the user would prefer based on their context. Example answer format: "The user would prefer recommendations for [category] that focus on [their interest]. They would not prefer [what to avoid based on context]."
+For questions asking for help or instructions, consider the users' recent memories and previous interactions with the assistant to understand their current situation better (recent purchases, specific product models used..)
+For specific number/value questions, use the context to understand what is the most up-to-date number based on recency, but also include the reasoning (in the answer) on previous possible values and why you think are less relevant.
+For open questions, include as much details as possible from different sources that are relevant.
+For questions where a specific entity/role is mentioned and it's different from your memory, just say the truth, don't make up anything just to fulfill the question. For example, if the question is about a specific sport, you should consider if the memories and the question are about the same sport. (e.g. american football vs soccer, shows vs podcasts)
+For comparative questions, say you don't know the answer if you don't have information about both sides. (or more sides)
+For questions related to time/date, carefully review the question date and the memories date to correctly answer the question.
+For questions related to time/date calculation (e.g. How many days passed between X and Y?), carefully review the memories date to correctly answer the question and only provide an answer if you have information about both X and Y, otherwise say it's not possible to calculate and why.
+
+Consider assistant's previous actions (e.g., bookings, reminders) as impactful to the user experiences.
+
+Question: {question}
+Question Date: {formatted_question_date}
+
+Retrieved Context:
+{context}
+
+Answer:"""
+
+    async def suggest_prompt_improvements(self, failures: List[FailureCase], category: str, correct_answers: List[str]) -> str:
+        """Use LLM to suggest improvements to the benchmark prompt."""
+
+        original_prompt = self.get_original_prompt_template()
+
+        # Create examples of failures
+        failure_examples = ""
+        for i, (failure, correct_answer) in enumerate(zip(failures, correct_answers), 1):
+            # Create summary of both memories and chunks for brevity
+            context_summary = []
+
+            # Add memory summaries
+            for j, memory in enumerate(failure.retrieved_memories[:3], 1):  # Limit to first 3 memories for brevity
+                context_summary.append(f"  Memory {j}: {memory.get('text', '')[:100]}...")
+
+            # Add chunk summaries
+            chunk_count = len(failure.chunks)
+            if chunk_count > 0:
+                context_summary.append(f"  + {chunk_count} conversation chunks with raw dialogue")
+
+            failure_examples += f"""
+Example {i}:
+Question: {failure.question}
+Expected Answer: {failure.correct_answer}
+Model's Incorrect Answer: {failure.predicted_answer}
+Why Incorrect: {failure.correctness_reasoning}
+Context Available to Model:
+{chr(10).join(context_summary)}
+What Should Have Been Generated: {correct_answer}
+
+---
+"""
+
+        prompt = f"""You are an expert at designing prompts for language models that need to work across diverse question types and topics in memory-based QA benchmarks.
+
+CONTEXT:
+I'm working on improving a benchmark prompt for the LongMemEval dataset. The current prompt is used to generate answers for questions across many different categories and topics. The category I'm focusing on is "{category}" but the prompt needs to work well for ALL categories in the benchmark.
+
+CURRENT BENCHMARK PROMPT:
+{original_prompt}
+
+FAILURE ANALYSIS:
+The current prompt is producing incorrect answers for questions in the "{category}" category. Here are the specific failure cases:
+
+{failure_examples}
+
+TASK:
+Analyze these failure patterns and suggest improvements to the benchmark prompt. Your improvements should:
+
+1. **Be general**: Work for the "{category}" category AND other categories in the benchmark
+2. **Address root causes**: Fix the underlying issues causing these specific failures
+3. **Maintain compatibility**: Don't break existing good performance on other question types
+4. **Be specific**: Provide concrete guidance that helps the model generate more accurate answers
+
+Focus on:
+- What specific guidance is missing that would have prevented these errors?
+- How can the prompt better handle the nuances of "{category}" questions?
+- What instructions would help the model be more precise in its reasoning?
+
+Please provide:
+1. A brief analysis of the failure patterns you observe
+2. Specific recommended changes to the prompt
+3. The improved prompt with your changes highlighted
+
+IMPROVED PROMPT:"""
+
+        try:
+            # Use retry function for prompt improvement generation
+            response = await make_api_call_with_retry(
+                client=self.client,
+                model=self.llm_config.model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=2000,
+                temperature=0.2,
+                seed=4242  # Same seed as used by LLMConfig
+            )
+
+            content = response.choices[0].message.content
+            return content.strip() if content else "No improvements generated"
+        except Exception as e:
+            return f"Error generating prompt improvements: {str(e)}"
+
+    async def test_prompt_performance(self, failures: List[FailureCase], prompt_template: str) -> float:
+        """Test a prompt template against the failure cases and return success rate."""
+        correct_count = 0
+        total_count = len(failures)
+
+        print(f"  Testing prompt performance on {total_count} failure cases...")
+
+        for i, failure in enumerate(failures, 1):
+            print(f"    Testing case {i}/{total_count}...", end="")
+
+            try:
+                # Recreate the exact same recall_result structure that the benchmark uses
+                recall_result = {
+                    "results": failure.retrieved_memories,
+                    "chunks": failure.chunks,
+                    "entities": {}
+                }
+
+                # Use the exact same context formatting as the benchmark (JSON format)
+                context = json.dumps(recall_result)
+
+                # Format the prompt exactly like the benchmark does
+                prompt_content = prompt_template.format(
+                    context_instructions="",  # Empty for default benchmark
+                    question=failure.question,
+                    formatted_question_date="Not specified",
+                    context=context
+                )
+
+                # Generate answer using the test prompt with retry
+                response = await make_api_call_with_retry(
+                    client=self.client,
+                    model=self.llm_config.model,
+                    messages=[{"role": "user", "content": prompt_content}],
+                    max_tokens=1000,
+                    temperature=0.1,
+                    seed=4242
+                )
+
+                candidate_answer = response.choices[0].message.content
+                if not candidate_answer:
+                    print(" SKIP (empty response)")
+                    continue
+
+                candidate_answer = candidate_answer.strip()
+
+                # Validate with judge
+                is_correct, judge_reasoning = await self.judge_answer(
+                    failure.question,
+                    failure.correct_answer,
+                    candidate_answer,
+                    failure.category
+                )
+
+                if is_correct:
+                    correct_count += 1
+                    print(" ✓")
+                else:
+                    print(" ✗")
+
+            except Exception as e:
+                print(f" ERROR: {str(e)}")
+                continue
+
+        success_rate = correct_count / total_count if total_count > 0 else 0.0
+        print(f"  Prompt performance: {correct_count}/{total_count} = {success_rate:.2%}")
+        return success_rate
+
+
+async def main():
+    """Main function to run the analysis."""
+    parser = argparse.ArgumentParser(description="Analyze benchmark failures and suggest prompt improvements")
+    parser.add_argument(
+        "--category",
+        type=str,
+        required=True,
+        help="Category to analyze (e.g., 'single-session-preference', 'multi-session')"
+    )
+    parser.add_argument(
+        "--results-file",
+        type=str,
+        default="hindsight-dev/benchmarks/longmemeval/results/benchmark_results.json",
+        help="Path to benchmark results JSON file"
+    )
+    parser.add_argument(
+        "--max-examples",
+        type=int,
+        default=10,
+        help="Maximum number of failure examples to analyze"
+    )
+    parser.add_argument(
+        "--optimize-prompt",
+        action="store_true",
+        help="Run iterative prompt optimization (25 rounds) to find best performing prompt"
+    )
+
+    args = parser.parse_args()
+
+    # Convert to absolute path
+    script_dir = Path(__file__).parent.parent
+    results_path = script_dir / args.results_file
+
+    if not results_path.exists():
+        print(f"Error: Results file not found at {results_path}")
+        return
+
+    print(f"🔍 Analyzing failures for category: {args.category}")
+    print(f"📄 Reading results from: {results_path}")
+
+    analyzer = BenchmarkAnalyzer()
+
+    # Load and analyze results
+    results = analyzer.load_benchmark_results(results_path)
+    failures = analyzer.extract_failures_by_category(results, args.category)
+
+    if not failures:
+        print(f"✅ No failures found for category '{args.category}'")
+        return
+
+    # Limit examples if requested (but not during optimization)
+    total_failures = len(failures)
+    if not args.optimize_prompt and len(failures) > args.max_examples:
+        failures = failures[:args.max_examples]
+        print(f"📊 Found {total_failures} failures (limited to {args.max_examples} for analysis)")
+    else:
+        print(f"📊 Found {len(failures)} failures")
+
+    print("\n🤖 Generating judge-validated correct answers for failure cases...")
+    correct_answers = []
+    validation_stats = {"successful": 0, "failed": 0, "errors": 0}
+
+    for i, failure in enumerate(failures, 1):
+        print(f"  Processing example {i}/{len(failures)} (Question: {failure.question[:50]}...)")
+        try:
+            correct_answer = await analyzer.generate_correct_answer(failure)
+
+            if correct_answer.startswith("Failed to generate"):
+                print(f"    ❌ {correct_answer}")
+                validation_stats["failed"] += 1
+            else:
+                print(f"    ✅ Judge-validated answer generated successfully")
+                validation_stats["successful"] += 1
+
+            correct_answers.append(correct_answer)
+        except Exception as e:
+            print(f"    💥 Error generating correct answer: {e}")
+            correct_answers.append(f"Error: {str(e)}")
+            validation_stats["errors"] += 1
+
+    # Report validation statistics
+    print(f"\n📈 Judge validation results:")
+    print(f"  ✅ Successfully validated: {validation_stats['successful']}")
+    print(f"  ❌ Failed to validate: {validation_stats['failed']}")
+    print(f"  💥 Errors: {validation_stats['errors']}")
+
+    print("\n💡 Generating prompt improvements...")
+    try:
+        improvements = await analyzer.suggest_prompt_improvements(failures, args.category, correct_answers)
+        print(f"  Generated improvements: {len(improvements)} characters")
+    except Exception as e:
+        print(f"  Error generating improvements: {e}")
+        improvements = f"Error generating improvements: {str(e)}"
+
+    # Run iterative prompt optimization if requested
+    best_prompt = None
+    best_performance = 0.0
+    optimization_history = []
+
+    if args.optimize_prompt:
+        print(f"\n🔄 Starting iterative prompt optimization (25 rounds)...")
+
+        # Start with the original benchmark prompt
+        current_prompt = analyzer.get_original_prompt_template()
+        baseline_performance = await analyzer.test_prompt_performance(failures, current_prompt)
+
+        print(f"\n📊 Baseline performance: {baseline_performance:.2%}")
+        best_prompt = current_prompt
+        best_performance = baseline_performance
+        optimization_history.append({
+            "round": 0,
+            "performance": baseline_performance,
+            "prompt": "Original benchmark prompt",
+            "is_best": True
+        })
+
+        for round_num in range(1, 26):
+            print(f"\n🔄 Round {round_num}/25: Analyzing current failures and generating improved prompt...")
+
+            try:
+                # Test current best prompt and collect new failures
+                current_failures = []
+                current_correct_answers = []
+
+                print(f"  📊 Testing current prompt to identify remaining failures...")
+                for i, failure in enumerate(failures, 1):
+                    # Recreate the exact same recall_result structure
+                    recall_result = {
+                        "results": failure.retrieved_memories,
+                        "chunks": failure.chunks,
+                        "entities": {}
+                    }
+                    context = json.dumps(recall_result)
+
+                    # Format the prompt exactly like the benchmark does
+                    prompt_content = current_prompt.format(
+                        context_instructions="",
+                        question=failure.question,
+                        formatted_question_date="Not specified",
+                        context=context
+                    )
+
+                    # Generate answer using current prompt with retry
+                    response = await make_api_call_with_retry(
+                        client=analyzer.client,
+                        model=analyzer.llm_config.model,
+                        messages=[{"role": "user", "content": prompt_content}],
+                        max_tokens=1000,
+                        temperature=0.1,
+                        seed=4242
+                    )
+
+                    candidate_answer = response.choices[0].message.content
+                    if not candidate_answer:
+                        # Still a failure
+                        current_failures.append(failure)
+                        current_correct_answers.append(failure.correct_answer)
+                        continue
+
+                    candidate_answer = candidate_answer.strip()
+
+                    # Validate with judge
+                    is_correct, judge_reasoning = await analyzer.judge_answer(
+                        failure.question,
+                        failure.correct_answer,
+                        candidate_answer,
+                        failure.category
+                    )
+
+                    if not is_correct:
+                        # This is still a failure - add it to current failures for analysis
+                        current_failures.append(failure)
+
+                        # Generate what the correct answer should have been for this failure
+                        correct_answer = await analyzer.generate_correct_answer(failure, max_attempts=3)
+                        current_correct_answers.append(correct_answer)
+
+                print(f"  🔍 Found {len(current_failures)} remaining failures to analyze")
+
+                if len(current_failures) == 0:
+                    print(f"  🎉 Perfect! No more failures - optimization complete!")
+                    break
+
+                # Generate a new improved prompt based on CURRENT failures
+                improved_prompt = await analyzer.suggest_prompt_improvements(current_failures, args.category, current_correct_answers)
+
+                # Extract the actual prompt from the improvements text
+                if "IMPROVED PROMPT:" in improved_prompt:
+                    prompt_start = improved_prompt.find("IMPROVED PROMPT:") + len("IMPROVED PROMPT:")
+                    new_prompt = improved_prompt[prompt_start:].strip()
+                elif "You are a helpful assistant" in improved_prompt:
+                    prompt_start = improved_prompt.find("You are a helpful assistant")
+                    new_prompt = improved_prompt[prompt_start:].strip()
+                else:
+                    print(f"  ❌ Could not extract prompt from improvements, using current prompt")
+                    new_prompt = current_prompt
+
+                # Test the new prompt on ALL original failures
+                performance = await analyzer.test_prompt_performance(failures, new_prompt)
+
+                is_better = performance > best_performance
+                optimization_history.append({
+                    "round": round_num,
+                    "performance": performance,
+                    "prompt": "Generated improved prompt",
+                    "is_best": is_better,
+                    "failures_analyzed": len(current_failures)
+                })
+
+                if is_better:
+                    print(f"  🎉 New best performance: {performance:.2%} (improved by {(performance - best_performance):.1%})")
+                    best_prompt = new_prompt
+                    best_performance = performance
+                    current_prompt = new_prompt
+                else:
+                    print(f"  📉 Performance: {performance:.2%} (no improvement)")
+                    # Don't update current_prompt - keep using the best one
+
+            except Exception as e:
+                print(f"  💥 Error in round {round_num}: {str(e)}")
+                optimization_history.append({
+                    "round": round_num,
+                    "performance": 0.0,
+                    "prompt": f"Error: {str(e)}",
+                    "is_best": False,
+                    "failures_analyzed": 0
+                })
+
+        print(f"\n🏆 Optimization complete!")
+        print(f"  📈 Best performance achieved: {best_performance:.2%}")
+        print(f"  📊 Improvement over baseline: {(best_performance - baseline_performance):.1%}")
+
+    # Save results
+    output_dir = script_dir / "results" / "prompt_analysis"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    output_file = output_dir / f"{args.category}_analysis.md"
+
+    # Create detailed report
+    report = f"""# Benchmark Failure Analysis: {args.category}
+
+## Summary
+- **Category**: {args.category}
+- **Total Failures Analyzed**: {len(failures)}
+- **Judge-Validated Answers Generated**: {validation_stats['successful']}/{len(failures)}
+- **Results File**: {results_path}
+
+## Judge Validation Process
+This analysis uses a new approach where generated "correct" answers are validated using the exact same judge logic as the benchmark. We attempt up to 10 times to generate an answer that the judge accepts as correct, ensuring that our prompt improvements are based on truly correct examples.
+
+- ✅ **Successfully validated**: {validation_stats['successful']} answers
+- ❌ **Failed to validate**: {validation_stats['failed']} answers
+- 💥 **Errors**: {validation_stats['errors']} answers
+
+## Failure Cases
+
+"""
+
+    for i, (failure, correct_answer) in enumerate(zip(failures, correct_answers), 1):
+        report += f"""### Example {i}
+
+**Question**: {failure.question}
+
+**Expected Answer**: {failure.correct_answer}
+
+**Model's Answer**: {failure.predicted_answer}
+
+**Why Incorrect**: {failure.correctness_reasoning}
+
+**Generated Correct Answer**: {correct_answer}
+
+**Retrieved Memories**:
+"""
+        for j, memory in enumerate(failure.retrieved_memories, 1):
+            memory_text = memory.get('text', '')
+            # Truncate memory text for markdown output only
+            truncated_text = memory_text[:200] + "..." if len(memory_text) > 200 else memory_text
+            report += f"- **Memory {j}**: {truncated_text}\n"
+
+        # Add chunks information
+        report += f"\n**Conversation Chunks**: {len(failure.chunks)} chunks containing raw dialogue\n"
+        if failure.chunks:
+            report += "- Chunk IDs: " + ", ".join(list(failure.chunks.keys())[:5])  # Show first 5 chunk IDs
+            if len(failure.chunks) > 5:
+                report += f" ... and {len(failure.chunks) - 5} more"
+            report += "\n"
+
+        report += "\n---\n\n"
+
+    report += f"""## Prompt Improvement Analysis
+
+{improvements}
+"""
+
+    # Add optimization results if optimization was run
+    if args.optimize_prompt and optimization_history:
+        report += f"""
+## Iterative Prompt Optimization Results
+
+### Performance Summary
+- **Baseline Performance**: {optimization_history[0]['performance']:.2%}
+- **Best Performance Achieved**: {best_performance:.2%}
+- **Improvement**: {(best_performance - optimization_history[0]['performance']):.1%}
+- **Total Optimization Rounds**: 25
+
+### Optimization History
+| Round | Performance | Improvement | Failures Analyzed | Best So Far |
+|-------|-------------|-------------|-------------------|-------------|
+"""
+        for entry in optimization_history:
+            is_best_marker = "🏆" if entry['is_best'] else ""
+            improvement = f"+{(entry['performance'] - optimization_history[0]['performance']):.1%}" if entry['performance'] > optimization_history[0]['performance'] else f"{(entry['performance'] - optimization_history[0]['performance']):.1%}"
+            failures_count = entry.get('failures_analyzed', 'N/A')
+            report += f"| {entry['round']} | {entry['performance']:.2%} | {improvement} | {failures_count} | {is_best_marker} |\n"
+
+        report += f"""
+### Best Performing Prompt
+```
+{best_prompt}
+```
+"""
+
+    with open(output_file, 'w') as f:
+        f.write(report)
+
+    print(f"\n✅ Analysis complete!")
+    print(f"📄 Report saved to: {output_file}")
+    print(f"\n🔍 Key findings:")
+    print(f"  - {len(failures)} failure cases analyzed")
+    print(f"  - {validation_stats['successful']}/{len(failures)} judge-validated correct answers generated")
+    print(f"  - Prompt improvements generated using validated answers")
+    if args.optimize_prompt and optimization_history:
+        print(f"  - Iterative optimization completed: {best_performance:.2%} best performance")
+        improvement_pct = (best_performance - optimization_history[0]['performance']) * 100
+        print(f"  - Performance improvement: {improvement_pct:+.1f} percentage points")
+    print(f"  - Detailed report available at {output_file}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())