Increase graph neighbor limit and benchmark improvements (#18)

* Improve LongMemEval benchmark with structured prompts and better options

- Add --context-format option with 'json' (original) and 'structured' modes
- Structured format groups facts with source chunks for better LLM comprehension
- Add detailed instructions for date calculations, relative time handling, and abstention
- Add --source-results flag to read failed questions from a different file
- Allow --category to be combined with --max-instances for sampling
- Fix Gemini structured output by passing response_schema parameter
- Add retry logic for empty Gemini responses with block reason logging
- Add judge prompt comparison documentation

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

Co-Authored-By: Claude <noreply@anthropic.com>

* fix recall in benchmarks

* Improve LongMemEval prompt and Gemini error handling

- Add JSONDecodeError retry for Gemini truncated responses
- Increase max_tokens to 32768 for thinking models
- Add counting/disambiguation guidance to structured prompt
- Add "when in doubt, undercount" and overlap detection rules

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

Co-Authored-By: Claude <noreply@anthropic.com>

* Add connection error retry and preference question guidance

- 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 <noreply@anthropic.com>

* Make reasoning optional

* Seed for LLM through Groq

* fix entity and observations

* Increase graph retrieval neighbor limit for expanded entities

Doubled the neighbor limit multiplier from 10 to 20 in graph retrieval.
With expanded entity extraction (now including objects and concepts like
"kitchen"), facts share more common entities, causing the previous limit
to arbitrarily exclude relevant results. This fix ensures better recall
for questions about related items (e.g., kitchen items).

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

Co-Authored-By: Claude <noreply@anthropic.com>

* Expand entity extraction to include objects and concepts

Updated entity extraction prompt to include:
- Specific objects (coffee maker, toaster, car, laptop, kitchen)
- Abstract concepts/themes (friendship, career growth, loss, celebration)
- Places and organizations (IKEA, Goodwill, New York)

This enables better fact linking through shared entities. For example,
kitchen appliances now share a "kitchen" entity, allowing graph traversal
to find related facts like "replaced coffee maker" when querying about
"kitchen items".

Works in conjunction with the increased neighbor limit to improve recall.

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

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Chris Bartholomew <chris.bartholomew@vectorize.io>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: andrew <andrew.neeser@me.com>
This commit is contained in:
Nicolò Boschi 2025-12-08 15:24:13 +01:00 committed by GitHub
parent cf2f739469
commit 3bb0a58ded
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 643 additions and 64 deletions

View file

@ -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"
)

View file

