connectai/src/retrieval/semanticRerank.ts

/**
 * LLM Semantic Re-ranking — TF-IDF / 임베딩이 놓치는 *의도* 매치를 작은 LLM 호출
 * 한 번으로 잡는다.
 *
 * 동작:
 *   1. 1차 검색(TF-IDF + embedding + 부스트들) 결과의 *상위 K* (기본 15) 후보를 추출
 *   2. 가벼운 프롬프트로 LLM 에게 "이 중 query 의도에 가장 부합하는 순서로 ID 나열" 요청
 *   3. LLM 응답을 파싱해 순서 적용 — 응답 실패/누락 ID 는 원순서 유지
 *
 * 비용·위험 관리:
 *   - 기본 OFF (g1nation.semanticRerankEnabled). 사용자가 latency 감수할 의지 있을 때만.
 *   - 짧은 timeout (기본 8초) — 초과 시 원순서 그대로 반환, 검색 실패 안 됨.
 *   - 후보 K 제한 — 토큰 비용 cap.
 *   - 별도 빠른 모델 지정 가능 (`g1nation.semanticRerankModel`) — 메인 모델 외 작은 모델.
 *
 * 인터페이스: input chunks 순서는 *원본 score 내림차순* 으로 들어와야 함.
 * 반환: re-rank 가 성공하면 새 순서의 RetrievalChunk[], 실패하면 원순서.
 */

import { RetrievalChunk } from './types';

export interface SemanticRerankOptions {
    ollamaUrl: string;
    /** Re-rank 전용 모델 ID. 비면 fallback model 사용. */
    model: string;
    /** 후보로 LLM 에 넘길 최대 chunk 개수. 기본 15. */
    candidateK: number;
    /** LLM 호출 타임아웃 (ms). 기본 8000. */
    timeoutMs: number;
    /** 각 chunk 미리보기 길이. 기본 240 chars. */
    excerptLength: number;
}

export const DEFAULT_SEMANTIC_RERANK_OPTIONS: Omit<SemanticRerankOptions, 'ollamaUrl' | 'model'> = {
    candidateK: 15,
    timeoutMs: 8000,
    excerptLength: 240,
};

export interface SemanticRerankResult {
    rerankedChunks: RetrievalChunk[];
    /** true 면 LLM 응답으로 순서 변경됨. false 면 원순서 (실패/타임아웃/파싱 실패). */
    success: boolean;
    durationMs: number;
    /** 디버그·footer 표시용 — re-rank 가 어떻게 동작했는지. */
    note: string;
}

function shortExcerpt(text: string, n: number): string {
    if (!text) return '';
    const cleaned = text.replace(/\s+/g, ' ').trim();
    return cleaned.length <= n ? cleaned : cleaned.slice(0, n) + '…';
}

function buildRerankPrompt(query: string, candidates: RetrievalChunk[], excerptLength: number): { system: string; user: string } {
    const lines: string[] = [];
    for (let i = 0; i < candidates.length; i++) {
        const c = candidates[i];
        lines.push(`[C${i + 1}] (${c.source}) ${c.title || '(제목 없음)'}`);
        lines.push(`  ${shortExcerpt(c.content, excerptLength)}`);
    }

    const system = [
        '당신은 검색 결과 재정렬기 (re-ranker). 사용자 질의의 *의도* 와 각 후보 문서의 *내용 부합도* 를 평가해 가장 유용한 순서로 정렬.',
        '',
        '[규칙]',
        '1. 응답은 *반드시* 한 줄의 JSON: `{"ranking":[3,1,5,2,4,...]}` 형식.',
        '2. ranking 배열 원소 = 입력 [C1], [C2] 의 *번호* (1-based).',
        '3. 모든 입력 후보를 한 번씩만 포함. 누락·중복·번호 외 값 금지.',
        '4. 다른 설명·코드 블록·텍스트 출력 절대 금지 — JSON 한 줄만.',
        '5. 평가 기준: (a) 질의 의도와의 직접 부합도 > (b) 키워드 매치 > (c) 문맥 풍부도.',
    ].join('\n');

    const user = [
        `[사용자 질의]\n${query}`,
        '',
        `[후보 ${candidates.length}개]`,
        ...lines,
        '',
        '위 후보를 가장 부합도 높은 순서로 정렬한 ranking 배열만 JSON 한 줄로 출력.',
    ].join('\n');

    return { system, user };
}

