connectai/src/retrieval/contextBudget.ts

/**
 * ============================================================
 * Context Budget Manager (컨텍스트 예산 관리)
 * 
 * 시스템 프롬프트의 토큰 예산을 관리하여
 * 로컬 모델의 context window를 효율적으로 활용합니다.
 * ============================================================
 */

import { RetrievalChunk, ContextBudgetConfig } from './types';

const DEFAULT_BUDGET: ContextBudgetConfig = {
    totalBudget: 8000,     // ~32K context 중 retrieval에 할당
    retrievalRatio: 0.4,   // 40%
    minChunks: 2,
    maxChunks: 12
};

// [코어 수렴] 토큰 추정은 lib/contextManager 의 단일 구현을 쓴다 (CJK 전반 커버·
// 보수적 계수). 예전 이 파일의 자체 추정(한글만 1.5/0.25)은 제거 — 기존 소비자
// 호환을 위해 같은 이름으로 재수출.
import { estimateTokens } from '../lib/contextManager';
export { estimateTokens };


/**
 * 검색 결과 청크들을 토큰 예산 내에서 선택합니다.
 * 
 * 선택 전략:
 * 1. 스코어 내림차순 정렬
 * 2. 중복 제거 (같은 filePath를 가진 청크)
 * 3. 토큰 예산 내에서 순서대로 선택
 * 4. 최소 청크 수 보장
 */
export function selectWithinBudget(
    chunks: RetrievalChunk[],
    config: Partial<ContextBudgetConfig> = {}
): { selected: RetrievalChunk[]; dropped: RetrievalChunk[]; tokensUsed: number } {
    const cfg = { ...DEFAULT_BUDGET, ...config };
    const budget = Math.floor(cfg.totalBudget * cfg.retrievalRatio);

    // 1. Sort by score descending
    const sorted = [...chunks].sort((a, b) => b.score - a.score);

    // 2. [Structural Fix] 파일당 청크 수 제한 완화 (Deduplication -> Multi-context)
    const fileChunkCounts = new Map<string, number>();
    const filtered = sorted.filter((chunk) => {
        const key = chunk.metadata.filePath || chunk.id;
        const count = fileChunkCounts.get(key) || 0;
        
        // 파일당 최대 3개까지의 주요 맥락 허용 (정보 유실 방지)
        if (count >= 3) return false;
        
        fileChunkCounts.set(key, count + 1);
        return true;
    });

    // 3. Select within budget
    const selected: RetrievalChunk[] = [];
    const dropped: RetrievalChunk[] = [];
    let tokensUsed = 0;

    for (const chunk of filtered) {
        const chunkTokens = chunk.tokenEstimate || estimateTokens(chunk.content);

        if (selected.length >= cfg.maxChunks) {
            dropped.push(chunk);
            continue;
        }

        if (tokensUsed + chunkTokens > budget && selected.length >= cfg.minChunks) {
            dropped.push(chunk);
            continue;
        }

        selected.push(chunk);
        tokensUsed += chunkTokens;
    }

    return { selected, dropped, tokensUsed };
}

/**
 * 청크의 '주제(Subject)' 태그를 도출한다 — 서로 다른 프로젝트/주제의 정보가 한
 * 컨텍스트에 섞일 때 모델이 경계를 인지하도록(무성 교차오염 방지). category 가 있으면
 * 그걸, 없으면 title/filePath 의 최상위 폴더 세그먼트를 주제로 본다. 파일명만 있으면 ''.
 */
function deriveSubject(chunk: RetrievalChunk): string {
    const cat = (chunk.metadata.category || '').trim();
    if (cat) return cat;
    const ref = (chunk.title || chunk.metadata.filePath || '').replace(/\\/g, '/');
    const seg = ref.split('/').filter(Boolean);
    return seg.length >= 2 ? seg[0] : '';
}

/**
 * 선택된 청크들을 하나의 컨텍스트 문자열로 조립합니다.
 * 소스별로 그룹화하여 가독성을 높입니다.
 */
export function assembleContext(chunks: RetrievalChunk[]): string {
    if (chunks.length === 0) return '';

    const sourceLabels: Record<string, string> = {
        'brain-trace': '📚 Second Brain Knowledge',
        'brain-memory': '📚 Brain Knowledge',
        'long-term-memory': '🧠 Long-Term Memory (사용자 규칙/결정)',
        'medium-term-memory': '🗂️ Medium-Term Memory (최근 세션 요약)',
        'project-memory': '📂 Project Memory (프로젝트 컨텍스트)',
        'procedural-memory': '📋 Procedural Memory (반복 절차)',
        'episodic-memory': '📖 Episodic Memory (과거 대화 흐름)',
        'project-scan': '🔍 Project Scan',
        'recent-knowledge': '📄 Recent Project Knowledge'
    };

    // Group by source
    const groups = new Map<string, RetrievalChunk[]>();
    for (const chunk of chunks) {
        const key = chunk.source;
        if (!groups.has(key)) groups.set(key, []);
        groups.get(key)!.push(chunk);
    }

    const sections: string[] = [];
    for (const [source, groupChunks] of groups) {
        const label = sourceLabels[source] || source;
        const items = groupChunks
            .map((c) => {
                const metadata = c.metadata;
                const subject = deriveSubject(c);
                const subjectTag = subject ? `[${subject}] ` : '';
                const conflictTag = metadata.conflictDetected ? ` [⚠️ CONFLICT: ${metadata.conflictSeverity}]` : '';
                const coverageTag = metadata.queryCoverage !== undefined ? ` (Coverage: ${metadata.queryCoverage.toFixed(2)})` : '';
                return `- ${subjectTag}${c.title}${conflictTag}${coverageTag}: ${c.content}`;
            })
            .join('\n');
        sections.push(`### ${label}\n${items}`);
    }

    return [
        '[MEMORY CONTEXT]',
        'Review this layered memory before preparing the answer. Use it only when relevant, and prefer the current user request when there is conflict.',
        '각 항목 앞의 [주제] 태그와 섹션 출처를 확인하라. **현재 요청과 다른 프로젝트·주제의 항목은 사용하지 마라** — 서로 다른 프로젝트의 규칙·결정·수치·고유명사를 섞지 말 것. 어느 항목이 현재 작업과 관련 있는지 불확실하면 그 항목에 의존하지 마라.',
        '',
        sections.join('\n\n')
    ].join('\n');
}