From 12ba08b4a8948233646c117a9a6a2999f39114c9 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 25 Jun 2026 12:59:28 +0100 Subject: [PATCH 01/20] Extra config + dry run option --- .../adapters/database/engram_loader.py | 57 ++++++++++++++++--- .../internal/core/services/populate_db.py | 14 ++++- 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 0088bd2..59112d7 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -7,10 +7,13 @@ import os import time -from typing import Optional, Callable from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Callable, Optional -from engram import EngramClient, ConversationInput, MessageInput +from engram import ConversationInput, EngramClient, MessageInput + +PRINT_INTERVAL = 10 def _parse_session_text(session_text: str) -> ConversationInput: @@ -177,7 +180,7 @@ def _submit_all( "msg_index": msg_idx, }) - if verbose and submitted % 50 == 0: + if verbose and submitted % PRINT_INTERVAL == 0: print(f" Submitted {submitted} user messages") if ingest_delay > 0: @@ -217,7 +220,7 @@ def _submit_all( "session_id": session.get("session_id", ""), }) - if verbose and submitted % 50 == 0: + if verbose and submitted % PRINT_INTERVAL == 0: print(f" Submitted {submitted}/{total}") if ingest_delay > 0: @@ -331,7 +334,7 @@ def _poll_and_collect( "memories_created": total_created, }) - if verbose and completed % 50 == 0: + if verbose and completed % PRINT_INTERVAL == 0: print(f" Completed {completed}/{len(records)}") # Set elapsed_seconds to wall-clock time from first submission to last completion @@ -359,6 +362,7 @@ def engram_ingest_all_tenants( verbose: bool = True, ingestion_mode: str = "conversation", on_progress: Optional[Callable[[dict], None]] = None, + dry_run: bool = False, ) -> IngestionResult: """ Ingest sessions for all tenants into Engram. @@ -381,15 +385,23 @@ def engram_ingest_all_tenants( ingestion_mode: ``"conversation"`` (default) submits the full parsed conversation per session. ``"user_messages"`` submits each user message individually as a single-message conversation. + dry_run: If True, count requests without submitting to Engram. + No API key is required. Returns an ``IngestionResult`` with + synthetic run records reflecting the would-be request count. Returns: An ``IngestionResult`` with run records and submission timing. If ``poll=True``, ``result.stats`` contains per-tenant stats. """ - client = EngramClient( - api_key=engram_api_key or os.environ["ENGRAM_API_KEY"], - base_url=engram_base_url, - ) + if dry_run: + client = _DryRunEngramClient() + ingest_delay = 0.0 + poll = False + else: + client = EngramClient( + api_key=engram_api_key or os.environ["ENGRAM_API_KEY"], + base_url=engram_base_url, + ) t0 = time.time() @@ -405,6 +417,14 @@ def engram_ingest_all_tenants( for rec in records: tenant_session_counts[rec.tenant_id] = tenant_session_counts.get(rec.tenant_id, 0) + 1 + if dry_run and verbose: + print( + f"\n[Dry run] Would submit {len(records)} requests across {len(tenant_session_counts)} tenants" + ) + for tid in sorted(tenant_session_counts): + print(f" {tid}: {tenant_session_counts[tid]} requests") + print("[Dry run] No data was sent to Engram.") + result = IngestionResult( run_records=records, tenant_session_counts=tenant_session_counts, @@ -440,3 +460,22 @@ def engram_ingest_all_tenants( ) return result + + +class _DryRunEngramClient: + """No-op Engram client for dry-run counting — no HTTP calls made.""" + + def __init__(self): + self._count = 0 + self.memories = self + self.runs = self + + def add(self, *args, **kwargs): + self._count += 1 + return SimpleNamespace(run_id=f"dry-run-{self._count}") + + def get(self, *args, **kwargs): + return SimpleNamespace( + status="completed", + committed_operations=SimpleNamespace(created=[], updated=[], deleted=[]), + ) diff --git a/query_agent_benchmarking/internal/core/services/populate_db.py b/query_agent_benchmarking/internal/core/services/populate_db.py index a6c4ee1..4e39120 100644 --- a/query_agent_benchmarking/internal/core/services/populate_db.py +++ b/query_agent_benchmarking/internal/core/services/populate_db.py @@ -72,8 +72,20 @@ def _run_engram_loader(config: dict, dataset_name: str) -> dict: tenant_ids=tenant_ids, ) + engram_base_url = config.get("engram_base_url", "https://dev-engram.labs.weaviate.io") ingestion_mode = config.get("engram_ingestion_mode", "conversation") - result = engram_ingest_all_tenants(docs_by_tenant, poll=True, ingestion_mode=ingestion_mode) + user_id_prefix = config.get("user_id_prefix", "longmemeval-") + engram_dry_run = config.get("engram_dry_run", False) + result = engram_ingest_all_tenants( + docs_by_tenant, + engram_base_url=engram_base_url, + poll=True, + ingestion_mode=ingestion_mode, + user_id_prefix=user_id_prefix, + dry_run=engram_dry_run, + ) + if engram_dry_run: + return {} manifest = save_engram_manifest(result, dataset_name) return manifest From 15aa6746bf70f981e507530291c55674b0057898 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 25 Jun 2026 14:50:43 +0100 Subject: [PATCH 02/20] Add message_turn input mode --- .../adapters/database/engram_loader.py | 131 +++++++++--------- 1 file changed, 66 insertions(+), 65 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 59112d7..135ddee 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -100,6 +100,46 @@ class TenantIngestionStats: run_durations: list[float] = field(default_factory=list) +def _get_inputs_from_conversation( + conversation: ConversationInput, + ingestion_mode: str, +) -> list[ConversationInput]: + """Split a parsed conversation into the list of inputs to submit to Engram. + + - ``conversation``: one input — the full conversation. + - ``user_messages``: one input per user message (assistant turns discarded). + - ``message_turn``: one input per assistant message encountered, accumulating + all messages since the last flush; any trailing messages after the final + assistant turn become one last input. + """ + if ingestion_mode == "conversation": + return [conversation] + + if ingestion_mode == "user_messages": + return [ + ConversationInput(messages=[m]) + for m in conversation.messages + if m.role == "user" + ] + + if ingestion_mode == "message_turn": + turns: list[ConversationInput] = [] + current: list[MessageInput] = [] + for msg in conversation.messages: + current.append(msg) + if msg.role == "assistant": + turns.append(ConversationInput(messages=current)) + current = [] + if current: + turns.append(ConversationInput(messages=current)) + return turns + + raise ValueError( + f"Unsupported ingestion_mode '{ingestion_mode}'. " + f"Supported: conversation, user_messages, message_turn" + ) + + def _submit_all( client: EngramClient, docs_by_tenant: dict[str, list[dict]], @@ -113,9 +153,10 @@ def _submit_all( """Submit every session across all tenants. Returns run records for polling. Args: - ingestion_mode: ``"conversation"`` submits the full parsed conversation - per session. ``"user_messages"`` submits each user message - individually as a single-message conversation (no assistant context). + ingestion_mode: How to split each session into Engram ``add()`` calls. + ``"conversation"`` submits the full parsed conversation per session. + ``"user_messages"`` submits each user message individually. + ``"message_turn"`` submits each user/assistant exchange as one input. """ records: list[RunRecord] = [] total = sum(len(sessions) for sessions in docs_by_tenant.values()) @@ -135,62 +176,19 @@ def _submit_all( print(f" Skipped session {sid} for tenant {tenant_id}: no valid messages after parsing") continue - if ingestion_mode == "user_messages": - # Submit each user message as its own single-message conversation - user_messages = [m for m in conversation.messages if m.role == "user"] - if not user_messages: - skipped += 1 - if verbose: - sid = session.get("session_id", "?") - print(f" Skipped session {sid} for tenant {tenant_id}: no user messages") - continue - for msg_idx, msg in enumerate(user_messages): - single = ConversationInput(messages=[msg]) - t_submit = time.time() - try: - run = client.memories.add( - single, - user_id=user_id, - group=group, - properties={"conversation_id": session.get("session_id")}, - ) - except Exception as e: - skipped += 1 - if verbose: - sid = session.get("session_id", "?") - print(f" Skipped session {sid} msg {msg_idx} for tenant {tenant_id}: {e}") - continue - records.append(RunRecord( - run_id=run.run_id, - tenant_id=tenant_id, - submitted_at=t_submit, - session_id=session.get("session_id", ""), - session_date=session.get("session_date", ""), - )) - submitted += 1 - - if on_progress: - on_progress({ - "phase": "submit", - "submitted": submitted, - "skipped": skipped, - "total": total, - "tenant_id": tenant_id, - "session_id": session.get("session_id", ""), - "msg_index": msg_idx, - }) - - if verbose and submitted % PRINT_INTERVAL == 0: - print(f" Submitted {submitted} user messages") - - if ingest_delay > 0: - time.sleep(ingest_delay) - else: - # Default: submit the full conversation + inputs = _get_inputs_from_conversation(conversation, ingestion_mode) + if not inputs: + skipped += 1 + if verbose: + sid = session.get("session_id", "?") + print(f" Skipped session {sid} for tenant {tenant_id}: no inputs after splitting") + continue + + for input_idx, conv_input in enumerate(inputs): t_submit = time.time() try: run = client.memories.add( - conversation, + conv_input, user_id=user_id, group=group, properties={"conversation_id": session.get("session_id")}, @@ -199,7 +197,7 @@ def _submit_all( skipped += 1 if verbose: sid = session.get("session_id", "?") - print(f" Skipped session {sid} for tenant {tenant_id}: {e}") + print(f" Skipped session {sid} input {input_idx} for tenant {tenant_id}: {e}") continue records.append(RunRecord( run_id=run.run_id, @@ -218,19 +216,19 @@ def _submit_all( "total": total, "tenant_id": tenant_id, "session_id": session.get("session_id", ""), + "input_index": input_idx, }) if verbose and submitted % PRINT_INTERVAL == 0: - print(f" Submitted {submitted}/{total}") + print(f" Submitted {submitted}") if ingest_delay > 0: time.sleep(ingest_delay) if verbose: - mode_label = "user messages" if ingestion_mode == "user_messages" else "sessions" - print(f" All {submitted} {mode_label} submitted across {len(docs_by_tenant)} tenants") + print(f" All {submitted} inputs submitted across {len(docs_by_tenant)} tenants (mode: {ingestion_mode})") if skipped: - print(f" Skipped {skipped} due to errors") + print(f" Skipped {skipped} due to errors or empty splits") return records @@ -382,9 +380,12 @@ def engram_ingest_all_tenants( poll: Whether to poll runs to completion and collect stats. poll_interval: Seconds between run-status polls (only used when poll=True). verbose: Print progress updates. - ingestion_mode: ``"conversation"`` (default) submits the full parsed - conversation per session. ``"user_messages"`` submits each user - message individually as a single-message conversation. + ingestion_mode: How to split each session into Engram ``add()`` calls. + ``"conversation"`` (default) submits the full conversation per session. + ``"user_messages"`` submits each user message individually. + ``"message_turn"`` submits each user/assistant exchange as one input, + accumulating until each assistant message; trailing user-only messages + become a final input. dry_run: If True, count requests without submitting to Engram. No API key is required. Returns an ``IngestionResult`` with synthetic run records reflecting the would-be request count. @@ -408,7 +409,7 @@ def engram_ingest_all_tenants( # Phase 1: submit everything if verbose: total = sum(len(s) for s in docs_by_tenant.values()) - mode_label = f" (mode: {ingestion_mode})" if ingestion_mode != "conversation" else "" + mode_label = f" (mode: {ingestion_mode})" if ingestion_mode != "conversation" else " (mode: conversation)" print(f"Submitting {total} sessions across {len(docs_by_tenant)} tenants{mode_label}...") records = _submit_all(client, docs_by_tenant, group, user_id_prefix, ingest_delay, verbose, ingestion_mode, on_progress) From 863d348bb2b567cfd2acb703d7d08739f6ea9ed6 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 10:04:58 +0100 Subject: [PATCH 03/20] Print manifest location on save --- .../internal/adapters/results/engram_manifest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/query_agent_benchmarking/internal/adapters/results/engram_manifest.py b/query_agent_benchmarking/internal/adapters/results/engram_manifest.py index a9fccab..88b6db1 100644 --- a/query_agent_benchmarking/internal/adapters/results/engram_manifest.py +++ b/query_agent_benchmarking/internal/adapters/results/engram_manifest.py @@ -79,4 +79,5 @@ def save_engram_manifest(result: IngestionResult, dataset_name: str) -> dict: with open(filepath, "w") as f: json.dump(manifest, f, indent=2) + print(f"Manifest: {filename}") return manifest From 030cfdf9567b7a7dd664191048ac26fc30e6604a Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 10:06:59 +0100 Subject: [PATCH 04/20] Flag to disable engram run polling on submit --- .../internal/core/services/populate_db.py | 10 ++++++---- scripts/populate-db.py | 7 ++++++- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/query_agent_benchmarking/internal/core/services/populate_db.py b/query_agent_benchmarking/internal/core/services/populate_db.py index 4e39120..f4a58f1 100644 --- a/query_agent_benchmarking/internal/core/services/populate_db.py +++ b/query_agent_benchmarking/internal/core/services/populate_db.py @@ -30,7 +30,7 @@ ) -def populate_db(recreate: bool = True, tag: str = "Default") -> None: +def populate_db(recreate: bool = True, tag: str = "Default", poll: bool = True) -> None: """ Load dataset from config and populate the target database. @@ -41,6 +41,8 @@ def populate_db(recreate: bool = True, tag: str = "Default") -> None: Args: recreate: Whether to drop existing collection before creating. tag: Suffix to add to collection name. + poll: Whether to poll Engram runs to completion. When False, submits + fire-and-forget and returns immediately. """ config = load_config(DEFAULT_CONFIG_PATH) @@ -49,14 +51,14 @@ def populate_db(recreate: bool = True, tag: str = "Default") -> None: database_target = config.get("database_target", "weaviate") if database_target == "engram": - manifest = _run_engram_loader(config, dataset_name) + manifest = _run_engram_loader(config, dataset_name, poll=poll) print(f"\nEngram manifest saved with {len(manifest.get('runs', []))} run records") return _run_weaviate_loader(config, dataset_name, recreate=recreate, tag=tag) -def _run_engram_loader(config: dict, dataset_name: str) -> dict: +def _run_engram_loader(config: dict, dataset_name: str, poll: bool = True) -> dict: """Ingest LongMemEval sessions into Engram and persist a run manifest. Returns: @@ -79,7 +81,7 @@ def _run_engram_loader(config: dict, dataset_name: str) -> dict: result = engram_ingest_all_tenants( docs_by_tenant, engram_base_url=engram_base_url, - poll=True, + poll=poll, ingestion_mode=ingestion_mode, user_id_prefix=user_id_prefix, dry_run=engram_dry_run, diff --git a/scripts/populate-db.py b/scripts/populate-db.py index 2638f8e..22cabc2 100644 --- a/scripts/populate-db.py +++ b/scripts/populate-db.py @@ -1,3 +1,8 @@ +import argparse import query_agent_benchmarking -query_agent_benchmarking.populate_db() +parser = argparse.ArgumentParser() +parser.add_argument("--no-poll", action="store_true", help="Submit to Engram without waiting for runs to complete") +args = parser.parse_args() + +query_agent_benchmarking.populate_db(poll=not args.no_poll) From 83c745c37023deed54b29adfb0b5bb18dc2a5094 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 10:28:46 +0100 Subject: [PATCH 05/20] Prompt --- .../adapters/agents/engram_dspy_agent.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 4884e76..e1b3b25 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -43,36 +43,35 @@ class AnswerUserQueryWithMemory(dspy.Signature): Determine the question type, then apply the matching strategy: - FACTUAL questions ("What is my cat's name?", "Where do I work?", "How many engineers do I lead?"): + FACTUAL questions: - Answer using ONLY facts explicitly stated in the memories. - Only mention specific names, numbers, products, or details if they appear in the memories. - Do NOT fabricate or guess facts. - PREFERENCE / RECOMMENDATION questions ("What should I cook?", "What should I serve?", "What would I enjoy?"): + PREFERENCE / RECOMMENDATION questions: - These ask you to USE the user's stored preferences, habits, and interests to make a personalized suggestion. - You SHOULD synthesize across memories to form a recommendation grounded in what the user likes, has, or has done. - - Example: if memories say the user grows basil and cherry tomatoes, and they ask what to cook, suggest dishes featuring those ingredients. This is the correct behavior, not fabrication. - Do NOT abstain just because no memory literally answers the question — the memories provide the ingredients for your recommendation. BOTH types: - Synthesizing facts across multiple memories is expected and encouraged. - BEFORE answering, check the question's premises against the memories (see premise_check). - - If the memories contain NO relevant information at all, say "The information provided is not enough." + - If the memories contain NO relevant information at all, say "The information provided is not enough." and explain why fully. Contradiction vs. gap: - - A CONTRADICTION is when the question assumes a specific fact (role title, name, date, event, location) and the memories state a DIFFERENT fact. Example: the question says "Software Engineer Manager" but memories say "Senior Software Engineer" — these are two different roles, so this is a contradiction. + - A CONTRADICTION is when the question assumes a specific fact (role title, name, date, event, location) and the memories state a DIFFERENT fact. - A GAP is when the memories are simply silent on a topic — they neither confirm nor deny it. - If you find a contradiction: you MUST abstain. Say "The information provided is not enough." then explain what the question assumed vs. what the memories actually say. - If you find only gaps: answer with whatever relevant information IS available. Do NOT abstain just because the answer is incomplete. Temporal reasoning: - When memories describe the same thing at different points in time, resolve the timeline. Look for language that distinguishes past states, plans/intentions, and current states: - * Past: "stored under the bed", "used to", "previously" + * Past: "used to", "previously" * Plans: "plans to", "intends to", "wants to", "will" - * Current: present tense statements of fact ("keeps sneakers in a shoe rack", "is in a shoe rack") - - A current-state memory supersedes a past-state or plan memory about the same topic. If one memory says "sneakers are under the bed" (past) and another says "sneakers are in a shoe rack" (current), the answer is the shoe rack — do NOT hedge or say "unclear." + * Current: present tense statements of fact + - A current-state memory supersedes a past-state or plan memory about the same topic — do NOT hedge or say "unclear." - If a memory describes both a past and current state, treat them as distinct facts at different points in time — do not collapse them. - - When the question asks about the current state ("Where do I currently keep..."), give the most recent state, not the full history. + - When the question asks about the current state, give the most recent state, not the full history. - The timestamps on the memories reflect storage time, NOT when the events originally occurred. Reason about temporal order from the content of the memories, not from the timestamps.""" user_question: str = dspy.InputField() From 468c5edbfcb9182ba2ca05f275c19ac728d788dd Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 11:15:29 +0100 Subject: [PATCH 06/20] Make engram loading async --- .../adapters/database/engram_loader.py | 181 ++++++++++++------ 1 file changed, 124 insertions(+), 57 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 135ddee..3e41c09 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -5,13 +5,14 @@ sessions into Engram's memory system on a per-tenant basis. """ +import asyncio import os import time from dataclasses import dataclass, field from types import SimpleNamespace from typing import Callable, Optional -from engram import ConversationInput, EngramClient, MessageInput +from engram import AsyncEngramClient, ConversationInput, EngramClient, MessageInput PRINT_INTERVAL = 10 @@ -140,27 +141,29 @@ def _get_inputs_from_conversation( ) -def _submit_all( - client: EngramClient, +@dataclass +class _SubmitItem: + """One pending ``memories.add`` call with all metadata needed to record results.""" + user_id: str + conv_input: ConversationInput + tenant_id: str + session_id: str + session_date: str + input_idx: int + + +def _build_per_user_items( docs_by_tenant: dict[str, list[dict]], - group: str, user_id_prefix: str, - ingest_delay: float, + ingestion_mode: str, verbose: bool, - ingestion_mode: str = "conversation", - on_progress: Optional[Callable[[dict], None]] = None, -) -> list[RunRecord]: - """Submit every session across all tenants. Returns run records for polling. +) -> tuple[dict[str, list[_SubmitItem]], int]: + """Pre-process all sessions into per-user ordered submit queues. - Args: - ingestion_mode: How to split each session into Engram ``add()`` calls. - ``"conversation"`` submits the full parsed conversation per session. - ``"user_messages"`` submits each user message individually. - ``"message_turn"`` submits each user/assistant exchange as one input. + Returns ``(per_user_items, skipped_count)`` where ``per_user_items`` maps + ``user_id -> [_SubmitItem, ...]`` in strict session → split order. """ - records: list[RunRecord] = [] - total = sum(len(sessions) for sessions in docs_by_tenant.values()) - submitted = 0 + per_user: dict[str, list[_SubmitItem]] = {} skipped = 0 for tenant_id in sorted(docs_by_tenant.keys()): @@ -184,27 +187,86 @@ def _submit_all( print(f" Skipped session {sid} for tenant {tenant_id}: no inputs after splitting") continue + items = per_user.setdefault(user_id, []) for input_idx, conv_input in enumerate(inputs): - t_submit = time.time() - try: - run = client.memories.add( - conv_input, - user_id=user_id, - group=group, - properties={"conversation_id": session.get("session_id")}, - ) - except Exception as e: - skipped += 1 - if verbose: - sid = session.get("session_id", "?") - print(f" Skipped session {sid} input {input_idx} for tenant {tenant_id}: {e}") - continue - records.append(RunRecord( - run_id=run.run_id, + items.append(_SubmitItem( + user_id=user_id, + conv_input=conv_input, tenant_id=tenant_id, - submitted_at=t_submit, session_id=session.get("session_id", ""), session_date=session.get("session_date", ""), + input_idx=input_idx, + )) + + return per_user, skipped + + +async def _submit_all( + client: "AsyncEngramClient | _DryRunEngramClient", + docs_by_tenant: dict[str, list[dict]], + group: str, + user_id_prefix: str, + ingest_delay: float, + verbose: bool, + ingestion_mode: str = "conversation", + on_progress: Optional[Callable[[dict], None]] = None, +) -> list[RunRecord]: + """Submit every session across all tenants concurrently per user. + + Within each user, items are submitted in strict session → split order. + Across users, items at the same index position are submitted concurrently. + + Args: + ingestion_mode: How to split each session into Engram ``add()`` calls. + ``"conversation"`` submits the full parsed conversation per session. + ``"user_messages"`` submits each user message individually. + ``"message_turn"`` submits each user/assistant exchange as one input. + """ + per_user, skipped = _build_per_user_items( + docs_by_tenant, user_id_prefix, ingestion_mode, verbose + ) + + total = sum(len(items) for items in per_user.values()) + submitted = 0 + max_len = max((len(items) for items in per_user.values()), default=0) + records: list[RunRecord] = [] + + for i in range(max_len): + batch = [ + (uid, items[i]) + for uid, items in per_user.items() + if i < len(items) + ] + + t_submit = time.time() + results = await asyncio.gather( + *[ + client.memories.add( + item.conv_input, + user_id=uid, + group=group, + properties={"conversation_id": item.session_id}, + ) + for uid, item in batch + ], + return_exceptions=True, + ) + + for (uid, item), result in zip(batch, results): + if isinstance(result, Exception): + skipped += 1 + if verbose: + print( + f" Skipped session {item.session_id} input {item.input_idx}" + f" for tenant {item.tenant_id}: {result}" + ) + else: + records.append(RunRecord( + run_id=result.run_id, + tenant_id=item.tenant_id, + submitted_at=t_submit, + session_id=item.session_id, + session_date=item.session_date, )) submitted += 1 @@ -214,19 +276,19 @@ def _submit_all( "submitted": submitted, "skipped": skipped, "total": total, - "tenant_id": tenant_id, - "session_id": session.get("session_id", ""), - "input_index": input_idx, + "tenant_id": item.tenant_id, + "session_id": item.session_id, + "input_index": item.input_idx, }) if verbose and submitted % PRINT_INTERVAL == 0: print(f" Submitted {submitted}") - if ingest_delay > 0: - time.sleep(ingest_delay) + if ingest_delay > 0: + await asyncio.sleep(ingest_delay) if verbose: - print(f" All {submitted} inputs submitted across {len(docs_by_tenant)} tenants (mode: {ingestion_mode})") + print(f" All {submitted} inputs submitted across {len(per_user)} tenants (mode: {ingestion_mode})") if skipped: print(f" Skipped {skipped} due to errors or empty splits") @@ -234,7 +296,7 @@ def _submit_all( def _poll_and_collect( - client: EngramClient, + client, records: list[RunRecord], poll_interval: float, verbose: bool, @@ -354,7 +416,7 @@ def engram_ingest_all_tenants( engram_api_key: Optional[str] = None, group: str = "default", user_id_prefix: str = "longmemeval-", - ingest_delay: float = 0.1, + ingest_delay: float = 0.0, poll: bool = False, poll_interval: float = 2.0, verbose: bool = True, @@ -365,10 +427,12 @@ def engram_ingest_all_tenants( """ Ingest sessions for all tenants into Engram. - Submits all sessions across every tenant first. If ``poll=True``, also - polls every run to completion and populates per-tenant stats with - operation counts and run durations. Otherwise returns immediately after - submission (fire-and-forget). + Submits all sessions across every tenant first. Within each user the + submission order is strictly preserved (session N before N+1, split M + before M+1). Across users, submissions at the same position are issued + concurrently. If ``poll=True``, also polls every run to completion and + populates per-tenant stats with operation counts and run durations. + Otherwise returns immediately after submission (fire-and-forget). Args: docs_by_tenant: Dict mapping tenant_id -> list of session dicts. @@ -376,7 +440,7 @@ def engram_ingest_all_tenants( engram_api_key: Engram API key. Falls back to ``ENGRAM_API_KEY`` env var. group: Engram memory group name. user_id_prefix: Prefix for Engram user IDs. - ingest_delay: Seconds to sleep between session submissions. + ingest_delay: Seconds to sleep between submission rounds. poll: Whether to poll runs to completion and collect stats. poll_interval: Seconds between run-status polls (only used when poll=True). verbose: Print progress updates. @@ -394,24 +458,26 @@ def engram_ingest_all_tenants( An ``IngestionResult`` with run records and submission timing. If ``poll=True``, ``result.stats`` contains per-tenant stats. """ + api_key = engram_api_key or (None if dry_run else os.environ["ENGRAM_API_KEY"]) + if dry_run: - client = _DryRunEngramClient() + submit_client = _DryRunEngramClient() ingest_delay = 0.0 poll = False else: - client = EngramClient( - api_key=engram_api_key or os.environ["ENGRAM_API_KEY"], + submit_client = AsyncEngramClient( + api_key=api_key, base_url=engram_base_url, ) t0 = time.time() - # Phase 1: submit everything + # Phase 1: submit everything (async, concurrent across users) if verbose: total = sum(len(s) for s in docs_by_tenant.values()) mode_label = f" (mode: {ingestion_mode})" if ingestion_mode != "conversation" else " (mode: conversation)" print(f"Submitting {total} sessions across {len(docs_by_tenant)} tenants{mode_label}...") - records = _submit_all(client, docs_by_tenant, group, user_id_prefix, ingest_delay, verbose, ingestion_mode, on_progress) + records = asyncio.run(_submit_all(submit_client, docs_by_tenant, group, user_id_prefix, ingest_delay, verbose, ingestion_mode, on_progress)) submit_elapsed = time.time() - t0 tenant_session_counts = {} @@ -440,10 +506,11 @@ def engram_ingest_all_tenants( print(" Polling disabled — returning immediately (fire-and-forget)") return result - # Phase 2: poll all runs to completion + # Phase 2: poll all runs to completion (sync client — polling is not parallelised) + poll_client = EngramClient(api_key=api_key, base_url=engram_base_url) if verbose: print(f"Polling {len(records)} runs for completion...") - stats_map, completions = _poll_and_collect(client, records, poll_interval, verbose, on_progress) + stats_map, completions = _poll_and_collect(poll_client, records, poll_interval, verbose, on_progress) result.stats = [stats_map[tid] for tid in sorted(stats_map)] result.run_completions = completions @@ -464,18 +531,18 @@ def engram_ingest_all_tenants( class _DryRunEngramClient: - """No-op Engram client for dry-run counting — no HTTP calls made.""" + """No-op async Engram client for dry-run counting — no HTTP calls made.""" def __init__(self): self._count = 0 self.memories = self self.runs = self - def add(self, *args, **kwargs): + async def add(self, *args, **kwargs): self._count += 1 return SimpleNamespace(run_id=f"dry-run-{self._count}") - def get(self, *args, **kwargs): + async def get(self, *args, **kwargs): return SimpleNamespace( status="completed", committed_operations=SimpleNamespace(created=[], updated=[], deleted=[]), From f05434ff472ac0a792305557594bfd63fe0eed4e Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 11:31:38 +0100 Subject: [PATCH 07/20] Make EngramDSPyAgent support async --- .../adapters/agents/engram_dspy_agent.py | 69 ++++++++++++++++++- .../internal/core/services/ask_benchmark.py | 1 - 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index e1b3b25..88d71b9 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -14,6 +14,7 @@ from pydantic import BaseModel from engram import ( + AsyncEngramClient, EngramClient, BM25Retrieval, FetchRetrieval, @@ -125,8 +126,13 @@ def __init__( engram_group: str = "default", user_id_prefix: str = "longmemeval-", ): + api_key = engram_api_key or os.environ["ENGRAM_API_KEY"] self.engram_client = EngramClient( - api_key=engram_api_key or os.environ["ENGRAM_API_KEY"], + api_key=api_key, + base_url=engram_base_url, + ) + self.async_engram_client = AsyncEngramClient( + api_key=api_key, base_url=engram_base_url, ) self.retrieval_limit = retrieval_limit @@ -195,3 +201,64 @@ def run( "memories": [r.model_dump() for r in retrieved], }, ) + + async def run_async( + self, + query: str, + tenant_id: Optional[str] = None, + oracle_context_id: Optional[str] = None, + ) -> EngramAskResponse: + """ + Retrieve memories from Engram and answer the question using DSPy, both async. + + Args: + query: The user's question. + tenant_id: Tenant whose memories to search. + oracle_context_id: Unused, kept for interface compatibility. + """ + if tenant_id is None: + raise ValueError("tenant_id is required for EngramDSPyAgent") + + user_id = f"{self.user_id_prefix}{tenant_id}" + + try: + retrieval_cls = _RETRIEVAL_CLASSES[self.retrieval_type] + except KeyError: + raise ValueError( + f"Unsupported retrieval_type '{self.retrieval_type}'. " + f"Supported: {sorted(_RETRIEVAL_CLASSES)}" + ) + + memories = await self.async_engram_client.memories.search( + query=query, + user_id=user_id, + group=self.engram_group, + retrieval_config=retrieval_cls(limit=self.retrieval_limit), + ) + + retrieved = [ + MemoryWithTimestamp( + memory=m.content, + time_added=str(m.created_at), + ) + for m in memories + ] + + response = await self.qa_system.acall( + user_question=query, + retrieved_memories=retrieved, + ) + + return EngramAskResponse( + final_answer=response.answer, + raw_response={ + "n_memories_retrieved": len(retrieved), + "memories": [r.model_dump() for r in retrieved], + }, + ) + + async def initialize_async(self) -> None: + pass + + async def close_async(self) -> None: + pass diff --git a/query_agent_benchmarking/internal/core/services/ask_benchmark.py b/query_agent_benchmarking/internal/core/services/ask_benchmark.py index 6ea68dd..c898f8e 100644 --- a/query_agent_benchmarking/internal/core/services/ask_benchmark.py +++ b/query_agent_benchmarking/internal/core/services/ask_benchmark.py @@ -141,7 +141,6 @@ async def _run_ask_eval(config: dict[str, Any]) -> dict[str, Any]: config["agent_name"] = agent_name if agent_name == "engram": - use_async = False agent_cfg = _load_agent_config(agent_name) ask_agent = EngramDSPyAgent( engram_base_url=agent_cfg.get("engram_base_url", "https://dev-engram.labs.weaviate.io"), From 5751870444971c963cc109fd6aed70adf5fb987a Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Fri, 26 Jun 2026 12:54:21 +0100 Subject: [PATCH 08/20] Add script to poll run status after loading from manifest --- scripts/poll_engram.py | 223 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) create mode 100644 scripts/poll_engram.py diff --git a/scripts/poll_engram.py b/scripts/poll_engram.py new file mode 100644 index 0000000..8435a43 --- /dev/null +++ b/scripts/poll_engram.py @@ -0,0 +1,223 @@ +"""Poll Engram run status from an ingestion manifest. + +Manifest runs are in submission order and Engram processes each user's runs +in FIFO order, so we can binary-search per tenant to find the latest completed +run rather than checking every run individually. Multiple tenants are searched +concurrently. + +Usage: + # Single pass — print counts and exit + uv run python3 scripts/poll_engram.py engram-ingest-longmemeval-s-20260625-103721.json + + # Loop every N seconds with a tqdm progress bar until all runs finish + uv run python3 scripts/poll_engram.py engram-ingest-longmemeval-s-20260625-103721.json --poll 10 +""" + +import argparse +import json +import os +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +from dotenv import load_dotenv +from engram import EngramClient +from tqdm import tqdm + +from query_agent_benchmarking.internal.adapters.results.serialization import RESULTS_DIR +from query_agent_benchmarking.internal.config.loader import load_config + +load_dotenv() + +_TERMINAL = {"completed", "failed", "deleted"} + +_DB_CONFIG_PATH = ( + Path(__file__).resolve().parent.parent + / "query_agent_benchmarking" + / "internal" + / "config" + / "database_loader_config.yml" +) + + +def _make_client() -> EngramClient: + cfg = load_config(_DB_CONFIG_PATH) + base_url = cfg.get("engram_base_url", "https://dev-engram.labs.weaviate.io") + return EngramClient(api_key=os.environ["ENGRAM_API_KEY"], base_url=base_url) + + +def _is_terminal(client: EngramClient, run_id: str) -> bool: + try: + return client.runs.get(run_id).status in _TERMINAL + except Exception: + return True # treat errors as done so they don't block the search + + +def _binary_search_frontier(client: EngramClient, runs: list[dict], known_frontier: int) -> int: + """Find the rightmost terminal run index using binary search. + + Exploits the FIFO ordering guarantee: if run[k] is terminal then + run[0..k-1] are also terminal; if run[k] is still running then + run[k+1..] are also still running. + + Args: + runs: Tenant's runs in submission order. + known_frontier: Index of the last confirmed-terminal run (-1 = none). + + Returns: + Updated frontier (rightmost terminal index, or known_frontier unchanged). + """ + lo = known_frontier + 1 + hi = len(runs) - 1 + result = known_frontier + + while lo <= hi: + mid = (lo + hi) // 2 + if _is_terminal(client, runs[mid]["run_id"]): + result = mid + lo = mid + 1 # could be more completed to the right + else: + hi = mid - 1 # nothing to the right can be complete yet + + return result + + +def _group_by_tenant(all_runs: list[dict]) -> dict[str, list[dict]]: + """Group runs by tenant_id, preserving manifest (submission) order.""" + groups: dict[str, list[dict]] = {} + for r in all_runs: + groups.setdefault(r["tenant_id"], []).append(r) + return groups + + +def _seed_frontiers(tenant_runs: dict[str, list[dict]]) -> dict[str, int]: + """Pre-seed frontiers from manifest status fields — no API calls needed. + + Scans from the start of each tenant's runs and advances the frontier for + every run that already carries a terminal status in the manifest. + Stops at the first non-terminal entry (monotonicity guarantee). + """ + frontiers: dict[str, int] = {} + for tid, runs in tenant_runs.items(): + frontier = -1 + for i, r in enumerate(runs): + if r.get("status") in _TERMINAL: + frontier = i + else: + break + frontiers[tid] = frontier + return frontiers + + +def _advance_all_frontiers( + client: EngramClient, + tenant_runs: dict[str, list[dict]], + frontiers: dict[str, int], +) -> dict[str, int]: + """Binary-search all incomplete tenants concurrently.""" + pending_tenants = [ + tid for tid, f in frontiers.items() + if f < len(tenant_runs[tid]) - 1 + ] + if not pending_tenants: + return frontiers + + with ThreadPoolExecutor(max_workers=len(pending_tenants)) as executor: + future_to_tid = { + executor.submit( + _binary_search_frontier, client, tenant_runs[tid], frontiers[tid] + ): tid + for tid in pending_tenants + } + for future in as_completed(future_to_tid): + tid = future_to_tid[future] + frontiers[tid] = future.result() + + return frontiers + + +def _count_done(frontiers: dict[str, int]) -> int: + return sum(f + 1 for f in frontiers.values()) + + +def _total_runs(tenant_runs: dict[str, list[dict]]) -> int: + return sum(len(runs) for runs in tenant_runs.values()) + + +def _all_complete(tenant_runs: dict[str, list[dict]], frontiers: dict[str, int]) -> bool: + return all(frontiers[tid] == len(runs) - 1 for tid, runs in tenant_runs.items()) + + +def run_single_pass( + client: EngramClient, + tenant_runs: dict[str, list[dict]], + frontiers: dict[str, int], +) -> None: + frontiers = _advance_all_frontiers(client, tenant_runs, frontiers) + total = _total_runs(tenant_runs) + done = _count_done(frontiers) + print(f"{done}/{total} runs done (still running: {total - done})") + + +def run_progress_loop( + client: EngramClient, + tenant_runs: dict[str, list[dict]], + frontiers: dict[str, int], + interval: float, +) -> None: + total = _total_runs(tenant_runs) + + frontiers = _advance_all_frontiers(client, tenant_runs, frontiers) + bar = tqdm(total=total, initial=_count_done(frontiers), desc="Engram runs", unit="run") + + try: + while not _all_complete(tenant_runs, frontiers): + time.sleep(interval) + prev_done = _count_done(frontiers) + frontiers = _advance_all_frontiers(client, tenant_runs, frontiers) + new_done = _count_done(frontiers) + if new_done > prev_done: + bar.update(new_done - prev_done) + finally: + bar.close() + + done = _count_done(frontiers) + print(f"\nAll {done}/{total} runs finished") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Poll Engram run status from an ingestion manifest") + parser.add_argument( + "manifest", + help="Manifest filename in console/results/ (e.g. engram-ingest-longmemeval-s-20260625-103721.json)", + ) + parser.add_argument( + "--poll", + type=float, + metavar="N", + help="Poll every N seconds with a live progress bar until all runs finish", + ) + args = parser.parse_args() + + manifest_path = RESULTS_DIR / args.manifest + if not manifest_path.exists(): + raise SystemExit(f"Manifest not found: {manifest_path}") + + manifest = json.loads(manifest_path.read_text()) + all_runs = manifest.get("runs", []) + if not all_runs: + raise SystemExit("Manifest contains no runs.") + + tenant_runs = _group_by_tenant(all_runs) + frontiers = _seed_frontiers(tenant_runs) + + client = _make_client() + + if args.poll is not None: + run_progress_loop(client, tenant_runs, frontiers, interval=args.poll) + else: + run_single_pass(client, tenant_runs, frontiers) + + +if __name__ == "__main__": + main() From 5cbf29d14db028ff26b4729f0989a3960b8439df Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 29 Jun 2026 09:47:44 +0100 Subject: [PATCH 09/20] Add topic filtering --- .../internal/adapters/agents/engram_dspy_agent.py | 4 ++++ .../internal/core/services/ask_benchmark.py | 1 + 2 files changed, 5 insertions(+) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 88d71b9..c2c2b1a 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -125,6 +125,7 @@ def __init__( retrieval_type: str = "hybrid", engram_group: str = "default", user_id_prefix: str = "longmemeval-", + search_topics: Optional[list[str]] = None, ): api_key = engram_api_key or os.environ["ENGRAM_API_KEY"] self.engram_client = EngramClient( @@ -139,6 +140,7 @@ def __init__( self.retrieval_type = retrieval_type self.engram_group = engram_group self.user_id_prefix = user_id_prefix + self.search_topics = search_topics # TODO: dspy.configure() sets a global LM — this will overwrite any # existing DSPy LM config in the process. Consider using dspy.context() @@ -179,6 +181,7 @@ def run( user_id=user_id, group=self.engram_group, retrieval_config=retrieval_cls(limit=self.retrieval_limit), + topics=self.search_topics, ) retrieved = [ @@ -234,6 +237,7 @@ async def run_async( user_id=user_id, group=self.engram_group, retrieval_config=retrieval_cls(limit=self.retrieval_limit), + topics=self.search_topics, ) retrieved = [ diff --git a/query_agent_benchmarking/internal/core/services/ask_benchmark.py b/query_agent_benchmarking/internal/core/services/ask_benchmark.py index c898f8e..f3f2221 100644 --- a/query_agent_benchmarking/internal/core/services/ask_benchmark.py +++ b/query_agent_benchmarking/internal/core/services/ask_benchmark.py @@ -150,6 +150,7 @@ async def _run_ask_eval(config: dict[str, Any]) -> dict[str, Any]: retrieval_type=agent_cfg.get("retrieval_type", "hybrid"), engram_group=agent_cfg.get("engram_group", "default"), user_id_prefix=agent_cfg.get("user_id_prefix", "longmemeval-"), + search_topics=agent_cfg.get("search_topics"), ) elif docs_collection or dataset_name: agent_cfg = _load_agent_config(agent_name) From 8d7e20ea94fd6309295311bd6f5ea0ff98c3a889 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 29 Jun 2026 09:48:25 +0100 Subject: [PATCH 10/20] Add --tenants flag to print completed/running tenants --- scripts/poll_engram.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/scripts/poll_engram.py b/scripts/poll_engram.py index 8435a43..69d3b31 100644 --- a/scripts/poll_engram.py +++ b/scripts/poll_engram.py @@ -148,15 +148,26 @@ def _all_complete(tenant_runs: dict[str, list[dict]], frontiers: dict[str, int]) return all(frontiers[tid] == len(runs) - 1 for tid, runs in tenant_runs.items()) +def _print_user_lists(tenant_runs: dict[str, list[dict]], frontiers: dict[str, int]) -> None: + completed = [tid for tid, runs in tenant_runs.items() if frontiers[tid] == len(runs) - 1] + in_progress = [tid for tid in tenant_runs if tid not in completed] + q = '"' + print(f"completed: {','.join(q + tid + q for tid in completed)}") + print(f"in_progress: {','.join(q + tid + q for tid in in_progress)}") + + def run_single_pass( client: EngramClient, tenant_runs: dict[str, list[dict]], frontiers: dict[str, int], + show_tenants: bool = False, ) -> None: frontiers = _advance_all_frontiers(client, tenant_runs, frontiers) total = _total_runs(tenant_runs) done = _count_done(frontiers) print(f"{done}/{total} runs done (still running: {total - done})") + if show_tenants: + _print_user_lists(tenant_runs, frontiers) def run_progress_loop( @@ -164,6 +175,7 @@ def run_progress_loop( tenant_runs: dict[str, list[dict]], frontiers: dict[str, int], interval: float, + show_tenants: bool = False, ) -> None: total = _total_runs(tenant_runs) @@ -183,6 +195,8 @@ def run_progress_loop( done = _count_done(frontiers) print(f"\nAll {done}/{total} runs finished") + if show_tenants: + _print_user_lists(tenant_runs, frontiers) def main() -> None: @@ -197,6 +211,11 @@ def main() -> None: metavar="N", help="Poll every N seconds with a live progress bar until all runs finish", ) + parser.add_argument( + "--tenants", + action="store_true", + help="On exit, print comma-separated lists of completed and in-progress tenants", + ) args = parser.parse_args() manifest_path = RESULTS_DIR / args.manifest @@ -214,9 +233,9 @@ def main() -> None: client = _make_client() if args.poll is not None: - run_progress_loop(client, tenant_runs, frontiers, interval=args.poll) + run_progress_loop(client, tenant_runs, frontiers, interval=args.poll, show_tenants=args.tenants) else: - run_single_pass(client, tenant_runs, frontiers) + run_single_pass(client, tenant_runs, frontiers, show_tenants=args.tenants) if __name__ == "__main__": From f39a53ef039867678eea1c29c0b6e257de6d49ff Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 29 Jun 2026 11:09:03 +0100 Subject: [PATCH 11/20] Order memories by time --- .../internal/adapters/agents/engram_dspy_agent.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index c2c2b1a..7e7ee97 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -72,8 +72,7 @@ class AnswerUserQueryWithMemory(dspy.Signature): * Current: present tense statements of fact - A current-state memory supersedes a past-state or plan memory about the same topic — do NOT hedge or say "unclear." - If a memory describes both a past and current state, treat them as distinct facts at different points in time — do not collapse them. - - When the question asks about the current state, give the most recent state, not the full history. - - The timestamps on the memories reflect storage time, NOT when the events originally occurred. Reason about temporal order from the content of the memories, not from the timestamps.""" + - Memories are ordered from oldest to newest. If multiple memories appear to give different information, you should interpret this as information which has been updated over time. This is not a contradiction - the memory nearest the end is the newest, and so should be considered the current state.""" user_question: str = dspy.InputField() retrieved_memories: list[MemoryWithTimestamp] = dspy.InputField() @@ -187,9 +186,9 @@ def run( retrieved = [ MemoryWithTimestamp( memory=m.content, - time_added=str(m.created_at), + time_added=str(m.updated_at), ) - for m in memories + for m in sorted(memories, key=lambda m: m.updated_at) ] response = self.qa_system( @@ -243,9 +242,9 @@ async def run_async( retrieved = [ MemoryWithTimestamp( memory=m.content, - time_added=str(m.created_at), + time_added=str(m.updated_at), ) - for m in memories + for m in sorted(memories, key=lambda m: m.updated_at) ] response = await self.qa_system.acall( From 8f2e77f4e18d14d8e59b0c6450bb09d9f7aa773d Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 2 Jul 2026 12:21:21 +0100 Subject: [PATCH 12/20] Parse input dates and include in engram inputs --- .../adapters/database/engram_loader.py | 19 +++++- tests/adapters/test_engram_loader.py | 63 +++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 3e41c09..22802ed 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -9,6 +9,7 @@ import os import time from dataclasses import dataclass, field +from datetime import datetime, timezone from types import SimpleNamespace from typing import Callable, Optional @@ -17,6 +18,17 @@ PRINT_INTERVAL = 10 +def _parse_session_date(session_date: str) -> "str | None": + """Parse "2023/05/20 (Sat) 10:58" → "2023-05-20T10:58:00Z". Returns None on failure.""" + if not session_date: + return None + try: + dt = datetime.strptime(session_date, "%Y/%m/%d (%a) %H:%M") + return dt.replace(tzinfo=timezone.utc).isoformat().replace("+00:00", "Z") + except ValueError: + return None + + def _parse_session_text(session_text: str) -> ConversationInput: """Parse a ``user: ... \\nassistant: ...`` session string into a ConversationInput. @@ -118,7 +130,7 @@ def _get_inputs_from_conversation( if ingestion_mode == "user_messages": return [ - ConversationInput(messages=[m]) + ConversationInput(messages=[m], updated_at=conversation.updated_at) for m in conversation.messages if m.role == "user" ] @@ -129,10 +141,10 @@ def _get_inputs_from_conversation( for msg in conversation.messages: current.append(msg) if msg.role == "assistant": - turns.append(ConversationInput(messages=current)) + turns.append(ConversationInput(messages=current, updated_at=conversation.updated_at)) current = [] if current: - turns.append(ConversationInput(messages=current)) + turns.append(ConversationInput(messages=current, updated_at=conversation.updated_at)) return turns raise ValueError( @@ -172,6 +184,7 @@ def _build_per_user_items( for session in sessions: conversation = _parse_session_text(session["session_text"]) + conversation.updated_at = _parse_session_date(session.get("session_date", "")) if not conversation.messages: skipped += 1 if verbose: diff --git a/tests/adapters/test_engram_loader.py b/tests/adapters/test_engram_loader.py index 080c53c..f591888 100644 --- a/tests/adapters/test_engram_loader.py +++ b/tests/adapters/test_engram_loader.py @@ -1,6 +1,10 @@ """Tests for the Engram ingestion loader's session-text parser.""" +import pytest + from query_agent_benchmarking.internal.adapters.database.engram_loader import ( + _get_inputs_from_conversation, + _parse_session_date, _parse_session_text, ) @@ -85,3 +89,62 @@ def test_whitespace_only_message_filtered(self): result = _parse_session_text(text) assert len(result.messages) == 1 assert result.messages[0].role == "assistant" + + +class TestParseSessionDate: + @pytest.mark.parametrize("session_date, expected", [ + ("2023/05/20 (Sat) 10:58", "2023-05-20T10:58:00Z"), + ("2023/05/27 (Sat) 05:39", "2023-05-27T05:39:00Z"), + ("2023/05/28 (Sun) 02:58", "2023-05-28T02:58:00Z"), + ("2000/01/01 (Sat) 00:00", "2000-01-01T00:00:00Z"), + ("", None), + ("not a date", None), + ("2023-05-20", None), + ]) + def test_parse_session_date(self, session_date, expected): + assert _parse_session_date(session_date) == expected + + +class TestGetInputsFromConversation: + """Verify updated_at is propagated to all split ConversationInputs.""" + + _SESSION = "user: Hi\nassistant: Hello!\nuser: How are you?\nassistant: Fine." + _UPDATED_AT = "2023-05-20T10:58:00Z" + + def _make_conversation(self, updated_at=None): + conv = _parse_session_text(self._SESSION) + conv.updated_at = updated_at + return conv + + def test_conversation_mode_preserves_updated_at(self): + conv = self._make_conversation(self._UPDATED_AT) + results = _get_inputs_from_conversation(conv, "conversation") + assert len(results) == 1 + assert results[0].updated_at == self._UPDATED_AT + + def test_conversation_mode_none_updated_at(self): + conv = self._make_conversation(None) + results = _get_inputs_from_conversation(conv, "conversation") + assert results[0].updated_at is None + + def test_user_messages_mode_propagates_updated_at(self): + conv = self._make_conversation(self._UPDATED_AT) + results = _get_inputs_from_conversation(conv, "user_messages") + assert len(results) == 2 + assert all(r.updated_at == self._UPDATED_AT for r in results) + + def test_user_messages_mode_none_updated_at(self): + conv = self._make_conversation(None) + results = _get_inputs_from_conversation(conv, "user_messages") + assert all(r.updated_at is None for r in results) + + def test_message_turn_mode_propagates_updated_at(self): + conv = self._make_conversation(self._UPDATED_AT) + results = _get_inputs_from_conversation(conv, "message_turn") + assert len(results) == 2 + assert all(r.updated_at == self._UPDATED_AT for r in results) + + def test_message_turn_mode_none_updated_at(self): + conv = self._make_conversation(None) + results = _get_inputs_from_conversation(conv, "message_turn") + assert all(r.updated_at is None for r in results) From ae77bc6975e2a18c62abc5fe8ebfe28a71a6afb0 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 2 Jul 2026 12:32:47 +0100 Subject: [PATCH 13/20] Include query date when answering --- .../internal/adapters/agents/engram_dspy_agent.py | 8 ++++++++ .../internal/adapters/agents/external_service.py | 2 ++ .../internal/adapters/agents/weaviate_query_agent.py | 2 ++ .../internal/adapters/dataset/huggingface_loader.py | 1 + query_agent_benchmarking/internal/core/domain/models.py | 1 + .../internal/core/domain/query_execution.py | 2 ++ query_agent_benchmarking/internal/core/ports/ask_agent.py | 4 ++++ query_agent_benchmarking/internal/mocks/agents.py | 4 ++-- tests/domain/test_query_execution.py | 4 ++-- tests/integration/test_ask_e2e.py | 6 +++--- 10 files changed, 27 insertions(+), 7 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 7e7ee97..83c7c00 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -66,6 +66,7 @@ class AnswerUserQueryWithMemory(dspy.Signature): - If you find only gaps: answer with whatever relevant information IS available. Do NOT abstain just because the answer is incomplete. Temporal reasoning: + - The question was asked on the date provided in `question_date`. Treat that as "today" when interpreting any time-relative terms ("now", "yesterday", "this year", etc.). - When memories describe the same thing at different points in time, resolve the timeline. Look for language that distinguishes past states, plans/intentions, and current states: * Past: "used to", "previously" * Plans: "plans to", "intends to", "wants to", "will" @@ -75,6 +76,9 @@ class AnswerUserQueryWithMemory(dspy.Signature): - Memories are ordered from oldest to newest. If multiple memories appear to give different information, you should interpret this as information which has been updated over time. This is not a contradiction - the memory nearest the end is the newest, and so should be considered the current state.""" user_question: str = dspy.InputField() + question_date: str = dspy.InputField( + desc="The date on which this question was asked. Treat this as today's date when interpreting any time-relative language in the question or memories." + ) retrieved_memories: list[MemoryWithTimestamp] = dspy.InputField() premise_check: str = dspy.OutputField( desc="List the key premises embedded in the question (role titles, names, dates, events, quantities). For each, state whether the memories SUPPORT it, CONTRADICT it (memories state a different fact), or are SILENT (memories don't mention it). Only mark CONTRADICT when the memories provide a conflicting fact — silence is not contradiction. If all premises are supported or the memories are simply silent, say 'Premises verified.'" @@ -153,6 +157,7 @@ def run( query: str, tenant_id: Optional[str] = None, oracle_context_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> EngramAskResponse: """ Retrieve memories from Engram and answer the question using DSPy. @@ -193,6 +198,7 @@ def run( response = self.qa_system( user_question=query, + question_date=question_date or "unknown", retrieved_memories=retrieved, ) @@ -209,6 +215,7 @@ async def run_async( query: str, tenant_id: Optional[str] = None, oracle_context_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> EngramAskResponse: """ Retrieve memories from Engram and answer the question using DSPy, both async. @@ -249,6 +256,7 @@ async def run_async( response = await self.qa_system.acall( user_question=query, + question_date=question_date or "unknown", retrieved_memories=retrieved, ) diff --git a/query_agent_benchmarking/internal/adapters/agents/external_service.py b/query_agent_benchmarking/internal/adapters/agents/external_service.py index d19ffab..9c3c57a 100644 --- a/query_agent_benchmarking/internal/adapters/agents/external_service.py +++ b/query_agent_benchmarking/internal/adapters/agents/external_service.py @@ -88,6 +88,7 @@ def run( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: payload = {"question": query} if oracle_context_id is not None: @@ -105,6 +106,7 @@ async def run_async( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: payload = {"question": query} if oracle_context_id is not None: diff --git a/query_agent_benchmarking/internal/adapters/agents/weaviate_query_agent.py b/query_agent_benchmarking/internal/adapters/agents/weaviate_query_agent.py index fc244b3..661aef4 100644 --- a/query_agent_benchmarking/internal/adapters/agents/weaviate_query_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/weaviate_query_agent.py @@ -243,6 +243,7 @@ def run( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: if self._agent is None: self.initialize_sync() @@ -254,6 +255,7 @@ async def run_async( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: response = await self._agent.ask(query) return AskResponse(final_answer=response.final_answer, raw_response=response) diff --git a/query_agent_benchmarking/internal/adapters/dataset/huggingface_loader.py b/query_agent_benchmarking/internal/adapters/dataset/huggingface_loader.py index 93c49e3..b2c157b 100644 --- a/query_agent_benchmarking/internal/adapters/dataset/huggingface_loader.py +++ b/query_agent_benchmarking/internal/adapters/dataset/huggingface_loader.py @@ -423,6 +423,7 @@ def load_ask_longmemeval( ground_truth_answer=item["answer"], tenant_id=tid, question_type=qtype_map.get(qid), + question_date=item.get("question_date"), ) ) diff --git a/query_agent_benchmarking/internal/core/domain/models.py b/query_agent_benchmarking/internal/core/domain/models.py index efa098f..71c6775 100644 --- a/query_agent_benchmarking/internal/core/domain/models.py +++ b/query_agent_benchmarking/internal/core/domain/models.py @@ -46,6 +46,7 @@ class InMemoryAskQuery(BaseModel): oracle_context_id: Optional[str] = None # Optional: for oracle/hard-negative experiments tenant_id: Optional[str] = None # Optional: for multi-tenant datasets (e.g., LongMemEval) question_type: Optional[str] = None # Optional: for type-specific evaluation (e.g., LongMemEval) + question_date: Optional[str] = None # Optional: date the question was asked (e.g., LongMemEval) # ============================================================================ diff --git a/query_agent_benchmarking/internal/core/domain/query_execution.py b/query_agent_benchmarking/internal/core/domain/query_execution.py index 3b55757..5781cdb 100644 --- a/query_agent_benchmarking/internal/core/domain/query_execution.py +++ b/query_agent_benchmarking/internal/core/domain/query_execution.py @@ -181,6 +181,7 @@ def run_ask_queries( query=query.question, oracle_context_id=query.oracle_context_id, tenant_id=query.tenant_id, + question_date=query.question_date, ) query_time_taken = time.time() - query_start_time results.append(AskResult( @@ -233,6 +234,7 @@ async def process_query(query: InMemoryAskQuery, index: int): query=query.question, oracle_context_id=query.oracle_context_id, tenant_id=query.tenant_id, + question_date=query.question_date, ) query_time_taken = time.time() - query_start_time return AskResult( diff --git a/query_agent_benchmarking/internal/core/ports/ask_agent.py b/query_agent_benchmarking/internal/core/ports/ask_agent.py index f5f1c34..c9f8874 100644 --- a/query_agent_benchmarking/internal/core/ports/ask_agent.py +++ b/query_agent_benchmarking/internal/core/ports/ask_agent.py @@ -23,6 +23,7 @@ def run( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: """Execute a synchronous ask query. @@ -30,6 +31,7 @@ def run( query: The question to answer. oracle_context_id: Optional document ID for oracle context. tenant_id: Optional tenant ID for multi-tenant datasets. + question_date: Optional date the question was asked (for temporal grounding). Returns: AskResponse containing the answer. @@ -41,6 +43,7 @@ async def run_async( query: str, oracle_context_id: Optional[str] = None, tenant_id: Optional[str] = None, + question_date: Optional[str] = None, ) -> AskResponse: """Execute an asynchronous ask query. @@ -48,6 +51,7 @@ async def run_async( query: The question to answer. oracle_context_id: Optional document ID for oracle context. tenant_id: Optional tenant ID for multi-tenant datasets. + question_date: Optional date the question was asked (for temporal grounding). Returns: AskResponse containing the answer. diff --git a/query_agent_benchmarking/internal/mocks/agents.py b/query_agent_benchmarking/internal/mocks/agents.py index e170705..d21151b 100644 --- a/query_agent_benchmarking/internal/mocks/agents.py +++ b/query_agent_benchmarking/internal/mocks/agents.py @@ -33,12 +33,12 @@ def __init__(self, responses: dict[str, str] | None = None): self.responses = responses or {} self.call_count = 0 - def run(self, query: str, oracle_context_id=None, tenant_id=None) -> AskResponse: + def run(self, query: str, oracle_context_id=None, tenant_id=None, question_date=None) -> AskResponse: self.call_count += 1 answer = self.responses.get(query, "I don't know.") return AskResponse(final_answer=answer) - async def run_async(self, query: str, oracle_context_id=None, tenant_id=None) -> AskResponse: + async def run_async(self, query: str, oracle_context_id=None, tenant_id=None, question_date=None) -> AskResponse: return self.run(query, oracle_context_id, tenant_id) async def initialize_async(self): diff --git a/tests/domain/test_query_execution.py b/tests/domain/test_query_execution.py index e04ac60..8908971 100644 --- a/tests/domain/test_query_execution.py +++ b/tests/domain/test_query_execution.py @@ -26,7 +26,7 @@ class MockAskAgent: def __init__(self): self.queries_received = [] - def run(self, query, oracle_context_id=None, tenant_id=None): + def run(self, query, oracle_context_id=None, tenant_id=None, question_date=None): self.queries_received.append((query, oracle_context_id, tenant_id)) return AskResponse(final_answer="Test answer") @@ -90,7 +90,7 @@ def test_oracle_context_propagation(self): def test_error_handling(self): class ErrorAgent: - def run(self, query, oracle_context_id=None, tenant_id=None): + def run(self, query, oracle_context_id=None, tenant_id=None, question_date=None): raise RuntimeError("Test error") queries = [ diff --git a/tests/integration/test_ask_e2e.py b/tests/integration/test_ask_e2e.py index 4205f0c..7ae5869 100644 --- a/tests/integration/test_ask_e2e.py +++ b/tests/integration/test_ask_e2e.py @@ -46,7 +46,7 @@ class MockAskAgent: def __init__(self, answers: dict[str, str]): self.answers = answers - def run(self, query, oracle_context_id=None, tenant_id=None): + def run(self, query, oracle_context_id=None, tenant_id=None, question_date=None): answer = self.answers.get(query, "I don't know") return AskResponse(final_answer=answer) @@ -100,7 +100,7 @@ def test_full_pipeline_with_errors(self, sample_ask_queries): """Agent raises errors on some queries -> errors are handled gracefully.""" class ErrorAgent: - def run(self, query, oracle_context_id=None, tenant_id=None): + def run(self, query, oracle_context_id=None, tenant_id=None, question_date=None): if "2+2" in query: raise ValueError("Simulated error") return AskResponse(final_answer="Paris") @@ -186,7 +186,7 @@ def test_tenant_propagation(self): received_tenants = [] class TenantCapturingAgent: - def run(self, query, oracle_context_id=None, tenant_id=None): + def run(self, query, oracle_context_id=None, tenant_id=None, question_date=None): received_tenants.append(tenant_id) return AskResponse(final_answer="Meeting") From 4b054a0809fb1c30f3f06dfc9e8530b7b1d64892 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 2 Jul 2026 12:37:09 +0100 Subject: [PATCH 14/20] Remove memory timestamps from prompt as they might not match artificial benchmark dataset times --- .../adapters/agents/engram_dspy_agent.py | 36 +++++++------------ 1 file changed, 13 insertions(+), 23 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 83c7c00..9c107ce 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -11,7 +11,6 @@ from dataclasses import dataclass, field import dspy -from pydantic import BaseModel from engram import ( AsyncEngramClient, @@ -34,11 +33,6 @@ # DSPy signatures # --------------------------------------------------------------------------- -class MemoryWithTimestamp(BaseModel): - memory: str - time_added: str - - class AnswerUserQueryWithMemory(dspy.Signature): """Answer the user's question using the retrieved memories as your knowledge base. @@ -79,7 +73,7 @@ class AnswerUserQueryWithMemory(dspy.Signature): question_date: str = dspy.InputField( desc="The date on which this question was asked. Treat this as today's date when interpreting any time-relative language in the question or memories." ) - retrieved_memories: list[MemoryWithTimestamp] = dspy.InputField() + retrieved_memories: list[str] = dspy.InputField() premise_check: str = dspy.OutputField( desc="List the key premises embedded in the question (role titles, names, dates, events, quantities). For each, state whether the memories SUPPORT it, CONTRADICT it (memories state a different fact), or are SILENT (memories don't mention it). Only mark CONTRADICT when the memories provide a conflicting fact — silence is not contradiction. If all premises are supported or the memories are simply silent, say 'Premises verified.'" ) @@ -188,13 +182,8 @@ def run( topics=self.search_topics, ) - retrieved = [ - MemoryWithTimestamp( - memory=m.content, - time_added=str(m.updated_at), - ) - for m in sorted(memories, key=lambda m: m.updated_at) - ] + sorted_memories = sorted(memories, key=lambda m: m.updated_at) + retrieved = [m.content for m in sorted_memories] response = self.qa_system( user_question=query, @@ -206,7 +195,10 @@ def run( final_answer=response.answer, raw_response={ "n_memories_retrieved": len(retrieved), - "memories": [r.model_dump() for r in retrieved], + "memories": [ + {"memory": m.content, "time_added": str(m.updated_at)} + for m in sorted_memories + ], }, ) @@ -246,13 +238,8 @@ async def run_async( topics=self.search_topics, ) - retrieved = [ - MemoryWithTimestamp( - memory=m.content, - time_added=str(m.updated_at), - ) - for m in sorted(memories, key=lambda m: m.updated_at) - ] + sorted_memories = sorted(memories, key=lambda m: m.updated_at) + retrieved = [m.content for m in sorted_memories] response = await self.qa_system.acall( user_question=query, @@ -264,7 +251,10 @@ async def run_async( final_answer=response.answer, raw_response={ "n_memories_retrieved": len(retrieved), - "memories": [r.model_dump() for r in retrieved], + "memories": [ + {"memory": m.content, "time_added": str(m.updated_at)} + for m in sorted_memories + ], }, ) From 32122f3eeb7964b3905263ca42de7c50312effe1 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 2 Jul 2026 12:52:07 +0100 Subject: [PATCH 15/20] Optional flag to disable conversation_id in engram loading --- .../internal/adapters/database/engram_loader.py | 12 ++++++++++-- .../internal/core/services/populate_db.py | 2 ++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 22802ed..9116960 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -223,6 +223,7 @@ async def _submit_all( verbose: bool, ingestion_mode: str = "conversation", on_progress: Optional[Callable[[dict], None]] = None, + include_conversation_id: bool = True, ) -> list[RunRecord]: """Submit every session across all tenants concurrently per user. @@ -234,6 +235,8 @@ async def _submit_all( ``"conversation"`` submits the full parsed conversation per session. ``"user_messages"`` submits each user message individually. ``"message_turn"`` submits each user/assistant exchange as one input. + include_conversation_id: Whether to include ``{"conversation_id": ...}`` + in the ``properties`` of each ``memories.add()`` call. """ per_user, skipped = _build_per_user_items( docs_by_tenant, user_id_prefix, ingestion_mode, verbose @@ -258,7 +261,7 @@ async def _submit_all( item.conv_input, user_id=uid, group=group, - properties={"conversation_id": item.session_id}, + properties={"conversation_id": item.session_id} if include_conversation_id else {}, ) for uid, item in batch ], @@ -436,6 +439,7 @@ def engram_ingest_all_tenants( ingestion_mode: str = "conversation", on_progress: Optional[Callable[[dict], None]] = None, dry_run: bool = False, + include_conversation_id: bool = True, ) -> IngestionResult: """ Ingest sessions for all tenants into Engram. @@ -466,6 +470,10 @@ def engram_ingest_all_tenants( dry_run: If True, count requests without submitting to Engram. No API key is required. Returns an ``IngestionResult`` with synthetic run records reflecting the would-be request count. + include_conversation_id: Whether to include ``{"conversation_id": ...}`` + in the ``properties`` of each ``memories.add()`` call. Defaults to + True. Set to False when the Engram instance does not expect this + property. Returns: An ``IngestionResult`` with run records and submission timing. @@ -490,7 +498,7 @@ def engram_ingest_all_tenants( total = sum(len(s) for s in docs_by_tenant.values()) mode_label = f" (mode: {ingestion_mode})" if ingestion_mode != "conversation" else " (mode: conversation)" print(f"Submitting {total} sessions across {len(docs_by_tenant)} tenants{mode_label}...") - records = asyncio.run(_submit_all(submit_client, docs_by_tenant, group, user_id_prefix, ingest_delay, verbose, ingestion_mode, on_progress)) + records = asyncio.run(_submit_all(submit_client, docs_by_tenant, group, user_id_prefix, ingest_delay, verbose, ingestion_mode, on_progress, include_conversation_id)) submit_elapsed = time.time() - t0 tenant_session_counts = {} diff --git a/query_agent_benchmarking/internal/core/services/populate_db.py b/query_agent_benchmarking/internal/core/services/populate_db.py index f4a58f1..b3b614b 100644 --- a/query_agent_benchmarking/internal/core/services/populate_db.py +++ b/query_agent_benchmarking/internal/core/services/populate_db.py @@ -78,6 +78,7 @@ def _run_engram_loader(config: dict, dataset_name: str, poll: bool = True) -> di ingestion_mode = config.get("engram_ingestion_mode", "conversation") user_id_prefix = config.get("user_id_prefix", "longmemeval-") engram_dry_run = config.get("engram_dry_run", False) + include_conversation_id = config.get("engram_include_conversation_id", True) result = engram_ingest_all_tenants( docs_by_tenant, engram_base_url=engram_base_url, @@ -85,6 +86,7 @@ def _run_engram_loader(config: dict, dataset_name: str, poll: bool = True) -> di ingestion_mode=ingestion_mode, user_id_prefix=user_id_prefix, dry_run=engram_dry_run, + include_conversation_id=include_conversation_id, ) if engram_dry_run: return {} From 1bfea16ba0e82ba3d10fb8c577400b5ceb1de905 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 6 Jul 2026 11:37:39 +0100 Subject: [PATCH 16/20] Add reasoning field to longmemeval judge --- .../adapters/metrics/ask_metrics_calculator.py | 4 ++++ .../adapters/metrics/longmemeval_judge.py | 18 +++++++++++------- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py index dc02635..226ebae 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py +++ b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py @@ -221,6 +221,7 @@ def compute(self, results: list[AskResult]) -> dict: print(f"Judge model: {self.model}") alignment_scores = [] + judge_reasonings: list[str | None] = [] query_times = [] misaligned_indices = [] total_input_tokens = 0 @@ -231,6 +232,7 @@ def compute(self, results: list[AskResult]) -> dict: if result.system_answer.startswith("[ERROR]"): print(f"\n\033[91mSkipping evaluation for query {i} due to error.\033[0m") alignment_scores.append(None) + judge_reasonings.append(None) continue qtype = result.query.question_type or "multi-session" @@ -244,6 +246,7 @@ def compute(self, results: list[AskResult]) -> dict: aligned = judge_result["aligned"] score = 1 if aligned else 0 alignment_scores.append(score) + judge_reasonings.append(judge_result.get("reasoning")) query_times.append(result.time_taken) total_input_tokens += judge_result.get("input_tokens", 0) total_output_tokens += judge_result.get("output_tokens", 0) @@ -275,6 +278,7 @@ def compute(self, results: list[AskResult]) -> dict: "metric": "longmemeval_judge", "avg_alignment_score": float(np.mean(valid_scores)) if valid_scores else 0, "alignment_score_scores": alignment_scores, + "judge_reasonings": judge_reasonings, "judge_model": self.model, "ensemble_k": 1, "type_accuracy": type_accuracy, diff --git a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py index 506c7ca..df1cd80 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py +++ b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py @@ -120,12 +120,15 @@ def _build_prompt( class LongMemEvalJudgment(dspy.Signature): """Judge whether a model response is correct given type-specific evaluation criteria. - Answer yes or no only. + First reason through each rule in the criteria step by step, then answer yes or no. """ evaluation_prompt: str = dspy.InputField( description="The full evaluation prompt with type-specific instructions, question, correct answer, and model response." ) + reasoning: str = dspy.OutputField( + description="Go through each rule in the evaluation criteria one by one. For each rule, determine whether the response satisfies it and why. Then give the overall reason for your final verdict." + ) judgment: str = dspy.OutputField( description="Answer yes or no only." ) @@ -168,15 +171,16 @@ def __init__( self.judge = dspy.Predict(LongMemEvalJudgment) - def _call_judge(self, prompt: str) -> tuple[bool, int, int]: + def _call_judge(self, prompt: str) -> tuple[bool, str, int, int]: """Call the LLM judge and parse yes/no response. Returns: - Tuple of (aligned, input_tokens, output_tokens). + Tuple of (aligned, reasoning, input_tokens, output_tokens). """ with dspy.context(lm=self.lm): response = self.judge(evaluation_prompt=prompt) content = response.judgment or "" + reasoning = response.reasoning or "" aligned = "yes" in content.strip().lower() input_tokens = 0 @@ -189,7 +193,7 @@ def _call_judge(self, prompt: str) -> tuple[bool, int, int]: except Exception: pass # Token tracking is best-effort - return aligned, input_tokens, output_tokens + return aligned, reasoning, input_tokens, output_tokens # ------------------------------------------------------------------ # LLMJudge protocol methods @@ -217,7 +221,7 @@ def evaluate( """ qtype = question_type or self.default_question_type prompt = _build_prompt(qtype, question, correct_answer, system_answer, question_id) - aligned, _, _ = self._call_judge(prompt) + aligned, _, _, _ = self._call_judge(prompt) return aligned def evaluate_with_details( @@ -245,7 +249,7 @@ def evaluate_with_details( is_abstention = question_id is not None and "_abs" in question_id prompt = _build_prompt(qtype, question, correct_answer, system_answer, question_id) - aligned, input_tokens, output_tokens = self._call_judge(prompt) + aligned, reasoning, input_tokens, output_tokens = self._call_judge(prompt) return { "aligned": aligned, @@ -255,5 +259,5 @@ def evaluate_with_details( "output_tokens": output_tokens, "votes": 1, "ensemble_k": 1, - "reasoning": None, # LongMemEval judge uses no chain-of-thought + "reasoning": reasoning, } From 633a4ce1115c9cb3974f46eb0590b2cbf637d4bf Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 6 Jul 2026 12:06:22 +0100 Subject: [PATCH 17/20] Wire up ensemble_k --- console/app/results/compare/page.tsx | 54 +++++++++++++---- .../[id]/trial/[trialNum]/page.tsx | 60 +++++++++++++------ console/lib/results.ts | 9 ++- .../metrics/ask_metrics_calculator.py | 9 +-- .../adapters/metrics/longmemeval_judge.py | 32 ++++++---- .../adapters/results/serialization.py | 2 +- .../internal/core/services/ask_benchmark.py | 2 +- 7 files changed, 122 insertions(+), 46 deletions(-) diff --git a/console/app/results/compare/page.tsx b/console/app/results/compare/page.tsx index 70ac1d8..b346b30 100644 --- a/console/app/results/compare/page.tsx +++ b/console/app/results/compare/page.tsx @@ -319,7 +319,16 @@ function buildMarkdownReport( : "✗ wrong"; L(`- **${name}** — ${status}${time}`); L(` - Answer: ${cell.system_answer ? cell.system_answer : "_(empty)_"}`); - if (cell.judge_reasoning) L(` - Judge: ${cell.judge_reasoning}`); + if (cell.judge_reasoning) { + const r = cell.judge_reasoning; + if (Array.isArray(r)) { + (r as { vote: boolean; reasoning: string }[]).forEach((v, i) => + L(` - Judge vote ${i + 1} (${v.vote ? "yes" : "no"}): ${v.reasoning}`) + ); + } else { + L(` - Judge: ${r as string}`); + } + } } else { const status = cell.correct ? "✓ hit" : "✗ miss"; const gt = new Set(row.ground_truth_ids ?? []); @@ -1232,16 +1241,39 @@ function CellDetail({