/**
 * Ollama / OpenAI 호환 endpoint 로 단발 호출. agents/factory.ts 의 BaseAgent.callLLM
 * 패턴 단순화. timeout, retry 1회만.
 */
async function callLlmForRerank(
    ollamaUrl: string,
    model: string,
    system: string,
    user: string,
    timeoutMs: number,
): Promise<string> {
    const isOllama = ollamaUrl.includes(':11434') || ollamaUrl.includes('ollama');
    const endpoint = isOllama ? `${ollamaUrl}/api/chat` : `${ollamaUrl}/v1/chat/completions`;
    const controller = new AbortController();
    const timer = setTimeout(() => controller.abort(), timeoutMs);
    try {
        const body = isOllama
            ? {
                model, stream: false,
                messages: [
                    { role: 'system', content: system },
                    { role: 'user', content: user },
                ],
                options: { temperature: 0.0, num_predict: 256 },
            }
            : {
                model,
                messages: [
                    { role: 'system', content: system },
                    { role: 'user', content: user },
                ],
                stream: false, temperature: 0.0, max_tokens: 256,
            };
        const res = await fetch(endpoint, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(body),
            signal: controller.signal,
        });
        if (!res.ok) throw new Error(`HTTP ${res.status}`);
        const data: any = await res.json();
        const content =
            data?.message?.content ??
            data?.choices?.[0]?.message?.content ??
            data?.choices?.[0]?.text ??
            data?.response ??
            '';
        return String(content || '');
    } finally {
        clearTimeout(timer);
    }
}

/** LLM 응답에서 ranking 배열 추출 + 검증. 실패 시 null. */
function parseRanking(raw: string, expectedSize: number): number[] | null {
    if (!raw) return null;
    // JSON 한 줄 추출 — { ... } 안에 ranking
    const match = raw.match(/\{[\s\S]*?\}/);
    if (!match) return null;
    try {
        const parsed = JSON.parse(match[0]);
        const arr = parsed?.ranking;
        if (!Array.isArray(arr)) return null;
        const seen = new Set<number>();
        const out: number[] = [];
        for (const v of arr) {
            const n = typeof v === 'number' ? v : parseInt(String(v), 10);
            if (!Number.isFinite(n) || n < 1 || n > expectedSize) continue;
            if (seen.has(n)) continue;
            seen.add(n);
            out.push(n);
        }
        // 누락 보충 — LLM 이 일부 빠뜨렸으면 원순서로 뒤에 붙임.
        for (let i = 1; i <= expectedSize; i++) {
            if (!seen.has(i)) out.push(i);
        }
        return out.length === expectedSize ? out : null;
    } catch {
        return null;
    }
}

export async function semanticRerank(
    query: string,
    chunks: RetrievalChunk[],
    options: SemanticRerankOptions,
): Promise<SemanticRerankResult> {
    const start = Date.now();
    const k = Math.max(2, Math.min(options.candidateK, chunks.length));
    if (chunks.length < 2 || k < 2) {
        return { rerankedChunks: chunks, success: false, durationMs: 0, note: 'too few candidates' };
    }
    // 입력은 score 내림차순 가정 — 상위 K 가 re-rank 대상, 나머지는 그대로 꼬리.
    const candidates = chunks.slice(0, k);
    const tail = chunks.slice(k);

    const { system, user } = buildRerankPrompt(query, candidates, options.excerptLength);

    let raw = '';
    try {
        raw = await callLlmForRerank(options.ollamaUrl, options.model, system, user, options.timeoutMs);
    } catch (e: any) {
        return {
            rerankedChunks: chunks,
            success: false,
            durationMs: Date.now() - start,
            note: `LLM call failed: ${e?.name || e?.message || 'unknown'}`,
        };
    }

    const ranking = parseRanking(raw, candidates.length);
    if (!ranking) {
        return {
            rerankedChunks: chunks,
            success: false,
            durationMs: Date.now() - start,
            note: 'unparseable LLM response',
        };
    }

    const reranked = ranking.map((i) => candidates[i - 1]);
    return {
        rerankedChunks: [...reranked, ...tail],
        success: true,
        durationMs: Date.now() - start,
        note: `re-ranked top ${k} (changed positions: ${ranking.filter((v, i) => v !== i + 1).length})`,
    };
}