Skip to content

reranker

zocrates.agents.common.reranker

Evidence reranking and budget trimming utilities for Zocrates research agents.

HARD_WORD_CEILING module-attribute

HARD_WORD_CEILING = 7500

MAX_RERANK_CHUNKS module-attribute

MAX_RERANK_CHUNKS = 15

logger module-attribute

logger = logging.getLogger(__name__)

_trim_to_budget

_trim_to_budget(chunks, word_ceiling=HARD_WORD_CEILING)

Trim lowest-ranked chunks so cumulative words <= word_ceiling.

Preserves at least the first chunk even if it alone exceeds word_ceiling.

Source code in zocrates/agents/common/reranker.py
def _trim_to_budget(
    chunks: list[dict[str, Any]], word_ceiling: int = HARD_WORD_CEILING
) -> list[dict[str, Any]]:
    """Trim lowest-ranked chunks so cumulative words <= word_ceiling.

    Preserves at least the first chunk even if it alone exceeds word_ceiling.
    """
    if not chunks:
        return []
    trimmed: list[dict[str, Any]] = []
    accumulated_words = 0
    for chunk in chunks:
        wc = chunk.get("word_count")
        if wc is None:
            wc = len(chunk.get("text", "").split())
        if trimmed and (accumulated_words + wc > word_ceiling):
            continue
        trimmed.append(chunk)
        accumulated_words += wc
    return trimmed

filter_substantive_evidence

filter_substantive_evidence(chunks)

Filter out non-substantive evidence chunks (frontmatter, TOC).

Fallback guard: If filtering yields 2 or fewer substantive chunks (regardless of whether any remain), returns the original chunks so the synthesizer is not starved of evidence. This preserves frontmatter chunks when there is not enough substantive content to stand on its own (e.g. the query only matched 1-2 real sections alongside a frontmatter hit).