{cell.system_answer || (empty)}

- {cell.judge_reasoning && ( -
- - Judge reasoning - -

- {cell.judge_reasoning} -

-
- )} + {cell.judge_reasoning && (() => { + const votes: { vote: boolean; reasoning: string }[] = Array.isArray(cell.judge_reasoning) + ? cell.judge_reasoning as { vote: boolean; reasoning: string }[] + : [{ vote: true, reasoning: cell.judge_reasoning as string }]; + return ( +
+ + Judge reasoning{votes.length > 1 ? ` (${votes.length} votes)` : ""} + +
+ {votes.map((v, i) => ( +
+ {votes.length > 1 && ( + + {v.vote ? "✓" : "✗"} Vote {i + 1} + + )} +

+ {v.reasoning} +

+
+ ))} +
+
+ ); + })()} ); } diff --git a/console/app/results/experiments/[id]/trial/[trialNum]/page.tsx b/console/app/results/experiments/[id]/trial/[trialNum]/page.tsx index 805c32c..3d14c72 100644 --- a/console/app/results/experiments/[id]/trial/[trialNum]/page.tsx +++ b/console/app/results/experiments/[id]/trial/[trialNum]/page.tsx @@ -471,23 +471,49 @@ function AskQueriesView({ queries }: { queries: AskQuery[] }) {

- {q.judge_reasoning && ( -
- - Judge Reasoning - -

- {q.judge_reasoning} -

-
- )} + {q.judge_reasoning && (() => { + const votes: { vote: boolean; reasoning: string }[] = Array.isArray(q.judge_reasoning) + ? q.judge_reasoning as { vote: boolean; reasoning: string }[] + : [{ vote: true, reasoning: q.judge_reasoning as string }]; + return ( +
+ + Judge Reasoning{votes.length > 1 ? ` (${votes.length} votes)` : ""} + +
+ {votes.map((v, i) => ( +
+ {votes.length > 1 && ( +
+ + {v.vote ? "✓" : "✗"} Vote {i + 1} + +
+ )} +

+ {v.reasoning} +

+
+ ))} +
+
+ ); + })()} {q.retrieved_context && ( )} diff --git a/console/lib/results.ts b/console/lib/results.ts index a1cb987..1e6aa2d 100644 --- a/console/lib/results.ts +++ b/console/lib/results.ts @@ -5,6 +5,11 @@ import path from "path"; // Types // ============================================================================ +export interface JudgeVote { + vote: boolean; + reasoning: string; +} + export interface TrialMetadata { dataset: string; agent_name: string; @@ -82,7 +87,7 @@ export interface AskQuery { is_error?: boolean; oracle_context_id?: string; tenant_id?: string; - judge_reasoning?: string; + judge_reasoning?: string | JudgeVote[]; question_type?: string; retrieved_context?: Record; } @@ -309,7 +314,7 @@ export interface ComparisonCell { system_answer?: string; score?: number; is_error?: boolean; - judge_reasoning?: string; + judge_reasoning?: string | JudgeVote[]; } /** One query aligned across all compared experiments (joined by question text). */ diff --git a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py index 226ebae..904c29c 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py +++ b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py @@ -212,16 +212,17 @@ class LongMemEvalAskCalculator: evaluation protocol from the LongMemEval paper (Wu et al., ICLR 2025). """ - def __init__(self, model: str = "openai/gpt-5.4", api_key: str | None = None): + def __init__(self, model: str = "openai/gpt-5.4", api_key: str | None = None, ensemble_k: int = 1): self.model = model - self.judge = LongMemEvalJudge(model=model, api_key=api_key) + self.ensemble_k = ensemble_k + self.judge = LongMemEvalJudge(model=model, api_key=api_key, ensemble_k=ensemble_k) def compute(self, results: list[AskResult]) -> dict: print(f"\n\033[94mAnalyzing {len(results)} ask results with LongMemEval judge...\033[0m") print(f"Judge model: {self.model}") alignment_scores = [] - judge_reasonings: list[str | None] = [] + judge_reasonings: list[list[dict] | None] = [] query_times = [] misaligned_indices = [] total_input_tokens = 0 @@ -280,7 +281,7 @@ def compute(self, results: list[AskResult]) -> dict: "alignment_score_scores": alignment_scores, "judge_reasonings": judge_reasonings, "judge_model": self.model, - "ensemble_k": 1, + "ensemble_k": self.ensemble_k, "type_accuracy": type_accuracy, "type_counts": {t: len(scores) for t, scores in sorted(type_scores.items())}, } diff --git a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py index df1cd80..a8d2062 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py +++ b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py @@ -139,8 +139,8 @@ class LongMemEvalJudge: Uses DSPy for LLM inference with type-specific prompts per question category, following the paper's setup: - - Simple yes/no output (no chain-of-thought) - Type-specific prompts per question category + - Optional ensemble voting: call the judge ensemble_k times, majority vote wins Satisfies the LLMJudge protocol from core/ports/llm_judge.py. """ @@ -151,6 +151,7 @@ def __init__( api_key: Optional[str] = None, default_question_type: str = "multi-session", cache: bool = False, + ensemble_k: int = 1, ): """Initialize the LongMemEval judge. @@ -162,6 +163,7 @@ def __init__( """ self.model = model self.default_question_type = default_question_type + self.ensemble_k = ensemble_k self.lm = dspy.LM( model, @@ -221,8 +223,8 @@ def evaluate( """ qtype = question_type or self.default_question_type prompt = _build_prompt(qtype, question, correct_answer, system_answer, question_id) - aligned, _, _, _ = self._call_judge(prompt) - return aligned + votes = sum(self._call_judge(prompt)[0] for _ in range(self.ensemble_k)) + return votes >= self.ensemble_k / 2 def evaluate_with_details( self, @@ -249,15 +251,25 @@ def evaluate_with_details( is_abstention = question_id is not None and "_abs" in question_id prompt = _build_prompt(qtype, question, correct_answer, system_answer, question_id) - aligned, reasoning, input_tokens, output_tokens = self._call_judge(prompt) + vote_results = [] + total_input_tokens = 0 + total_output_tokens = 0 + for _ in range(self.ensemble_k): + aligned_i, reasoning_i, in_tok, out_tok = self._call_judge(prompt) + vote_results.append({"vote": aligned_i, "reasoning": reasoning_i}) + total_input_tokens += in_tok + total_output_tokens += out_tok + + votes = sum(v["vote"] for v in vote_results) + majority_aligned = votes >= self.ensemble_k / 2 return { - "aligned": aligned, + "aligned": majority_aligned, "question_type": qtype, "is_abstention": is_abstention, - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "votes": 1, - "ensemble_k": 1, - "reasoning": reasoning, + "input_tokens": total_input_tokens, + "output_tokens": total_output_tokens, + "votes": votes, + "ensemble_k": self.ensemble_k, + "reasoning": vote_results, } diff --git a/query_agent_benchmarking/internal/adapters/results/serialization.py b/query_agent_benchmarking/internal/adapters/results/serialization.py index 2cf34b8..cc7baff 100644 --- a/query_agent_benchmarking/internal/adapters/results/serialization.py +++ b/query_agent_benchmarking/internal/adapters/results/serialization.py @@ -113,7 +113,7 @@ def save_ask_trial_results( config: dict[str, Any], trial_number: int, alignment_scores: list[int] | None = None, - judge_reasonings: list[str | None] | None = None, + judge_reasonings: list[str | list[dict] | None] | None = None, ) -> None: """ Save raw query results for a single ask trial. diff --git a/query_agent_benchmarking/internal/core/services/ask_benchmark.py b/query_agent_benchmarking/internal/core/services/ask_benchmark.py index f3f2221..28f7ea6 100644 --- a/query_agent_benchmarking/internal/core/services/ask_benchmark.py +++ b/query_agent_benchmarking/internal/core/services/ask_benchmark.py @@ -174,7 +174,7 @@ async def _run_ask_eval(config: dict[str, Any]) -> dict[str, Any]: metrics_across_trials = [] if use_longmemeval: - metrics_calculator = LongMemEvalAskCalculator(model=judge_model) + metrics_calculator = LongMemEvalAskCalculator(model=judge_model, ensemble_k=ensemble_k) elif use_officeqa_metric: metrics_calculator = OfficeQAAskCalculator(tolerance=officeqa_tolerance) elif use_exact_match: From db7f68b0e0dd697f7bc3efe3e835b03dadde5c69 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 6 Jul 2026 12:32:43 +0100 Subject: [PATCH 18/20] Async longmemeval judge in batches --- .../metrics/ask_metrics_calculator.py | 79 ++++++++++++++++++- .../adapters/metrics/longmemeval_judge.py | 52 ++++++++++++ .../internal/core/services/ask_benchmark.py | 10 ++- 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py index 904c29c..00abb89 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py +++ b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py @@ -7,8 +7,11 @@ - LongMemEvalAskCalculator: Uses LongMemEval's type-specific LLM judge """ +import asyncio + import numpy as np from tqdm import tqdm +from tqdm.asyncio import tqdm as atqdm from query_agent_benchmarking.internal.core.domain.models import AskResult from query_agent_benchmarking.internal.adapters.metrics.lmjudge_alignment import LMJudge @@ -212,9 +215,11 @@ class LongMemEvalAskCalculator: evaluation protocol from the LongMemEval paper (Wu et al., ICLR 2025). """ - def __init__(self, model: str = "openai/gpt-5.4", api_key: str | None = None, ensemble_k: int = 1): + def __init__(self, model: str = "openai/gpt-5.4", api_key: str | None = None, + ensemble_k: int = 1, max_concurrent_judge: int = 10): self.model = model self.ensemble_k = ensemble_k + self.max_concurrent_judge = max_concurrent_judge self.judge = LongMemEvalJudge(model=model, api_key=api_key, ensemble_k=ensemble_k) def compute(self, results: list[AskResult]) -> dict: @@ -290,6 +295,78 @@ def compute(self, results: list[AskResult]) -> dict: total_input_tokens, total_output_tokens, type_accuracy) return results_dict + async def compute_async(self, results: list[AskResult]) -> dict: + """Async version of compute() — runs all judge calls concurrently.""" + semaphore = asyncio.Semaphore(self.max_concurrent_judge) + + async def evaluate_one(i: int, result: AskResult): + if result.system_answer.startswith("[ERROR]"): + return i, None, None, result.time_taken + qtype = result.query.question_type or "multi-session" + async with semaphore: + judge_result = await self.judge.evaluate_with_details_async( + question=result.query.question, + system_answer=result.system_answer, + correct_answer=result.query.ground_truth_answer, + question_type=qtype, + ) + return i, judge_result, qtype, result.time_taken + + print(f"\n\033[94mAnalyzing {len(results)} ask results with LongMemEval judge (async, max_concurrent={self.max_concurrent_judge})...\033[0m") + print(f"Judge model: {self.model}, Ensemble K: {self.ensemble_k}") + + raw = await atqdm.gather( + *[evaluate_one(i, r) for i, r in enumerate(results)], + desc="Running LongMemEval judge", + ) + + alignment_scores: list[int | None] = [] + judge_reasonings: list[list[dict] | None] = [] + query_times = [] + misaligned_indices = [] + total_input_tokens = 0 + total_output_tokens = 0 + type_scores: dict[str, list[int]] = {} + + for i, judge_result, qtype, time_taken in sorted(raw, key=lambda x: x[0]): + if judge_result is None: + alignment_scores.append(None) + judge_reasonings.append(None) + continue + score = 1 if judge_result["aligned"] else 0 + alignment_scores.append(score) + judge_reasonings.append(judge_result.get("reasoning")) + query_times.append(time_taken) + total_input_tokens += judge_result.get("input_tokens", 0) + total_output_tokens += judge_result.get("output_tokens", 0) + type_scores.setdefault(qtype, []).append(score) + if not judge_result["aligned"]: + misaligned_indices.append(i) + + type_accuracy = { + t: float(np.mean(scores)) for t, scores in sorted(type_scores.items()) + } + valid_scores = [s for s in alignment_scores if s is not None] + results_dict = { + "avg_query_time": float(np.mean(query_times)) if query_times else 0, + "query_times": query_times, + "misaligned_indices": misaligned_indices, + "total_input_tokens": total_input_tokens, + "total_output_tokens": total_output_tokens, + "metric": "longmemeval_judge", + "avg_alignment_score": float(np.mean(valid_scores)) if valid_scores else 0, + "alignment_score_scores": alignment_scores, + "judge_reasonings": judge_reasonings, + "judge_model": self.model, + "ensemble_k": self.ensemble_k, + "type_accuracy": type_accuracy, + "type_counts": {t: len(scores) for t, scores in sorted(type_scores.items())}, + } + + self._print_summary(valid_scores, results_dict, misaligned_indices, + total_input_tokens, total_output_tokens, type_accuracy) + return results_dict + @staticmethod def _print_summary(alignment_scores, results_dict, misaligned_indices, total_input_tokens, total_output_tokens, type_accuracy): diff --git a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py index a8d2062..179ff07 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py +++ b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py @@ -12,6 +12,7 @@ Output format: simple yes/no (no chain-of-thought reasoning). """ +import asyncio import os from typing import Optional @@ -173,6 +174,26 @@ def __init__( self.judge = dspy.Predict(LongMemEvalJudgment) + async def _call_judge_async(self, prompt: str) -> tuple[bool, str, int, int]: + """Async version of _call_judge using dspy.Predict.acall().""" + with dspy.context(lm=self.lm): + response = await self.judge.acall(evaluation_prompt=prompt) + content = response.judgment or "" + reasoning = response.reasoning or "" + aligned = "yes" in content.strip().lower() + + input_tokens = 0 + output_tokens = 0 + try: + usage = response.get_lm_usage() + if usage and self.model in usage: + input_tokens = usage[self.model].get("prompt_tokens", 0) + output_tokens = usage[self.model].get("completion_tokens", 0) + except Exception: + pass + + return aligned, reasoning, input_tokens, output_tokens + def _call_judge(self, prompt: str) -> tuple[bool, str, int, int]: """Call the LLM judge and parse yes/no response. @@ -273,3 +294,34 @@ def evaluate_with_details( "ensemble_k": self.ensemble_k, "reasoning": vote_results, } + + async def evaluate_with_details_async( + self, + question: str, + system_answer: str, + correct_answer: str, + question_type: Optional[str] = None, + question_id: Optional[str] = None, + ) -> dict: + """Async version of evaluate_with_details. Runs all ensemble_k calls concurrently.""" + qtype = question_type or self.default_question_type + is_abstention = question_id is not None and "_abs" in question_id + prompt = _build_prompt(qtype, question, correct_answer, system_answer, question_id) + + call_results = await asyncio.gather(*[ + self._call_judge_async(prompt) for _ in range(self.ensemble_k) + ]) + + vote_results = [{"vote": a, "reasoning": r} for a, r, _, _ in call_results] + votes = sum(v["vote"] for v in vote_results) + + return { + "aligned": votes >= self.ensemble_k / 2, + "question_type": qtype, + "is_abstention": is_abstention, + "input_tokens": sum(c[2] for c in call_results), + "output_tokens": sum(c[3] for c in call_results), + "votes": votes, + "ensemble_k": self.ensemble_k, + "reasoning": vote_results, + } diff --git a/query_agent_benchmarking/internal/core/services/ask_benchmark.py b/query_agent_benchmarking/internal/core/services/ask_benchmark.py index 28f7ea6..788cdc1 100644 --- a/query_agent_benchmarking/internal/core/services/ask_benchmark.py +++ b/query_agent_benchmarking/internal/core/services/ask_benchmark.py @@ -169,12 +169,15 @@ async def _run_ask_eval(config: dict[str, Any]) -> dict[str, Any]: judge_model = config.get("judge_model", "openai/gpt-4.1") ensemble_k = config.get("ensemble_k", 3) use_reasoning = config.get("use_reasoning", False) + max_concurrent_judge = config.get("max_concurrent_judge", 10) num_trials = config.get("num_trials", 1) metrics_across_trials = [] if use_longmemeval: - metrics_calculator = LongMemEvalAskCalculator(model=judge_model, ensemble_k=ensemble_k) + metrics_calculator = LongMemEvalAskCalculator( + model=judge_model, ensemble_k=ensemble_k, max_concurrent_judge=max_concurrent_judge + ) elif use_officeqa_metric: metrics_calculator = OfficeQAAskCalculator(tolerance=officeqa_tolerance) elif use_exact_match: @@ -221,7 +224,10 @@ async def _run_ask_eval(config: dict[str, Any]) -> dict[str, Any]: sleep_between_requests=config.get("sleep_between_requests", 0.0), ) - metrics = metrics_calculator.compute(results) + if hasattr(metrics_calculator, 'compute_async'): + metrics = await metrics_calculator.compute_async(results) + else: + metrics = metrics_calculator.compute(results) print(f"\n\033[92mTrial {trial+1} Results:\033[0m") print(f" {score_key}: {metrics[score_key]:.2%}") From 1275ee7e7337cc0ab8de64c776589e0d245560c0 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 6 Jul 2026 12:37:09 +0100 Subject: [PATCH 19/20] Use bools for final judgement --- .../adapters/metrics/longmemeval_judge.py | 27 ++++++++++--------- tests/adapters/test_longmemeval_judge.py | 4 +-- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py index 179ff07..cfdb27d 100644 --- a/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py +++ b/query_agent_benchmarking/internal/adapters/metrics/longmemeval_judge.py @@ -29,7 +29,8 @@ "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. " "\n\nQuestion: {question}\n\nCorrect Answer: {answer}\n\nModel Response: {response}" - "\n\nIs the model response correct? Answer yes or no only." + "\n\nIs the model response correct? " + "First reason through each rule in the criteria step by step, then give a boolean verdict" ) _TEMPORAL_REASONING_TEMPLATE = ( @@ -42,7 +43,8 @@ "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. " "\n\nQuestion: {question}\n\nCorrect Answer: {answer}\n\nModel Response: {response}" - "\n\nIs the model response correct? Answer yes or no only." + "\n\nIs the model response correct? " + "First reason through each rule in the criteria step by step, then give a boolean verdict" ) _KNOWLEDGE_UPDATE_TEMPLATE = ( @@ -56,7 +58,8 @@ "vs. 'in a shoe rack in my closet'), it should be considered correct. Only answer no if " "the response gives a fundamentally different answer or misses the key updated fact." "\n\nQuestion: {question}\n\nCorrect Answer: {answer}\n\nModel Response: {response}" - "\n\nIs the model response correct? Answer yes or no only." + "\n\nIs the model response correct? " + "First reason through each rule in the criteria step by step, then give a boolean verdict" ) _PREFERENCE_TEMPLATE = ( @@ -65,7 +68,8 @@ "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." "\n\nQuestion: {question}\n\nRubric: {answer}\n\nModel Response: {response}" - "\n\nIs the model response correct? Answer yes or no only." + "\n\nIs the model response correct? " + "First reason through each rule in the criteria step by step, then give a boolean verdict" ) _ABSTENTION_TEMPLATE = ( @@ -74,7 +78,8 @@ "The model could say that the information is incomplete, or some other information is " "given but the asked information is not." "\n\nQuestion: {question}\n\nExplanation: {answer}\n\nModel Response: {response}" - "\n\nDoes the model correctly identify the question as unanswerable? Answer yes or no only." + "\n\nDoes the model correctly identify the question as unanswerable? " + "First reason through each rule in the criteria step by step, then give a boolean verdict" ) _TEMPLATE_MAP = { @@ -121,7 +126,7 @@ def _build_prompt( class LongMemEvalJudgment(dspy.Signature): """Judge whether a model response is correct given type-specific evaluation criteria. - First reason through each rule in the criteria step by step, then answer yes or no. + First reason through each rule in the criteria step by step, then give a boolean verdict. """ evaluation_prompt: str = dspy.InputField( @@ -130,8 +135,8 @@ class LongMemEvalJudgment(dspy.Signature): reasoning: str = dspy.OutputField( description="Go through each rule in the evaluation criteria one by one. For each rule, determine whether the response satisfies it and why. Then give the overall reason for your final verdict." ) - judgment: str = dspy.OutputField( - description="Answer yes or no only." + judgment: bool = dspy.OutputField( + description="True if the model response is correct, False otherwise." ) @@ -178,9 +183,8 @@ async def _call_judge_async(self, prompt: str) -> tuple[bool, str, int, int]: """Async version of _call_judge using dspy.Predict.acall().""" with dspy.context(lm=self.lm): response = await self.judge.acall(evaluation_prompt=prompt) - content = response.judgment or "" + aligned = bool(response.judgment) reasoning = response.reasoning or "" - aligned = "yes" in content.strip().lower() input_tokens = 0 output_tokens = 0 @@ -202,9 +206,8 @@ def _call_judge(self, prompt: str) -> tuple[bool, str, int, int]: """ with dspy.context(lm=self.lm): response = self.judge(evaluation_prompt=prompt) - content = response.judgment or "" + aligned = bool(response.judgment) reasoning = response.reasoning or "" - aligned = "yes" in content.strip().lower() input_tokens = 0 output_tokens = 0 diff --git a/tests/adapters/test_longmemeval_judge.py b/tests/adapters/test_longmemeval_judge.py index 957ad1b..b86c1fd 100644 --- a/tests/adapters/test_longmemeval_judge.py +++ b/tests/adapters/test_longmemeval_judge.py @@ -51,8 +51,8 @@ def test_prompt_contains_question_and_answer(self): assert self.A in prompt assert self.R in prompt - def test_all_prompts_end_with_yes_or_no(self): + def test_all_prompts_end_with_question(self): for qtype in ("single-session-user", "temporal-reasoning", "knowledge-update", "single-session-preference"): prompt = _build_prompt(qtype, self.Q, self.A, self.R) - assert prompt.strip().endswith("Answer yes or no only.") + assert prompt.strip().endswith("?") From 6d25a5cfcfe41b85d1c097ab1e336948c757e4fc Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Mon, 6 Jul 2026 17:05:22 +0100 Subject: [PATCH 20/20] Prompt changes --- .../internal/adapters/agents/engram_dspy_agent.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 9c107ce..aaf9944 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -50,14 +50,9 @@ class AnswerUserQueryWithMemory(dspy.Signature): BOTH types: - Synthesizing facts across multiple memories is expected and encouraged. - - BEFORE answering, check the question's premises against the memories (see premise_check). - - If the memories contain NO relevant information at all, say "The information provided is not enough." and explain why fully. - - Contradiction vs. gap: - - A CONTRADICTION is when the question assumes a specific fact (role title, name, date, event, location) and the memories state a DIFFERENT fact. - - A GAP is when the memories are simply silent on a topic — they neither confirm nor deny it. - - If you find a contradiction: you MUST abstain. Say "The information provided is not enough." then explain what the question assumed vs. what the memories actually say. - - If you find only gaps: answer with whatever relevant information IS available. Do NOT abstain just because the answer is incomplete. + - BEFORE answering, identify which memories are relevant to answering the question. Then, reason through these instructions, making sure to follow them exactly, and ensure that you have done all required temporal reasoning. Finally, decide whether you have the information you need to answer the question or should abstain, again making sure to follow the instructions about this decision. + - If you have partial information about a question, use the reasoning field to make the best possible deduction you can from the memories. You should not abstain from answering if you have partial information, as you can still provide a best-effort answer. + - You should only abstain if the memories contain NO relevant information at all. In this case, say "The information provided is not enough." and explain why fully. Information is only not relevant if the question assumes a specific fact (role title, name, date, event, location) and the memories state a DIFFERENT fact. Temporal reasoning: - The question was asked on the date provided in `question_date`. Treat that as "today" when interpreting any time-relative terms ("now", "yesterday", "this year", etc.). @@ -74,8 +69,8 @@ class AnswerUserQueryWithMemory(dspy.Signature): desc="The date on which this question was asked. Treat this as today's date when interpreting any time-relative language in the question or memories." ) retrieved_memories: list[str] = dspy.InputField() - premise_check: str = dspy.OutputField( - desc="List the key premises embedded in the question (role titles, names, dates, events, quantities). For each, state whether the memories SUPPORT it, CONTRADICT it (memories state a different fact), or are SILENT (memories don't mention it). Only mark CONTRADICT when the memories provide a conflicting fact — silence is not contradiction. If all premises are supported or the memories are simply silent, say 'Premises verified.'" + reasoning: str = dspy.OutputField( + desc="First, identify which memories are relevant to answering the question. Then, reason through the instructions above, making sure to follow them exactly, and ensure that you have done all required temporal reasoning. Finally, decide whether you have the information you need to answer the question or should abstain, again making sure to follow the instructions about this decision above." ) answer: str = dspy.OutputField( desc="If the premise check found a CONTRADICTION (memories state a different fact than the question assumes), abstain: say 'The information provided is not enough.' and explain the discrepancy. For factual questions, answer strictly from the memories. For preference/recommendation questions, synthesize the user's stored preferences, habits, and interests into a personalized suggestion — this is expected, not fabrication."