Spaces:

shreyask
/

qmd-web

Running

File size: 9,825 Bytes

import type { RRFResult, RerankedResult, FinalResult } from '../types';
import { InfoTooltip } from './PipelineView';

interface FusionColumnState {
  rrf: { status: 'idle' | 'done'; data?: { merged: RRFResult[] } };
  rerank: { status: 'idle' | 'running' | 'done'; data?: { before: RRFResult[]; after: RerankedResult[] } };
  finalResults?: FinalResult[];
}

interface FusionColumnProps {
  state: FusionColumnState;
  accent: string;
  info: string;
}

function Spinner({ color }: { color: string }) {
  return (
    <span style={{
      display: 'inline-block',
      width: '14px',
      height: '14px',
      border: '2px solid var(--border)',
      borderTopColor: color,
      borderRadius: '50%',
      animation: 'spin 0.7s linear infinite',
    }} />
  );
}

function SectionHeader({ label, color, badge }: { label: string; color: string; badge?: string }) {
  return (
    <div style={{
      fontSize: '0.68rem',
      fontWeight: 700,
      fontFamily: 'system-ui, -apple-system, sans-serif',
      color,
      textTransform: 'uppercase',
      letterSpacing: '0.06em',
      marginBottom: '0.35rem',
      display: 'flex',
      alignItems: 'center',
      gap: '0.4rem',
    }}>
      {label}
      {badge && (
        <span style={{ color: 'var(--text-muted)', fontWeight: 400, fontSize: '0.65rem' }}>{badge}</span>
      )}
    </div>
  );
}

function RRFRow({ result, rank }: { result: RRFResult; rank: number }) {
  return (
    <div style={{
      display: 'flex',
      alignItems: 'center',
      gap: '0.4rem',
      padding: '0.3rem 0.5rem',
      background: 'var(--bg-card)',
      border: '1px solid var(--border)',
      borderRadius: '5px',
      marginBottom: '0.2rem',
      fontSize: '0.72rem',
    }}>
      <span style={{
        fontFamily: "'SF Mono', 'Fira Code', 'Cascadia Code', monospace",
        color: 'var(--text-muted)',
        fontSize: '0.65rem',
        minWidth: '18px',
      }}>
        #{rank}
      </span>
      <span style={{
        flex: 1,
        fontFamily: 'system-ui, -apple-system, sans-serif',
        color: 'var(--text)',
        fontWeight: 500,
        overflow: 'hidden',
        textOverflow: 'ellipsis',
        whiteSpace: 'nowrap',
      }}>
        {result.title}
      </span>
      <span style={{
        fontFamily: "'SF Mono', 'Fira Code', 'Cascadia Code', monospace",
        fontSize: '0.65rem',
        color: '#2e7d32',
        fontWeight: 700,
        flexShrink: 0,
      }}>
        {result.score.toFixed(4)}
      </span>
    </div>
  );
}

// Rank badge: shows rank # with color coding for movement
function RankBadge({ label, rank, color }: { label: string; rank: number; color: string }) {
  return (
    <span style={{
      display: 'inline-flex',
      alignItems: 'center',
      gap: '0.15rem',
      fontFamily: "'SF Mono', 'Fira Code', 'Cascadia Code', monospace",
      fontSize: '0.6rem',
      color,
      fontWeight: 600,
    }}>
      <span style={{ color: 'var(--text-muted)', fontWeight: 400, fontSize: '0.55rem' }}>{label}</span>
      #{rank}
    </span>
  );
}

interface RankJourneyRow {
  docId: string;
  title: string;
  rrfRank?: number;
  rerankRank?: number;
  finalRank?: number;
}

