#!/usr/bin/env python """ Letta text-based retrieval MCQ evaluation. Builds one Letta agent per movie, loads pre-computed text about the movie (captions or world-state histories) into the agent's archival memory, then asks each MCQ one at a time. The agent calls `archival_memory_search` to retrieve the most relevant passages before answering. Public runner for the Letta text-retrieval condition. ## Requires (on the server running this script) - `pip install letta letta-client` - A running Letta server (`letta server`, default `http://localhost:8283`) - For Gemini backbone: `gcloud auth application-default login` - For `gpt` / `claude`: OPENAI_API_KEY / ANTHROPIC_API_KEY - For `qwen3vl` / `glm4v`: a vLLM OpenAI-compatible server (see `LETTA_LOCAL_LLM_URL` env var, default `http://localhost:8001/v1`). - Embeddings: `text-embedding-3-small` from OpenAI (requires OPENAI_API_KEY even when the agent's LLM isn't GPT). ## Usage python run.py --backbone gemini --source captions --caption-variant gemini_frames python run.py --backbone gpt --source wsh --caption-variant gemini_video_clip_shot_grouped """ from __future__ import annotations import argparse import os import re import sys from pathlib import Path _PROJECT = Path(__file__).resolve().parents[3] sys.path.insert(0, str(_PROJECT)) from util.benchmark import load_benchmark from util.config import ( ANTHROPIC_API_KEY_ENV, DEFAULT_LETTA_BASE_URL, DEFAULT_LETTA_LOCAL_LLM_API_KEY, DEFAULT_LETTA_LOCAL_LLM_URL, DEFAULT_MODELS, LETTA_API_KEY_ENV, LETTA_BASE_URL_ENV, LETTA_LOCAL_LLM_API_KEY_ENV, LETTA_LOCAL_LLM_URL_ENV, OPENAI_API_KEY_ENV, ) from util.eval_loop import run_sweep from util.parsers import parse_mcq_any from util.paths import ( BENCHMARK_PATH, CAPTIONS_DIR, RESULTS_DIR, WSH_DIR, ensure_dir, results_filename, ) EXPERIMENT_NAME = 'letta' SUPPORTED_BACKBONES = ('gemini', 'gpt', 'claude', 'qwen3vl', 'glm4v') SOURCES = ('captions', 'wsh') LETTA_PASSAGE_TOKEN_BUDGET = 3500 LETTA_PASSAGE_CHAR_BUDGET = 14_000 # Per-backbone Letta-specific overrides. None = use util.config.DEFAULT_MODELS # verbatim. Letta's agentic loop is multi-call (search + answer + tool reasoning), # so picking the cheaper / higher-rate-limit sibling matters; Haiku 4.5 is kept # as an explicit pin because the global default could conceivably move to Sonnet # (Sonnet's tier-1 30K TPM gets eaten alive by Letta's archival_search churn). _LETTA_DEFAULT_MODELS = { 'gemini': None, # use DEFAULT_MODELS['gemini'] (gemini-3-flash-preview) 'gpt': None, # use DEFAULT_MODELS['gpt'] (gpt-5-mini) 'claude': DEFAULT_MODELS['claude'], 'qwen3vl': None, 'glm4v': None, } def main() -> None: ap = argparse.ArgumentParser() ap.add_argument('--backbone', default='gemini', choices=SUPPORTED_BACKBONES) ap.add_argument('--model', default=None, help='Override the LLM the agent uses.') ap.add_argument('--source', default='captions', choices=SOURCES, help="'captions' loads per-chunk text from data/captions/{variant}/{vid}/; " "'wsh' loads the concatenated world-state-history cache instead.") ap.add_argument('--caption-variant', default='gemini_frames', help='Captioner dir name under data/captions/; also used for the WSH cache.') ap.add_argument('--benchmark', default=str(BENCHMARK_PATH)) ap.add_argument('--letta-base-url', default=os.environ.get(LETTA_BASE_URL_ENV, DEFAULT_LETTA_BASE_URL)) ap.add_argument('--letta-api-key', default=os.environ.get(LETTA_API_KEY_ENV)) ap.add_argument('--checkpoint-every', type=int, default=25) ap.add_argument('--refresh-passages', action='store_true', help='When an agent already exists in postgres, delete its ' 'archival passages and re-ingest from the current ' 'captions on disk. Use this after regenerating ' 'captions / WSH —- otherwise the agent answers based ' 'on the stale passages baked into its memory.') ap.add_argument('--refresh-llm', action='store_true', help="When an agent already exists, modify its llm_config " "to the current --model / Letta default. Letta bakes " "the LLM choice in at agent-creation time, so " "switching DEFAULT_MODELS or passing --model on " "re-run does NOT take effect for existing agents " "without this flag. Use this when you want to keep " "the agent's archival memory intact but swap models. ") args = ap.parse_args() # Reject obviously-malformed --caption-variant values. Without this, a # bash quoting accident like --caption-variant "gemini_frames --refresh-passages" # silently turns the second flag into the first flag's value, and you # end up looking for captions in `data/captions/gemini_frames --refresh-passages/`. # Fail loud at startup instead. if any(c.isspace() or c in '-=' for c in args.caption_variant[:1]) \ or any(c.isspace() for c in args.caption_variant): sys.exit( f'ERROR: --caption-variant contains whitespace or starts with -/= ' f'({args.caption_variant!r}).\n' f' Looks like a flag got swallowed into its value (a bash ' f'quoting accident). Each flag should be its own shell word:\n' f' bash run.sh --backbone gpt --source captions \\\n' f' --caption-variant gemini_frames --refresh-passages --refresh-llm\n' f' NOT --caption-variant "gemini_frames --refresh-passages" ' f'(that is one value, not two flags).' ) # Resolve effective model id (--model > backbone-specific Letta default > # global DEFAULT_MODELS). if args.model is None: args.model = _LETTA_DEFAULT_MODELS.get(args.backbone) # Hard import (lets the pure-argparse errors above fire before Letta is required) try: from letta_client import Letta except ImportError as e: sys.exit(f'letta_client not installed. Run `pip install letta letta-client`.\n' f' ({type(e).__name__}: {e})') _env_preflight(args.backbone) client = Letta(base_url=args.letta_base_url, token=args.letta_api_key) if args.letta_api_key \ else Letta(base_url=args.letta_base_url) _ping(client, args.letta_base_url) out_dir = ensure_dir(RESULTS_DIR / EXPERIMENT_NAME) df = load_benchmark(args.benchmark) print(f'Loaded {len(df)} questions across {df["video_file_id"].nunique()} movies') llm_config = _llm_config_for(args.backbone, model_override=args.model) print(f'Letta agent LLM: {llm_config["model"]} ({args.backbone})') emb_config = _embedding_config() extra_tool_ids = _resolve_tools(client, ['archival_memory_search', 'archival_memory_insert']) _run_one(df, args, out_dir, client, llm_config, emb_config, extra_tool_ids) def _env_preflight(backbone: str) -> None: from util.secrets import require, load_secrets load_secrets() if not os.environ.get(OPENAI_API_KEY_ENV): print(f' ⚠ {OPENAI_API_KEY_ENV} not set -- Letta uses OpenAI for embeddings even with a non-GPT agent.') if backbone == 'claude': require(ANTHROPIC_API_KEY_ENV) def _ping(client, base_url: str) -> None: """Confirm the Letta server is reachable before we start creating agents.""" for attempt in ( lambda: client.health.check(), lambda: client.health(), lambda: client.agents.list(limit=1), ): try: attempt() print(f'Letta server reachable at {base_url}') return except Exception: continue sys.exit(f'Letta server unreachable at {base_url}. Start it: `letta server`.') def _llm_config_for(backbone: str, model_override: Optional[str] = None) -> dict: """Build the LLM config dict sent to the Letta server.""" pick = lambda default: model_override or default if backbone == 'gemini': return { 'model': pick(DEFAULT_MODELS['gemini']), 'model_endpoint_type': 'google_vertex', 'model_endpoint': 'https://aiplatform.googleapis.com', 'context_window': 32768, } if backbone in ('qwen3vl', 'glm4v'): url = os.environ.get(LETTA_LOCAL_LLM_URL_ENV, DEFAULT_LETTA_LOCAL_LLM_URL) return { 'model': pick(DEFAULT_MODELS[backbone]), 'model_endpoint_type': 'openai', 'model_endpoint': url, 'context_window': 32768, # vLLM usually needs no key, but pass through if set for auth'd setups 'api_key': os.environ.get(LETTA_LOCAL_LLM_API_KEY_ENV) or DEFAULT_LETTA_LOCAL_LLM_API_KEY, } if backbone == 'gpt': return { 'model': pick(DEFAULT_MODELS['gpt']), 'model_endpoint_type': 'openai', 'model_endpoint': 'https://api.openai.com/v1', 'context_window': 128000, 'api_key': os.environ.get(OPENAI_API_KEY_ENV), } if backbone == 'claude': return { 'model': pick(DEFAULT_MODELS['claude']), 'model_endpoint_type': 'anthropic', 'model_endpoint': 'https://api.anthropic.com/v1', 'context_window': 200000, 'api_key': os.environ.get(ANTHROPIC_API_KEY_ENV), } raise ValueError(f'Unknown backbone: {backbone}') def _embedding_config() -> dict: return { 'embedding_model': 'text-embedding-3-small', 'embedding_endpoint_type': 'openai', 'embedding_endpoint': 'https://api.openai.com/v1', 'embedding_dim': 1536, 'embedding_chunk_size': 300, # Letta's server needs OPENAI_API_KEY for embedding calls. 'api_key': os.environ.get(OPENAI_API_KEY_ENV), } def _unwrap_page(page): for attr in ('items', 'data', 'results', 'agents'): v = getattr(page, attr, None) if v is not None: return list(v) try: return list(page) except TypeError: return [] def _resolve_tools(client, names: list[str]) -> list[str]: tools = _unwrap_page(client.tools.list(limit=200)) by_name = {t.name: t.id for t in tools} resolved = [by_name[n] for n in names if n in by_name] missing = [n for n in names if n not in by_name] if missing: print(f' ⚠ tools not found on Letta server: {missing}') else: print(f'Resolved Letta tools: {names}') return resolved # --------------------------------------------------------------------------- def _safe_agent_name(name: str) -> str: return re.sub(r"[^\w\s'\-]", '_', name) def _load_passages(vid: str, source: str, caption_variant: str) -> list[tuple[int, str]]: if source == 'captions': d = CAPTIONS_DIR / caption_variant / vid else: # 'wsh' wsh_file = WSH_DIR / caption_variant / f'{vid}.txt' if not wsh_file.exists(): return [] raw = wsh_file.read_text() blocks = re.split(r'(?=^\*\*Timestamp\*\*)', raw, flags=re.MULTILINE) return [(i, b.strip()) for i, b in enumerate(blocks) if b.strip()] if not d.exists(): return [] items: list[tuple[int, str]] = [] for p in sorted(d.glob('*.txt'), key=lambda x: int(x.stem) if x.stem.isdigit() else -1): if not p.stem.isdigit(): continue text = p.read_text().strip() if text: items.append((int(p.stem), text)) return items def _rough_token_count(text: str) -> int: # Conservative enough for Letta's 8192-token per-passage cap without # adding a hard tokenizer dependency to the experiment runner. return max(len(re.findall(r'\w+|[^\w\s]', text)), len(text) // 4) def _fits_letta_passage(text: str) -> bool: return ( len(text) <= LETTA_PASSAGE_CHAR_BUDGET and _rough_token_count(text) <= LETTA_PASSAGE_TOKEN_BUDGET ) def _split_long_unit(text: str) -> list[str]: words = re.findall(r'\S+\s*', text) out: list[str] = [] cur = '' for word in words: if not cur and not _fits_letta_passage(word): out.extend( word[i:i + LETTA_PASSAGE_CHAR_BUDGET].strip() for i in range(0, len(word), LETTA_PASSAGE_CHAR_BUDGET) if word[i:i + LETTA_PASSAGE_CHAR_BUDGET].strip() ) continue candidate = cur + word if cur and not _fits_letta_passage(candidate): out.append(cur.strip()) cur = word else: cur = candidate if cur.strip(): out.append(cur.strip()) return out def _split_passage_for_letta(text: str) -> list[str]: """Split one source caption/WSH chunk into Letta-safe archival passages.""" text = text.strip() if not text: return [] if _fits_letta_passage(text): return [text] units = [u.strip() for u in re.split(r'(?<=[.!?])\s+|\n{2,}', text) if u.strip()] out: list[str] = [] cur = '' for unit in units: if not _fits_letta_passage(unit): if cur: out.append(cur.strip()) cur = '' out.extend(_split_long_unit(unit)) continue candidate = f'{cur}\n\n{unit}' if cur else unit if cur and not _fits_letta_passage(candidate): out.append(cur.strip()) cur = unit else: cur = candidate if cur.strip(): out.append(cur.strip()) return out def _iter_letta_passages(passages: list[tuple[int, str]]): for chunk_idx, text in passages: parts = _split_passage_for_letta(text) if len(parts) <= 1: if parts: yield str(chunk_idx), f'chunk_{chunk_idx}', parts[0] continue for part_idx, part in enumerate(parts, start=1): yield ( f'{chunk_idx}.{part_idx}/{len(parts)}', f'chunk_{chunk_idx}_part_{part_idx}', part, ) def _create_passages(client, agent_id: str, vid: str, passages: list[tuple[int, str]]) -> int: created = 0 created_refs = [] try: for chunk_label, tag, text in _iter_letta_passages(passages): passage = client.agents.passages.create( agent_id=agent_id, text=f'[chunk {chunk_label}] {text}', tags=[vid, tag], ) created_refs.append(passage) created += 1 except Exception: for passage in created_refs: try: _delete_passage(client, agent_id, passage) except Exception: pass raise return created def _delete_passage(client, agent_id: str, passage) -> None: pid = getattr(passage, 'id', None) if not pid: return try: client.agents.passages.delete(agent_id=agent_id, memory_id=pid) except TypeError: client.agents.passages.delete(agent_id=agent_id, id=pid) def _get_or_create_agent( client, vid: str, movie_title: str, args, llm_config: dict, emb_config: dict, extra_tool_ids: list[str], ): agent_prefix = f'chnb-{args.backbone}-{args.source}-{args.caption_variant}' agent_name = _safe_agent_name(f'{agent_prefix}-{vid}') existing = _unwrap_page(client.agents.list(name=agent_name, limit=1)) if existing: agent = existing[0] # Top up tools on older agents that predate the required-tool list try: for tid in extra_tool_ids: client.agents.tools.attach(agent_id=agent.id, tool_id=tid) except Exception: pass # Optional: re-ingest captions when --refresh-passages is set. Without # this, the agent keeps the passages it was created with even after # captions on disk have been regenerated. We create the replacement # passages first, then delete the old IDs, so a failed refresh does not # leave the agent with no archival memory. if getattr(args, 'refresh_passages', False): try: fresh = _load_passages(vid, args.source, args.caption_variant) if not fresh: print(f' ⚠ refresh_passages skipped for {agent_name}: ' f'no current passages found on disk') return agent old = _unwrap_page(client.agents.passages.list(agent_id=agent.id, limit=10_000)) created = _create_passages(client, agent.id, vid, fresh) for p in old: _delete_passage(client, agent.id, p) print(f' refreshed {created} passages for {agent_name} ' f'from {len(fresh)} source chunks (was {len(old)})') except Exception as e: print(f' ⚠ refresh_passages failed for {agent_name}: ' f'{type(e).__name__}: {e!r}') # Optional: swap the agent's llm_config to the current --model / # backbone default. Letta bakes llm_config in at agent creation time, # so changing DEFAULT_MODELS / _LETTA_DEFAULT_MODELS or passing # --model on re-run is a no-op for existing agents without this. if getattr(args, 'refresh_llm', False): try: old_model = None old_cfg = getattr(agent, 'llm_config', None) if old_cfg is not None: old_model = getattr(old_cfg, 'model', None) or ( old_cfg.get('model') if isinstance(old_cfg, dict) else None) # Letta SDK has used both `update` and `modify` across # versions for agent edits. Try the newer `modify` first. try: client.agents.modify(agent_id=agent.id, llm_config=llm_config) except AttributeError: client.agents.update(agent_id=agent.id, llm_config=llm_config) # Re-fetch so the caller sees the updated config (some SDKs # don't mutate the local object in place). try: agent = client.agents.retrieve(agent_id=agent.id) except AttributeError: agent = client.agents.get(agent_id=agent.id) print(f' refreshed llm for {agent_name}: ' f'{old_model or "?"} -> {llm_config["model"]}') except Exception as e: print(f' ⚠ refresh_llm failed for {agent_name}: ' f'{type(e).__name__}: {e!r}') return agent passages = _load_passages(vid, args.source, args.caption_variant) if not passages: return None agent = client.agents.create( name=agent_name, llm_config=llm_config, embedding_config=emb_config, memory_blocks=[ {'label': 'persona', 'value': (f'You answer multiple-choice questions about the movie "{movie_title}". ' 'Use archival_memory_search to find relevant scenes before answering. ' 'Output ONLY valid JSON as instructed.')}, {'label': 'human', 'value': 'User is a researcher evaluating MCQ performance on hour-long movies.'}, ], include_base_tools=True, tool_ids=extra_tool_ids, ) _create_passages(client, agent.id, vid, passages) return agent # One-shot diagnostic: dump message types the FIRST time a movie returns # empty, so we can tell whether Letta is producing assistant_message frames # at all (vs. only tool calls / reasoning) without spamming logs forever. _DIAGNOSED_AGENTS: set = set() def _extract_text(content) -> str: """Extract any text from a Letta message's `content` field. Letta uses str / list-of-blocks / list-of-dicts depending on version + message type.""" if isinstance(content, str): return content if isinstance(content, list): parts = [] for block in content: t = getattr(block, 'text', None) or \ (block.get('text') if isinstance(block, dict) else None) if t: parts.append(t) if parts: return ''.join(parts) return '' def _extract_send_message_arg(m) -> str: """If `m` is a tool_call_message whose tool is `send_message`, return its `message` argument.""" tcs = getattr(m, 'tool_calls', None) or getattr(m, 'tool_call', None) if tcs is None and isinstance(m, dict): tcs = m.get('tool_calls') or m.get('tool_call') if tcs is None: return '' if not isinstance(tcs, list): tcs = [tcs] import json as _json for tc in tcs: name = getattr(tc, 'name', None) or \ (tc.get('name') if isinstance(tc, dict) else None) # Some SDKs nest under .function.name if not name: fn = getattr(tc, 'function', None) or \ (tc.get('function') if isinstance(tc, dict) else None) if fn is not None: name = getattr(fn, 'name', None) or \ (fn.get('name') if isinstance(fn, dict) else None) if name != 'send_message': continue args = getattr(tc, 'arguments', None) or \ (tc.get('arguments') if isinstance(tc, dict) else None) if args is None: fn = getattr(tc, 'function', None) or \ (tc.get('function') if isinstance(tc, dict) else None) if fn is not None: args = getattr(fn, 'arguments', None) or \ (fn.get('arguments') if isinstance(fn, dict) else None) if isinstance(args, str): try: args = _json.loads(args) except Exception: return args # raw JSON string is still parseable downstream if isinstance(args, dict): return str(args.get('message', '') or args.get('content', '') or '') return '' def _ask_agent(client, agent_id: str, prompt: str) -> str: response = client.agents.messages.create( agent_id=agent_id, messages=[{'role': 'user', 'content': prompt}], ) msgs = list(getattr(response, 'messages', None) or _unwrap_page(response)) # 1. Canonical for m in reversed(msgs): if getattr(m, 'message_type', None) == 'assistant_message': text = _extract_text(getattr(m, 'content', None)) if text: return text # 2. Fallback for m in reversed(msgs): if getattr(m, 'message_type', None) == 'tool_call_message': arg = _extract_send_message_arg(m) if arg: return arg # 3. Last-ditch for m in reversed(msgs): if getattr(m, 'message_type', None) == 'reasoning_message': text = _extract_text(getattr(m, 'content', None) or getattr(m, 'reasoning', None)) if text: return text # Diagnostic: first empty response per agent dumps the message-type # histogram so we can see whether assistant_message frames are even # being produced. Subsequent empties stay quiet to avoid log spam. if agent_id not in _DIAGNOSED_AGENTS: _DIAGNOSED_AGENTS.add(agent_id) types = [getattr(m, 'message_type', '?') for m in msgs] print(f' [empty-debug] agent={agent_id[:12]}… msg types: {types}') # Show the last 1-2 messages' attributes so we can see what's there for m in msgs[-2:]: mt = getattr(m, 'message_type', '?') preview = _extract_text(getattr(m, 'content', None))[:200] tcs = getattr(m, 'tool_calls', None) or getattr(m, 'tool_call', None) print(f' [{mt}] content_preview={preview!r} has_tool_calls={bool(tcs)}') return '' def _build_prompt(row: dict) -> str: return ( 'Use your archival_memory_search tool to find relevant passages about the movie, ' 'then answer the question below.\n\n' f"Question: {row['question']} Select the single best choice.\n" f"A: {row['option_A']}\n" f"B: {row['option_B']}\n" f"C: {row['option_C']}\n" f"D: {row['option_D']}\n\n" 'Rules:\n' '- You MUST pick exactly one of A/B/C/D. If retrieval comes back empty ' 'or inconclusive, commit to your best guess from the partial evidence.\n' '- Do NOT refuse, do NOT ask for clarification, do NOT include commentary ' 'in the final answer -— output ONLY the JSON below.\n\n' 'Output ONLY valid JSON:\n' f' {{"id": {row["qid"]}, "answer": "X"}}' ) # --------------------------------------------------------------------------- def _run_one(df, args, out_dir, client, llm_config, emb_config, extra_tool_ids) -> None: def build_movie_ctx(vid, movie_df): title = movie_df.iloc[0]['movie_title'] try: agent = _get_or_create_agent( client, vid, title, args, llm_config, emb_config, extra_tool_ids, ) except Exception as e: print(f' [agent create] {vid}: {type(e).__name__}: {e!r}') return {} if agent is None: print(f' SKIP {vid}: no passages under ' f'data/captions/{args.caption_variant}/{vid}/ (or wsh cache missing)') return {} return {'agent_id': agent.id} def call_fn(row, ctx): agent_id = ctx.get('agent_id') if not agent_id: return '' try: raw = _ask_agent(client, agent_id, _build_prompt(row)) except Exception as e: print(f' [ask] {row["qid"]}: {type(e).__name__}: {e!r}') return '' return parse_mcq_any(raw, qid=row['qid']) # Filename: `results_{source}_{caption_variant}_{backbone}.csv`. # Closed-source model defaults are pinned in util.config.DEFAULT_MODELS, # so the model id doesn't need to appear in the filename. model_id = llm_config['model'] results_path = out_dir / f'results_{args.source}_{args.caption_variant}_{args.backbone}.csv' run_sweep( df, call_fn=call_fn, build_movie_ctx=build_movie_ctx, results_path=results_path, experiment_name=f'{EXPERIMENT_NAME} {args.source}/{args.caption_variant}x{args.backbone}', extra_row_fields={ 'source': args.source, 'caption_variant': args.caption_variant, 'backbone': args.backbone, 'model': model_id, }, checkpoint_every=args.checkpoint_every, ) if __name__ == '__main__': main()