@ -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(

View file

@ -167,10 +167,10 @@ class ExtractedFact(BaseModel):
description="'world' = about the user/others (background, experiences). 'assistant' = experience with the assistant."
)
# Entities - extracted from 'who' field
# Entities - extracted from fact content
entities: Optional[List[Entity]] = Field(
default=None,
description="Named entities from 'who': people names, organizations, places. NOT generic relations."
description="Named entities, objects, AND abstract concepts from the fact. Include: people names, organizations, places, significant objects (e.g., 'coffee maker', 'car'), AND abstract concepts/themes (e.g., 'friendship', 'career growth', 'loss', 'celebration'). Extract anything that could help link related facts together."
)
causal_relations: Optional[List[CausalRelation]] = Field(
default=None,
@ -415,14 +415,26 @@ Example: "I love Italian food and prefer outdoor dining"
Fact 2: what="User prefers outdoor dining", who="user", why="This is a dining preference", entities=["user"]
ENTITIES - INCLUDE "user" (CRITICAL)
ENTITIES - INCLUDE PEOPLE, PLACES, OBJECTS, AND CONCEPTS (CRITICAL)
When a fact is ABOUT the user (their preferences, plans, experiences), ALWAYS include "user" in entities!
Extract entities that help link related facts together. Include:
1. "user" - when the fact is about the user
2. People names - Emily, Dr. Smith, etc.
3. Organizations/Places - IKEA, Goodwill, New York, etc.
4. Specific objects - coffee maker, toaster, car, laptop, kitchen, etc.
5. Abstract concepts - themes, values, emotions, or ideas that capture the essence of the fact:
- "friendship" for facts about friends helping each other, bonding, loyalty
- "career growth" for facts about promotions, learning new skills, job changes
- "loss" or "grief" for facts about death, endings, saying goodbye
- "celebration" for facts about parties, achievements, milestones
- "trust" or "betrayal" for facts involving those themes
CORRECT: entities=["user"] for "User loves coffee"
CORRECT: entities=["user", "Emily"] for "User attended Emily's wedding"
WRONG: entities=[] for facts about the user
CORRECT: entities=["user", "coffee maker", "Goodwill", "kitchen"] for "User donated their coffee maker to Goodwill"
CORRECT: entities=["user", "Emily", "friendship"] for "Emily helped user move to a new apartment"
CORRECT: entities=["user", "promotion", "career growth"] for "User got promoted to senior engineer"
CORRECT: entities=["user", "grandmother", "loss", "grief"] for "User's grandmother passed away last week"
WRONG: entities=["user", "Emily"] only - missing the "friendship" concept that links to other friendship facts!
EXAMPLES
@ -438,14 +450,14 @@ Output facts:
- who: "user"
- why: "User prefers intimate outdoor settings"
- fact_type: "world", fact_kind: "conversation"
- entities: ["user"]
- entities: ["user", "wedding", "outdoor ceremony"]
2. User planning wedding
- what: "User is planning their own wedding"
- who: "user"
- why: "Inspired by Emily's ceremony"
- fact_type: "world", fact_kind: "conversation"
- entities: ["user"]
- entities: ["user", "wedding"]
3. Emily's wedding (THE EVENT)
- what: "Emily got married to Sarah at a rooftop garden ceremony in the city"
@ -453,7 +465,7 @@ Output facts:
- why: "User found it romantic and beautiful"
- fact_type: "world", fact_kind: "event"
- occurred_start: "2024-06-09T00:00:00Z" (recently, user "just got back")
- entities: ["user", "Emily", "Sarah"]
- entities: ["user", "Emily", "Sarah", "wedding", "rooftop garden"]
Example 2 - Assistant Facts (Context: March 5, 2024):
Input: "User: My API is really slow when we have 1000+ concurrent users. What can I do?
@ -465,7 +477,22 @@ Output fact:
- who: "user, assistant"
- why: "User asked how to fix slow API performance with 1000+ concurrent users, expected 70-80% reduction in database load"
- fact_type: "assistant", fact_kind: "conversation"
- entities: ["user"]
- entities: ["user", "API", "Redis"]
Example 3 - Kitchen Items with Concept Inference (Context: May 30, 2024):
Input: "I finally donated my old coffee maker to Goodwill. I upgraded to that new espresso machine last month and the old one was just taking up counter space."
Output fact:
- what: "User donated their old coffee maker to Goodwill after upgrading to a new espresso machine"
- when: "May 30, 2024"
- who: "user"
- why: "The old coffee maker was taking up counter space after the upgrade"
- fact_type: "world", fact_kind: "event"
- occurred_start: "2024-05-30T00:00:00Z"
- entities: ["user", "coffee maker", "Goodwill", "espresso machine", "kitchen"]
Note: "kitchen" is inferred as a concept because coffee makers and espresso machines are kitchen appliances.
This links the fact to other kitchen-related facts (toaster, faucet, kitchen mat, etc.) via the shared "kitchen" entity.
Note how the "why" field captures the FULL STORY: what the user asked AND what outcome was expected!

View file

@ -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")

View file

@ -170,9 +170,9 @@ async def retrieve_graph(
batch_activations[unit_id] = activation
# Batch fetch neighbors for all nodes in this batch
# Fetch top weighted neighbors (batch_size * 10 = ~200 for good distribution)
# Fetch top weighted neighbors (batch_size * 20 = ~400 for good distribution)
if batch_nodes and budget_remaining > 0:
max_neighbors = len(batch_nodes) * 10
max_neighbors = len(batch_nodes) * 20
neighbors = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end, mu.mentioned_at,

View file

@ -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);

View file

@ -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);

View file

@ -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>

View file

@ -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()
@ -780,6 +839,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 +931,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 +1193,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 +1316,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
}

View file

@ -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 |")

View file

@ -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. |

View file

@ -128,16 +128,202 @@ 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"):
"""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
"""
self.llm_config = LLMConfig.for_answer_generation()
self.client = self.llm_config._client
self.model = self.llm_config.model
self.context_format = context_format
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_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,7 +345,13 @@ 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"
@ -172,21 +364,21 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
"role": "user",
"content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
**How to Answer:**
{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 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.
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 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 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 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.
@ -206,9 +398,13 @@ Answer:
],
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 +423,9 @@ 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
):
"""
Run the LongMemEval benchmark.
@ -247,14 +445,17 @@ 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.
"""
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 +517,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 +592,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)
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 +688,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 +777,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 +857,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 +912,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 +926,19 @@ 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."
)
args = parser.parse_args()
@ -663,13 +947,9 @@ 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")
results = asyncio.run(run_benchmark(
max_instances=args.max_instances,
@ -685,5 +965,7 @@ 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
))