// Shows each doc's rank journey: RRF #N → Reranker #N → Final #N
function RankJourney({ before, after, finalResults }: {
  before: RRFResult[];
  after: RerankedResult[];
  finalResults?: FinalResult[];
}) {
  const topLimit = 5;
  const topBefore = before.slice(0, topLimit);
  const topFinal = (finalResults ?? []).slice(0, topLimit);
  const rerankOrder = [...after].sort((a, b) => b.rerankScore - a.rerankScore);
  const titleMap = new Map<string, string>([
    ...before.map((result) => [result.docId, result.title] as const),
    ...after.map((result) => [result.docId, result.title] as const),
    ...topFinal.map((result) => [result.docId, result.title] as const),
  ]);
  const rrfRankMap = new Map(before.map((result, index) => [result.docId, index + 1]));
  const rerankRankMap = new Map(rerankOrder.map((r, i) => [r.docId, i + 1]));
  const finalRankMap = new Map((finalResults ?? []).map((result, index) => [result.docId, index + 1]));
  const rowMap = new Map<string, RankJourneyRow>();

  for (const result of [...topBefore, ...topFinal]) {
    const existing = rowMap.get(result.docId);
    rowMap.set(result.docId, {
      docId: result.docId,
      title: titleMap.get(result.docId) ?? result.title,
      rrfRank: rrfRankMap.get(result.docId),
      rerankRank: rerankRankMap.get(result.docId),
      finalRank: finalRankMap.get(result.docId),
      ...existing,
    });
  }

  const rows = [...rowMap.values()]
    .sort((a, b) =>
      (a.finalRank ?? Number.POSITIVE_INFINITY) - (b.finalRank ?? Number.POSITIVE_INFINITY) ||
      (a.rrfRank ?? Number.POSITIVE_INFINITY) - (b.rrfRank ?? Number.POSITIVE_INFINITY) ||
      (a.rerankRank ?? Number.POSITIVE_INFINITY) - (b.rerankRank ?? Number.POSITIVE_INFINITY),
    )
    .slice(0, topLimit);

  return (
    <div>
      {rows.map((row) => {
        return (
          <div key={row.docId} style={{
            padding: '0.3rem 0.5rem',
            background: 'var(--bg-card)',
            border: '1px solid var(--border)',
            borderRadius: '5px',
            marginBottom: '0.2rem',
            fontSize: '0.7rem',
          }}>
            <div style={{
              fontFamily: 'system-ui, -apple-system, sans-serif',
              color: 'var(--text)',
              fontWeight: 500,
              marginBottom: '0.2rem',
            }}>
              {row.title}
            </div>
            <div style={{
              display: 'flex',
              alignItems: 'center',
              gap: '0.3rem',
            }}>
              {row.rrfRank !== undefined && (
                <>
                  <RankBadge label="RRF" rank={row.rrfRank} color="var(--text-secondary)" />
                  <span style={{ color: 'var(--text-muted)', fontSize: '0.55rem' }}>{'\u2192'}</span>
                </>
              )}
              {row.rerankRank !== undefined && (
                <>
                  <RankBadge label="Reranker" rank={row.rerankRank} color="#f57f17" />
                  <span style={{ color: 'var(--text-muted)', fontSize: '0.55rem' }}>{'\u2192'}</span>
                </>
              )}
              {row.finalRank !== undefined ? (
                <RankBadge label="Final" rank={row.finalRank} color="#1b5e20" />
              ) : (
                <span style={{
                  fontSize: '0.55rem',
                  color: 'var(--text-muted)',
                  fontStyle: 'italic',
                }}>
                  blending...
                </span>
              )}
            </div>
          </div>
        );
      })}
      <div style={{
        fontSize: '0.62rem',
        fontFamily: 'system-ui, -apple-system, sans-serif',
        color: 'var(--text-muted)',
        marginTop: '0.3rem',
        fontStyle: 'italic',
        lineHeight: 1.4,
      }}>
        Final ranking blends reranker scores with retrieval position.
      </div>
    </div>
  );
}

export default function FusionColumn({ state, accent, info }: FusionColumnProps) {
  const rrfDone = state.rrf.status === 'done';
  const rerankRunning = state.rerank.status === 'running';
  const rerankDone = state.rerank.status === 'done';
  const isIdle = !rrfDone && !rerankRunning && !rerankDone;

  return (
    <div style={{ opacity: isIdle ? 0.45 : 1, transition: 'opacity 0.3s' }}>
      <div style={{
        display: 'flex',
        alignItems: 'center',
        gap: '0.4rem',
        marginBottom: '0.75rem',
        paddingBottom: '0.5rem',
        borderBottom: '1px solid var(--stage-divider)',
      }}>
        <span style={{
          width: '3px',
          height: '14px',
          borderRadius: '2px',
          background: accent,
          flexShrink: 0,
        }} />
        <h3 style={{
          margin: 0,
          fontSize: '0.78rem',
          fontFamily: 'system-ui, -apple-system, sans-serif',
          fontWeight: 700,
          color: accent,
          textTransform: 'uppercase',
          letterSpacing: '0.05em',
        }}>
          Fusion & Reranking
        </h3>
        {rerankRunning && <Spinner color={accent} />}
        <InfoTooltip text={info} />
      </div>

      {isIdle && (
        <p style={{ fontFamily: 'system-ui, -apple-system, sans-serif', fontSize: '0.75rem', color: 'var(--text-muted)', margin: 0 }}>
          Awaiting search...
        </p>
      )}

      {rrfDone && state.rrf.data && (
        <div style={{ marginBottom: '0.7rem' }}>
          <SectionHeader
            label="RRF Fusion"
            color="#558b2f"
            badge={`(${state.rrf.data.merged.length} docs)`}
          />
          {state.rrf.data.merged.slice(0, 5).map((r, i) => (
            <RRFRow key={r.docId} result={r} rank={i + 1} />
          ))}
          {state.rrf.data.merged.length > 5 && (
            <div style={{ fontSize: '0.68rem', color: 'var(--text-muted)', fontFamily: 'system-ui, -apple-system, sans-serif', paddingLeft: '0.25rem' }}>
              +{state.rrf.data.merged.length - 5} more
            </div>
          )}
        </div>
      )}

      {rerankRunning && !rerankDone && (
        <p style={{ fontFamily: 'system-ui, -apple-system, sans-serif', fontSize: '0.75rem', color: 'var(--text-secondary)', margin: '0 0 0.6rem 0', fontStyle: 'italic' }}>
          Reranking with cross-encoder...
        </p>
      )}

      {rerankDone && state.rerank.data && (
        <div>
          <SectionHeader label="Rank Journey" color="#33691e" />
          <RankJourney
            before={state.rerank.data.before}
            after={state.rerank.data.after}
            finalResults={state.finalResults}
          />
        </div>
      )}
    </div>
  );
}