candding 0.1.0

//! What a rerank entry feeds its model: a segment pair, a prompt, or a windowed pair.

use std::borrow::Cow;

use serde::Deserialize;
use serde::Serialize;

use crate::error::CanddingError;
use crate::error::Result;
use crate::tokenizer::PairTruncationPolicy;
use crate::tokenizer::PromptTemplate;

/// How a reranker's query and document become token ids, serialized with a `kind` tag such as
/// `{"kind": "pair", "truncation": "longest_first"}` on every surface that prints one.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
#[non_exhaustive]
pub enum RerankInput {
  /// The query and the document as the two segments of one pair, laid out by the published tokenizer's own
  /// pair post-processor, and cut by `truncation` when the pair is longer than the maximum length.
  Pair { truncation: PairTruncationPolicy },
  /// One prompt rendered from `template` with the instruction, the query and the document, tokenized whole; the
  /// default instruction is part of the scored input.
  Prompt {
    template: PromptTemplate,
    default_instruction: Cow<'static, str>,
  },
  /// The query, cut to `query_cap` ids, against each window of the document, consecutive windows sharing `overlap`
  /// tokens, at most a quarter of a window; the document scores as its best window.
  Windowed { query_cap: usize, overlap: usize },
}

impl RerankInput {
  /// The truncation policy of a pair input.
  pub fn pair_truncation(&self) -> Option<PairTruncationPolicy> {
    match self {
      RerankInput::Pair { truncation } => Some(*truncation),
      RerankInput::Prompt { .. } | RerankInput::Windowed { .. } => None,
    }
  }

  /// The template of a prompt input.
  pub fn template(&self) -> Option<PromptTemplate> {
    match self {
      RerankInput::Prompt { template, .. } => Some(*template),
      RerankInput::Pair { .. } | RerankInput::Windowed { .. } => None,
    }
  }

  /// The instruction a prompt input renders when the caller names none.
  pub fn default_instruction(&self) -> Option<&str> {
    match self {
      RerankInput::Prompt {
        default_instruction,
        ..
      } => Some(default_instruction),
      RerankInput::Pair { .. } | RerankInput::Windowed { .. } => None,
    }
  }

  /// The query cap and the overlap of a windowed input.
  pub fn windows(&self) -> Option<(usize, usize)> {
    match self {
      RerankInput::Windowed { query_cap, overlap } => Some((*query_cap, *overlap)),
      RerankInput::Pair { .. } | RerankInput::Prompt { .. } => None,
    }
  }

  /// Refuses an instruction for a pair or a windowed input, which has no instruction slot; a prompt input takes any,
  /// the empty one included.
  pub fn check_instruction(&self, model: &str, instruction: Option<&str>) -> Result<()> {
    match (self, instruction) {
      (RerankInput::Pair { .. } | RerankInput::Windowed { .. }, Some(_)) => {
        Err(CanddingError::InstructionOnPairInput {
          model: model.to_string(),
        })
      }
      _ => Ok(()),
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  /// The tag and the field names are the serialized contract every surface shares.
  #[test]
  fn a_pair_input_serializes_with_its_kind_and_truncation() {
    let input = RerankInput::Pair {
      truncation: PairTruncationPolicy::PreserveQueryToThreeQuarters,
    };
    let json = serde_json::to_value(&input).unwrap();
    assert_eq!(
      json,
      serde_json::json!({"kind": "pair", "truncation": "preserve_query_to_three_quarters"})
    );
    assert_eq!(serde_json::from_value::<RerankInput>(json).unwrap(), input);
    assert_eq!(
      input.pair_truncation(),
      Some(PairTruncationPolicy::PreserveQueryToThreeQuarters)
    );
    assert_eq!(input.template(), None);
    assert_eq!(input.default_instruction(), None);
  }

  #[test]
  fn a_prompt_input_serializes_with_its_kind_template_and_default_instruction() {
    let input = RerankInput::Prompt {
      template: PromptTemplate::Qwen3Reranker,
      default_instruction: Cow::Borrowed("Judge it"),
    };
    let json = serde_json::to_value(&input).unwrap();
    assert_eq!(
      json,
      serde_json::json!({"kind": "prompt", "template": "qwen3_reranker", "default_instruction": "Judge it"})
    );
    assert_eq!(serde_json::from_value::<RerankInput>(json).unwrap(), input);
    assert_eq!(input.pair_truncation(), None);
    assert_eq!(input.template(), Some(PromptTemplate::Qwen3Reranker));
    assert_eq!(input.default_instruction(), Some("Judge it"));
  }

  #[test]
  fn a_windowed_input_serializes_with_its_kind_query_cap_and_overlap() {
    let input = RerankInput::Windowed {
      query_cap: 512,
      overlap: 80,
    };
    let json = serde_json::to_value(&input).unwrap();
    assert_eq!(
      json,
      serde_json::json!({"kind": "windowed", "query_cap": 512, "overlap": 80})
    );
    assert_eq!(serde_json::from_value::<RerankInput>(json).unwrap(), input);
    assert_eq!(input.windows(), Some((512, 80)));
    assert_eq!(input.pair_truncation(), None);
    assert_eq!(input.template(), None);
    assert_eq!(input.default_instruction(), None);
  }

  #[test]
  fn a_windowed_input_refuses_an_instruction() {
    let input = RerankInput::Windowed {
      query_cap: 512,
      overlap: 80,
    };
    for instruction in ["judge it", ""] {
      assert!(matches!(
        input.check_instruction("m", Some(instruction)),
        Err(CanddingError::InstructionOnPairInput { .. })
      ));
    }
    assert!(input.check_instruction("m", None).is_ok());
  }

  /// `PromptTemplate::name` is what error messages print, so it must be the serde spelling every other surface uses.
  #[test]
  fn a_template_name_is_its_serde_name() {
    let template = PromptTemplate::Qwen3Reranker;
    assert_eq!(
      serde_json::to_value(template).unwrap(),
      serde_json::Value::String(template.name().to_string())
    );
  }
}