Source code in zocrates/agents/common/reranker.py
def filter_substantive_evidence(
    chunks: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Filter out non-substantive evidence chunks (frontmatter, TOC).

    Fallback guard: If filtering yields 2 or fewer substantive chunks (regardless
    of whether any remain), returns the original chunks so the synthesizer is not
    starved of evidence.  This preserves frontmatter chunks when there is not
    enough substantive content to stand on its own (e.g. the query only matched
    1-2 real sections alongside a frontmatter hit).
    """
    if not chunks:
        return []

    substantive = [c for c in chunks if is_substantive_chunk(c)]
    if len(substantive) <= 2:
        if len(chunks) > len(substantive):
            logger.debug(
                "Substantive filter left only %d chunk(s) (<= 2); falling back to original %d chunks.",
                len(substantive),
                len(chunks),
            )
        return chunks
    return substantive

get_secure_api_key

get_secure_api_key(key_name='openrouter_api_key', env_vars=('OPENROUTER_API_KEY', 'OPEN_ROUTER_API_KEY'))

Load a named API key from in-memory session environment variables.

Zotero Preferences / Login Manager serves as the secure persistent vault. Backend keys are in-memory session variables (os.environ) synced from Zotero.

Parameters:

Name Type Description Default
key_name str

Entry identifier (e.g., 'openrouter_api_key', 'openai_api_key').

'openrouter_api_key'
env_vars list[str] | tuple[str, ...] | None

Environment variable names to check.

('OPENROUTER_API_KEY', 'OPEN_ROUTER_API_KEY')

Returns:

Type Description
str | None

The API key string or None if not found.

Source code in zocrates/config.py
def get_secure_api_key(
    key_name: str = "openrouter_api_key",
    env_vars: list[str] | tuple[str, ...] | None = ("OPENROUTER_API_KEY", "OPEN_ROUTER_API_KEY"),
) -> str | None:
    """Load a named API key from in-memory session environment variables.

    Zotero Preferences / Login Manager serves as the secure persistent vault.
    Backend keys are in-memory session variables (os.environ) synced from Zotero.

    Args:
        key_name: Entry identifier (e.g., 'openrouter_api_key', 'openai_api_key').
        env_vars: Environment variable names to check.

    Returns:
        The API key string or None if not found.
    """
    if env_vars:
        for var in env_vars:
            val = os.getenv(var)
            if val:
                return val
    return None

is_substantive_chunk

is_substantive_chunk(chunk)

Determine whether a chunk contains substantive content rather than frontmatter or TOC.

Returns False if:
  - chunk.get("chunk_type") == "frontmatter"
  - str(chunk.get("chunk_id", "")).endswith(":s0") or chunk.get("asset_id") == "s0"
  - Heading matches frontmatter or TOC: ('frontmatter', 'table of contents', 'contents', 'toc')
  - Text is pure YAML frontmatter (starts with '---', contains 'title:' or 'authors:', and has '

---') Otherwise returns True.

Source code in zocrates/agents/common/reranker.py
def is_substantive_chunk(chunk: dict[str, Any]) -> bool:
    """Determine whether a chunk contains substantive content rather than frontmatter or TOC.

    Returns False if:
      - chunk.get("chunk_type") == "frontmatter"
      - str(chunk.get("chunk_id", "")).endswith(":s0") or chunk.get("asset_id") == "s0"
      - Heading matches frontmatter or TOC: ('frontmatter', 'table of contents', 'contents', 'toc')
      - Text is pure YAML frontmatter (starts with '---', contains 'title:' or 'authors:', and has '\n---')
    Otherwise returns True.
    """
    if chunk.get("chunk_type") == "frontmatter":
        return False

    chunk_id = str(chunk.get("chunk_id") or "")
    if chunk_id.endswith(":s0") or chunk.get("asset_id") == "s0":
        return False

    heading = str(chunk.get("heading") or "").strip().lower()
    if heading in ("frontmatter", "table of contents", "contents", "toc"):
        return False

    text = str(chunk.get("text") or "")
    clean_text = text.strip()
    if clean_text.startswith(("---\n", "---\r\n")) and (
        "\ntitle:" in clean_text or "\nauthors:" in clean_text or clean_text.startswith("---\ntitle:") or clean_text.startswith("---\nauthors:")
    ) and "\n---" in clean_text:
        return False

    return True

load_zocrates_config

load_zocrates_config(zotero_data_dir=None)

Load the Zocrates config from the data directory.

If zotero_data_dir is None, attempts to obtain it from the active LibrarySyncer service or the ZOTERO_DATA_DIR environment variable.

Returns config merged with standard defaults.

Parameters:

Name Type Description Default
zotero_data_dir Path | str | None

Root Zotero data directory.

None

Returns:

Type Description
dict

Parsed JSON config dict merged with defaults.

Source code in zocrates/config.py
def load_zocrates_config(zotero_data_dir: Path | str | None = None) -> dict:
    """Load the Zocrates config from the data directory.

    If ``zotero_data_dir`` is ``None``, attempts to obtain it from the active
    ``LibrarySyncer`` service or the ``ZOTERO_DATA_DIR`` environment variable.

    Returns config merged with standard defaults.

    Args:
        zotero_data_dir: Root Zotero data directory.

    Returns:
        Parsed JSON config dict merged with defaults.
    """
    defaults = copy.deepcopy(DEFAULT_CONFIG)

    resolved_dir = resolve_zotero_data_dir(zotero_data_dir)
    if resolved_dir is None:
        return defaults

    config_path = get_config_path(resolved_dir)
    resolved_path = config_path.resolve()
    if not config_path.exists():
        with _CONFIG_CACHE_LOCK:
            _CONFIG_CACHE.pop(resolved_path, None)
            _CONFIG_CACHE_FINGERPRINT.pop(resolved_path, None)
        return defaults

    try:
        stat_res = config_path.stat()
        current_fingerprint = (stat_res.st_mtime_ns, stat_res.st_size)
    except OSError:
        return defaults

    with _CONFIG_CACHE_LOCK:
        if resolved_path in _CONFIG_CACHE and _CONFIG_CACHE_FINGERPRINT.get(resolved_path) == current_fingerprint:
            return copy.deepcopy(_CONFIG_CACHE[resolved_path])

    if config_path.exists():
        try:
            lock_path = str(config_path) + ".lock"
            with FileLock(lock_path):
                try:
                    stat_res = config_path.stat()
                    current_fingerprint = (stat_res.st_mtime_ns, stat_res.st_size)
                except OSError:
                    return defaults

                with _CONFIG_CACHE_LOCK:
                    if resolved_path in _CONFIG_CACHE and _CONFIG_CACHE_FINGERPRINT.get(resolved_path) == current_fingerprint:
                        return copy.deepcopy(_CONFIG_CACHE[resolved_path])

                with open(config_path) as f:
                    loaded = json.load(f)

                # Backfill override_agent_models if missing or incomplete
                needs_save = False

                if "override_agent_models" not in loaded or not isinstance(loaded["override_agent_models"], dict):
                    loaded["override_agent_models"] = copy.deepcopy(DEFAULT_OVERRIDE_AGENT_MODELS)
                    needs_save = True
                else:
                    for role, default_role_cfg in DEFAULT_OVERRIDE_AGENT_MODELS.items():
                        if role not in loaded["override_agent_models"]:
                            loaded["override_agent_models"][role] = copy.deepcopy(default_role_cfg)
                            needs_save = True

                # Backfill search_harness if missing or incomplete in existing config file on disk
                if "search_harness" not in loaded or not isinstance(loaded["search_harness"], dict):
                    loaded["search_harness"] = copy.deepcopy(DEFAULT_SEARCH_HARNESS)
                    needs_save = True
                else:
                    for k, default_val in DEFAULT_SEARCH_HARNESS.items():
                        if k not in loaded["search_harness"]:
                            loaded["search_harness"][k] = copy.deepcopy(default_val)
                            needs_save = True

                # Backfill graph_relations if missing or incomplete in existing config file on disk
                if "graph_relations" not in loaded or not isinstance(loaded["graph_relations"], dict):
                    loaded["graph_relations"] = copy.deepcopy(DEFAULT_GRAPH_RELATIONS)
                    needs_save = True
                else:
                    for k, default_val in DEFAULT_GRAPH_RELATIONS.items():
                        if k not in loaded["graph_relations"]:
                            loaded["graph_relations"][k] = copy.deepcopy(default_val)
                            needs_save = True

                # Backfill llm_num_retries if missing from existing config file on disk
                if "llm_num_retries" not in loaded:
                    loaded["llm_num_retries"] = DEFAULT_LLM_NUM_RETRIES
                    needs_save = True

                # Backfill llm_base_url if missing from existing config file on disk
                if "llm_base_url" not in loaded:
                    loaded["llm_base_url"] = DEFAULT_LLM_BASE_URL
                    needs_save = True

                # Backfill citations options if missing from existing config file on disk
                if "citations_enabled" not in loaded:
                    loaded["citations_enabled"] = DEFAULT_CITATIONS_ENABLED
                    needs_save = True
                if "references_section_enabled" not in loaded:
                    loaded["references_section_enabled"] = DEFAULT_REFERENCES_SECTION_ENABLED
                    needs_save = True
                if "considered_sources_disclosure" not in loaded:
                    loaded["considered_sources_disclosure"] = DEFAULT_CONSIDERED_SOURCES_DISCLOSURE
                    needs_save = True
                if "auto_render_mermaid_diagrams" not in loaded:
                    loaded["auto_render_mermaid_diagrams"] = DEFAULT_AUTO_RENDER_MERMAID_DIAGRAMS
                    needs_save = True
                if "enable_mcp_server" not in loaded:
                    loaded["enable_mcp_server"] = DEFAULT_ENABLE_MCP_SERVER
                    needs_save = True
                if "auto_spawn_backend" not in loaded:
                    loaded["auto_spawn_backend"] = DEFAULT_AUTO_SPAWN_BACKEND
                    needs_save = True
                if "enable_predictive_ghost_input" not in loaded:
                    loaded["enable_predictive_ghost_input"] = DEFAULT_ENABLE_PREDICTIVE_GHOST_INPUT
                    needs_save = True
                if "predictive_prompt_model" not in loaded:
                    loaded["predictive_prompt_model"] = DEFAULT_PREDICTIVE_PROMPT_MODEL
                    needs_save = True
                if "sentence_attribution_model" not in loaded:
                    loaded["sentence_attribution_model"] = DEFAULT_SENTENCE_ATTRIBUTION_MODEL
                    needs_save = True
                if "scholar_polite_email" not in loaded:
                    loaded["scholar_polite_email"] = DEFAULT_SCHOLAR_POLITE_EMAIL
                    needs_save = True
                if "conversion_timeout_seconds" not in loaded:
                    loaded["conversion_timeout_seconds"] = DEFAULT_CONVERSION_TIMEOUT_SECONDS
                    needs_save = True
                if "datalab_merge_cross_page" not in loaded:
                    loaded["datalab_merge_cross_page"] = DEFAULT_DATALAB_MERGE_CROSS_PAGE
                    needs_save = True
                if "user_research_profile" not in loaded:
                    loaded["user_research_profile"] = DEFAULT_USER_RESEARCH_PROFILE
                    needs_save = True
                if "max_research_profile_chars" not in loaded:
                    loaded["max_research_profile_chars"] = DEFAULT_MAX_RESEARCH_PROFILE_CHARS
                    needs_save = True
                if "max_saved_sessions" not in loaded:
                    loaded["max_saved_sessions"] = DEFAULT_MAX_SAVED_SESSIONS
                    needs_save = True
                if "session_retention_days" not in loaded:
                    loaded["session_retention_days"] = DEFAULT_SESSION_RETENTION_DAYS
                    needs_save = True
                if "session_db_auto_vacuum_on_startup" not in loaded:
                    loaded["session_db_auto_vacuum_on_startup"] = DEFAULT_SESSION_DB_AUTO_VACUUM_ON_STARTUP
                    needs_save = True
                if "prune_empty_sessions" not in loaded:
                    loaded["prune_empty_sessions"] = DEFAULT_PRUNE_EMPTY_SESSIONS
                    needs_save = True
                if "backend_port" not in loaded:
                    loaded["backend_port"] = DEFAULT_BACKEND_PORT
                    needs_save = True

                # Backfill Phase 7 onboarding SSOT keys if missing from existing
                # config files on disk (embedding_dimension / sync_auto_abort_failures)
                if "embedding_dimension" not in loaded:
                    loaded["embedding_dimension"] = DEFAULT_CONFIG["embedding_dimension"]
                    needs_save = True
                if "sync_auto_abort_failures" not in loaded:
                    loaded["sync_auto_abort_failures"] = DEFAULT_SYNC_AUTO_ABORT_FAILURES
                    needs_save = True

                if "vlm_captioning_backup_model" not in loaded:
                    loaded["vlm_captioning_backup_model"] = DEFAULT_CONFIG["vlm_captioning_backup_model"]
                    needs_save = True
                if "vlm_captioning_model" not in loaded:
                    loaded["vlm_captioning_model"] = DEFAULT_CONFIG["vlm_captioning_model"]
                    needs_save = True
                if "vlm_captioning_provider" not in loaded:
                    loaded["vlm_captioning_provider"] = DEFAULT_CONFIG["vlm_captioning_provider"]
                    needs_save = True
                if "enable_vlm_captioning" not in loaded:
                    loaded["enable_vlm_captioning"] = DEFAULT_CONFIG["enable_vlm_captioning"]
                    needs_save = True
                if "vlm_max_images_per_doc" not in loaded:
                    loaded["vlm_max_images_per_doc"] = DEFAULT_CONFIG["vlm_max_images_per_doc"]
                    needs_save = True

                # Detect and migrate deprecated compaction keys from disk
                if "compaction_token_threshold" in loaded and "adaptive_compaction_token_ceiling" not in loaded:
                    loaded["adaptive_compaction_token_ceiling"] = loaded["compaction_token_threshold"]
                    needs_save = True
                if "enable_context_compaction" in loaded and "enable_adaptive_compaction" not in loaded:
                    loaded["enable_adaptive_compaction"] = loaded["enable_context_compaction"]
                    needs_save = True
                if "adaptive_decision_model" in loaded and "adaptive_compaction_decision_model" not in loaded:
                    loaded["adaptive_compaction_decision_model"] = loaded.pop("adaptive_decision_model")
                    needs_save = True
                elif "adaptive_decision_model" in loaded:
                    loaded.pop("adaptive_decision_model", None)
                    needs_save = True

                # Purge all 5 deprecated compaction keys
                deprecated_compaction_keys = (
                    "compaction_interval",
                    "compaction_overlap_size",
                    "compaction_token_threshold",
                    "compaction_event_retention_size",
                    "enable_context_compaction",
                )
                for dep_key in deprecated_compaction_keys:
                    if dep_key in loaded:
                        loaded.pop(dep_key, None)
                        needs_save = True

                # Backfill missing adaptive compaction options
                adaptive_compaction_defaults = {
                    "enable_adaptive_compaction": DEFAULT_ENABLE_ADAPTIVE_COMPACTION,
                    "adaptive_compaction_retention_turns": DEFAULT_ADAPTIVE_COMPACTION_RETENTION_TURNS,
                    "adaptive_compaction_min_turns": DEFAULT_ADAPTIVE_COMPACTION_MIN_TURNS,
                    "adaptive_compaction_token_floor": DEFAULT_ADAPTIVE_COMPACTION_TOKEN_FLOOR,
                    "adaptive_compaction_token_ceiling": DEFAULT_ADAPTIVE_COMPACTION_TOKEN_CEILING,
                    "adaptive_compaction_threshold": DEFAULT_ADAPTIVE_COMPACTION_THRESHOLD,
                    "adaptive_compaction_decision_model": DEFAULT_ADAPTIVE_COMPACTION_DECISION_MODEL,
                    "compaction_summarizer_model": DEFAULT_COMPACTION_SUMMARIZER_MODEL,
                }
                for k, default_val in adaptive_compaction_defaults.items():
                    if k not in loaded:
                        loaded[k] = default_val
                        needs_save = True

                # Disallow local Laya for compaction decisions on disk; reset to default remote model
                if is_local_decision_model(loaded.get("adaptive_compaction_decision_model")):
                    loaded["adaptive_compaction_decision_model"] = DEFAULT_ADAPTIVE_COMPACTION_DECISION_MODEL
                    needs_save = True

                # Backfill citation verification options if missing from existing config file on disk
                verification_defaults = {
                    "citation_verification_enabled": DEFAULT_CITATION_VERIFICATION_ENABLED,
                    "citation_verification_provider": DEFAULT_CITATION_VERIFICATION_PROVIDER,
                    "citation_verification_model": DEFAULT_CITATION_VERIFICATION_MODEL,
                    "citation_verification_threshold": DEFAULT_CITATION_VERIFICATION_THRESHOLD,
                    "citation_verification_local_url": DEFAULT_CITATION_VERIFICATION_LOCAL_URL,
                    "citation_verification_local_port": DEFAULT_CITATION_VERIFICATION_LOCAL_PORT,
                    "citation_verification_auto_start_local": DEFAULT_CITATION_VERIFICATION_AUTO_START_LOCAL,
                    "citation_verification_timeout_sec": DEFAULT_CITATION_VERIFICATION_TIMEOUT_SEC,
                    "citation_verification_turn_budget_sec": DEFAULT_CITATION_VERIFICATION_TURN_BUDGET_SEC,
                }
                for k, default_val in verification_defaults.items():
                    if k not in loaded:
                        loaded[k] = default_val
                        needs_save = True
                # Migrate legacy default 3.0s/10.0s timeout and budget to 10.0s/20.0s for CPU inference headroom
                try:
                    if float(loaded.get("citation_verification_timeout_sec", 0)) == 3.0:
                        loaded["citation_verification_timeout_sec"] = DEFAULT_CITATION_VERIFICATION_TIMEOUT_SEC
                        needs_save = True
                except (ValueError, TypeError):
                    pass
                try:
                    if float(loaded.get("citation_verification_turn_budget_sec", 0)) == 10.0:
                        loaded["citation_verification_turn_budget_sec"] = DEFAULT_CITATION_VERIFICATION_TURN_BUDGET_SEC
                        needs_save = True
                except (ValueError, TypeError):
                    pass

                if needs_save:
                    try:
                        with open(config_path, "w") as f_out:
                            json.dump(loaded, f_out, indent=2)
                        stamp_config_mtime(resolved_dir)
                    except Exception as save_err:
                        logger.warning(
                            "Failed to backfill default dialogue agent models to %s: %s",
                            config_path,
                            save_err,
                        )
                    try:
                        stat_res = config_path.stat()
                        current_fingerprint = (stat_res.st_mtime_ns, stat_res.st_size)
                    except OSError:
                        pass

                res = dict(defaults)
                res.update(loaded)
                with _CONFIG_CACHE_LOCK:
                    _CONFIG_CACHE[resolved_path] = copy.deepcopy(res)
                    _CONFIG_CACHE_FINGERPRINT[resolved_path] = current_fingerprint
                return copy.deepcopy(res)
        except Exception as e:
            logger.warning("Failed to load config from %s: %s", config_path, e)
    return defaults

rerank_evidence async

rerank_evidence(query, chunks, max_chunks=None, word_ceiling=None)

Rerank candidate chunks using OpenRouter cross-encoder with budget trimming and fallback.

Source code in zocrates/agents/common/reranker.py
async def rerank_evidence(
    query: str,
    chunks: list[dict[str, Any]],
    max_chunks: int | None = None,
    word_ceiling: int | None = None,
) -> list[dict[str, Any]]:
    """Rerank candidate chunks using OpenRouter cross-encoder with budget trimming and fallback."""
    chunks = filter_substantive_evidence(chunks)

    effective_max_chunks = max_chunks if max_chunks is not None else MAX_RERANK_CHUNKS
    effective_word_ceiling = word_ceiling if word_ceiling is not None else HARD_WORD_CEILING

    if not chunks:
        return []

    if len(chunks) == 1:
        single = dict(chunks[0])
        single["rerank_score"] = None
        return _trim_to_budget([single], effective_word_ceiling)

    api_key = get_secure_api_key()

    # Fail-open enforcement (Phase 7 onboarding, plan §5.1/§8): a missing or
    # invalid reranker model config must degrade to un-reranked retrieval with
    # a logged warning — never fail the turn or the search run.
    try:
        reranker_model_raw = load_zocrates_config().get("reranker_model")
    except Exception as cfg_exc:
        logger.warning(
            "Failed to read reranker_model from config (%s). Skipping cross-encoder rerank.",
            cfg_exc,
        )
        trimmed = [dict(c, rerank_score=None) for c in chunks][:effective_max_chunks]
        return _trim_to_budget(trimmed, effective_word_ceiling)

    if not isinstance(reranker_model_raw, str) or not reranker_model_raw.strip():
        logger.warning(
            "No reranker model configured (reranker_model unset or invalid). "
            "Skipping cross-encoder rerank; retrieval proceeds un-reranked."
        )
        trimmed = [dict(c, rerank_score=None) for c in chunks][:effective_max_chunks]
        return _trim_to_budget(trimmed, effective_word_ceiling)

    reranker_model = reranker_model_raw.strip()

    if not api_key:
        logger.debug(
            "No OpenRouter API key found. Skipping cross-encoder rerank."
        )
        trimmed = [dict(c, rerank_score=None) for c in chunks][:effective_max_chunks]
        return _trim_to_budget(trimmed, effective_word_ceiling)

    documents = [c.get("text", "") for c in chunks]
    payload = {
        "model": reranker_model,
        "query": query,
        "documents": documents,
        "top_n": len(documents),
    }

    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.post(
                "https://openrouter.ai/api/v1/rerank",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
            resp.raise_for_status()
            data = resp.json()

        results = data.get("results", [])
        score_by_idx = {
            r["index"]: r.get("relevance_score", 0.0) for r in results
        }

        ranked_chunks = []
        for idx, chunk in enumerate(chunks):
            chunk_copy = dict(chunk)
            chunk_copy["rerank_score"] = float(score_by_idx.get(idx, 0.0))
            ranked_chunks.append(chunk_copy)

        ranked_chunks.sort(
            key=lambda c: c.get("rerank_score", 0.0), reverse=True
        )
        top_candidates = ranked_chunks[:effective_max_chunks]
        return _trim_to_budget(top_candidates, effective_word_ceiling)

    except Exception as e:
        logger.warning(
            "OpenRouter rerank call failed (%s: %s). Falling back to original chunk order.",
            type(e).__name__,
            e,
        )
        trimmed = [dict(c, rerank_score=None) for c in chunks][:effective_max_chunks]
        return _trim_to_budget(trimmed, effective_word_ceiling)