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/agents/engram_dspy_agent.py b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py index 4884e76..aaf9944 100644 --- a/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py +++ b/query_agent_benchmarking/internal/adapters/agents/engram_dspy_agent.py @@ -11,9 +11,9 @@ from dataclasses import dataclass, field import dspy -from pydantic import BaseModel from engram import ( + AsyncEngramClient, EngramClient, BM25Retrieval, FetchRetrieval, @@ -33,52 +33,44 @@ # 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. 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." - - 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 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.). - 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. - - 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() - 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.'" + 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[str] = dspy.InputField() + 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." @@ -125,15 +117,22 @@ 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( - 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 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() @@ -147,6 +146,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. @@ -174,18 +174,71 @@ def run( user_id=user_id, group=self.engram_group, retrieval_config=retrieval_cls(limit=self.retrieval_limit), + topics=self.search_topics, ) - retrieved = [ - MemoryWithTimestamp( - memory=m.content, - time_added=str(m.created_at), - ) - for m in memories - ] + 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, + question_date=question_date or "unknown", + retrieved_memories=retrieved, + ) + + return EngramAskResponse( + final_answer=response.answer, + raw_response={ + "n_memories_retrieved": len(retrieved), + "memories": [ + {"memory": m.content, "time_added": str(m.updated_at)} + for m in sorted_memories + ], + }, + ) + + async def run_async( + self, + 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. + + 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), + topics=self.search_topics, + ) + + 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, + question_date=question_date or "unknown", retrieved_memories=retrieved, ) @@ -193,6 +246,15 @@ 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 + ], }, ) + + async def initialize_async(self) -> None: + pass + + async def close_async(self) -> None: + pass 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/database/engram_loader.py b/query_agent_benchmarking/internal/adapters/database/engram_loader.py index 0088bd2..9116960 100644 --- a/query_agent_benchmarking/internal/adapters/database/engram_loader.py +++ b/query_agent_benchmarking/internal/adapters/database/engram_loader.py @@ -5,12 +5,28 @@ sessions into Engram's memory system on a per-tenant basis. """ +import asyncio import os import time -from typing import Optional, Callable from dataclasses import dataclass, field +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Callable, Optional -from engram import EngramClient, ConversationInput, MessageInput +from engram import AsyncEngramClient, ConversationInput, EngramClient, MessageInput + +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: @@ -97,26 +113,69 @@ class TenantIngestionStats: run_durations: list[float] = field(default_factory=list) -def _submit_all( - client: EngramClient, +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], updated_at=conversation.updated_at) + 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, updated_at=conversation.updated_at)) + current = [] + if current: + turns.append(ConversationInput(messages=current, updated_at=conversation.updated_at)) + return turns + + raise ValueError( + f"Unsupported ingestion_mode '{ingestion_mode}'. " + f"Supported: conversation, user_messages, message_turn" + ) + + +@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: ``"conversation"`` submits the full parsed conversation - per session. ``"user_messages"`` submits each user message - individually as a single-message conversation (no assistant context). + 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()): @@ -125,6 +184,7 @@ def _submit_all( 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: @@ -132,78 +192,97 @@ 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 % 50 == 0: - print(f" Submitted {submitted} user messages") - - if ingest_delay > 0: - time.sleep(ingest_delay) - else: - # Default: submit the full conversation - t_submit = time.time() - try: - run = client.memories.add( - conversation, - 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} for tenant {tenant_id}: {e}") - continue - records.append(RunRecord( - run_id=run.run_id, + 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 + + items = per_user.setdefault(user_id, []) + for input_idx, conv_input in enumerate(inputs): + 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, + include_conversation_id: bool = True, +) -> 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. + 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 + ) + + 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} if include_conversation_id else {}, + ) + 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 @@ -213,27 +292,27 @@ def _submit_all( "submitted": submitted, "skipped": skipped, "total": total, - "tenant_id": tenant_id, - "session_id": session.get("session_id", ""), + "tenant_id": item.tenant_id, + "session_id": item.session_id, + "input_index": item.input_idx, }) - if verbose and submitted % 50 == 0: - print(f" Submitted {submitted}/{total}") + 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: - 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(per_user)} 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 def _poll_and_collect( - client: EngramClient, + client, records: list[RunRecord], poll_interval: float, verbose: bool, @@ -331,7 +410,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 @@ -353,20 +432,24 @@ 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, 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. - 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. @@ -374,37 +457,62 @@ 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. - 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. + 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. 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, - ) + api_key = engram_api_key or (None if dry_run else os.environ["ENGRAM_API_KEY"]) + + if dry_run: + submit_client = _DryRunEngramClient() + ingest_delay = 0.0 + poll = False + else: + 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_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, include_conversation_id)) submit_elapsed = time.time() - t0 tenant_session_counts = {} 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, @@ -419,10 +527,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 @@ -440,3 +549,22 @@ def engram_ingest_all_tenants( ) return result + + +class _DryRunEngramClient: + """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 + + async def add(self, *args, **kwargs): + self._count += 1 + return SimpleNamespace(run_id=f"dry-run-{self._count}") + + async def get(self, *args, **kwargs): + return SimpleNamespace( + status="completed", + committed_operations=SimpleNamespace(created=[], updated=[], deleted=[]), + ) 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/adapters/metrics/ask_metrics_calculator.py b/query_agent_benchmarking/internal/adapters/metrics/ask_metrics_calculator.py index dc02635..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,15 +215,19 @@ 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, max_concurrent_judge: int = 10): self.model = model - self.judge = LongMemEvalJudge(model=model, api_key=api_key) + 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: 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[list[dict] | None] = [] query_times = [] misaligned_indices = [] total_input_tokens = 0 @@ -231,6 +238,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 +252,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,8 +284,81 @@ 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": 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 + + 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": 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 506c7ca..cfdb27d 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 @@ -28,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 = ( @@ -41,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 = ( @@ -55,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 = ( @@ -64,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 = ( @@ -73,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 = { @@ -120,14 +126,17 @@ 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 give a boolean verdict. """ evaluation_prompt: str = dspy.InputField( description="The full evaluation prompt with type-specific instructions, question, correct answer, and model response." ) - judgment: str = dspy.OutputField( - description="Answer yes or no only." + 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: bool = dspy.OutputField( + description="True if the model response is correct, False otherwise." ) @@ -136,8 +145,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. """ @@ -148,6 +157,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. @@ -159,6 +169,7 @@ def __init__( """ self.model = model self.default_question_type = default_question_type + self.ensemble_k = ensemble_k self.lm = dspy.LM( model, @@ -168,16 +179,35 @@ def __init__( self.judge = dspy.Predict(LongMemEvalJudgment) - def _call_judge(self, prompt: str) -> tuple[bool, int, int]: + 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) + aligned = bool(response.judgment) + reasoning = response.reasoning or "" + + 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. 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 "" - aligned = "yes" in content.strip().lower() + aligned = bool(response.judgment) + reasoning = response.reasoning or "" input_tokens = 0 output_tokens = 0 @@ -189,7 +219,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,8 +247,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, @@ -245,15 +275,56 @@ 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) + 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": majority_aligned, + "question_type": qtype, + "is_abstention": is_abstention, + "input_tokens": total_input_tokens, + "output_tokens": total_output_tokens, + "votes": votes, + "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": aligned, + "aligned": votes >= self.ensemble_k / 2, "question_type": qtype, "is_abstention": is_abstention, - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "votes": 1, - "ensemble_k": 1, - "reasoning": None, # LongMemEval judge uses no chain-of-thought + "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/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 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/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/core/services/ask_benchmark.py b/query_agent_benchmarking/internal/core/services/ask_benchmark.py index 6ea68dd..788cdc1 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"), @@ -151,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) @@ -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) + 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%}") diff --git a/query_agent_benchmarking/internal/core/services/populate_db.py b/query_agent_benchmarking/internal/core/services/populate_db.py index a6c4ee1..b3b614b 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: @@ -72,8 +74,22 @@ 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) + include_conversation_id = config.get("engram_include_conversation_id", True) + result = engram_ingest_all_tenants( + docs_by_tenant, + engram_base_url=engram_base_url, + poll=poll, + 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 {} manifest = save_engram_manifest(result, dataset_name) return manifest 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/scripts/poll_engram.py b/scripts/poll_engram.py new file mode 100644 index 0000000..69d3b31 --- /dev/null +++ b/scripts/poll_engram.py @@ -0,0 +1,242 @@ +"""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 _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( + client: EngramClient, + tenant_runs: dict[str, list[dict]], + frontiers: dict[str, int], + interval: float, + show_tenants: bool = False, +) -> 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") + if show_tenants: + _print_user_lists(tenant_runs, frontiers) + + +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", + ) + 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 + 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, show_tenants=args.tenants) + else: + run_single_pass(client, tenant_runs, frontiers, show_tenants=args.tenants) + + +if __name__ == "__main__": + main() 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) 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) 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("?") 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")