project-nomad/admin/app/utils/rag_prompt.ts

124 lines
4.8 KiB
TypeScript

/**
* Pure prompt-budgeting helpers for the chat RAG pipeline.
*
* These live here rather than in RagPipelineService so they can be exercised
* under bare `node --experimental-strip-types` with no MySQL, Redis, Qdrant, or
* Ollama — the same reason `kb_ratio_lookup.ts` is shaped this way. The service
* supplies the config; everything below is a function of its arguments.
*/
/**
* Chars-per-token estimate used when budgeting the *prompt* (context trimming
* and the num_ctx ladder).
*
* NOTE: RagService.CHAR_TO_TOKEN_RATIO is 2, not 3.5 — the two halves of the
* system disagree about what a token costs, and RagService's own doc-comment
* says 3. That divergence is real and known; reconciling it changes chunk size
* and therefore retrieval results, so it is deliberately left alone here and
* tracked as its own measured change rather than folded into a refactor.
*/
export const PROMPT_CHARS_PER_TOKEN = 3.5
/**
* num_ctx is only requested once the system prompt is large enough to risk
* overflowing Ollama's silent 2048 default. Below the trigger we send nothing
* and inherit the server default.
*/
export const NUM_CTX_TRIGGER_TOKENS = 3000
export const NUM_CTX_RESPONSE_HEADROOM = 2048
export const NUM_CTX_LADDER = [8192, 16384, 32768, 65536]
export type ContextLimits = { maxResults: number; maxTokens: number }
export type ContextLimitTier = { maxParams: number; maxResults: number; maxTokens: number }
/** The minimum shape these helpers need; the real chunks carry more. */
export type BudgetableChunk = { text: string; metadata?: Record<string, any> }
export type BudgetableMessage = { role: 'system' | 'user' | 'assistant'; content: string }
/**
* Determine RAG context limits from the model size encoded in its name.
* Parses size indicators like "1b", "3b", "8b", "70b".
*
* PRESERVED QUIRK: an unparseable model name is treated as 8B. That is a guess,
* and for a name like "phi3" or a custom tag it can hand a small model far more
* context than it can actually use. Faithful to the pre-extraction behaviour.
*/
export function getContextLimitsForModel(
modelName: string,
tiers: readonly ContextLimitTier[]
): ContextLimits {
// e.g. "llama3.2:3b", "qwen2.5:1.5b", "gemma:7b"
const sizeMatch = modelName.match(/(\d+\.?\d*)[bB]/)
const paramBillions = sizeMatch ? Number.parseFloat(sizeMatch[1]) : 8 // default to 8B if unknown
for (const tier of tiers) {
if (paramBillions <= tier.maxParams) {
return { maxResults: tier.maxResults, maxTokens: tier.maxTokens }
}
}
return { maxResults: 5, maxTokens: 0 }
}
/**
* Apply the model-size context budget: cap the result count, then cap total
* characters.
*
* The first (most relevant) result is always kept — the token cap only gates
* subsequent results, so a single oversized chunk never starves the model of
* context entirely.
*/
export function trimToContextBudget<T extends BudgetableChunk>(
docs: T[],
limits: ContextLimits
): T[] {
const byCount = docs.slice(0, limits.maxResults)
if (limits.maxTokens <= 0) return byCount
const charCap = limits.maxTokens * PROMPT_CHARS_PER_TOKEN
let totalChars = 0
return byCount.filter((doc, idx) => {
totalChars += doc.text.length
return idx === 0 || totalChars <= charCap
})
}
/**
* Render retrieved chunks into the block the rag_context prompt wraps.
*
* Each block is labelled with its source title when one is available — a
* neutral, honest provenance signal — but never with the raw relevance score.
* nomic cosine scores for genuinely relevant passages sit around 0.4-0.6, and
* surfacing e.g. "42%" primes the model to distrust correct context. Scores
* stay in the logs and in the eval report.
*/
export function buildContextBlock(docs: BudgetableChunk[]): string {
return docs
.map((doc, idx) => {
const title = doc.metadata?.full_title || doc.metadata?.article_title
const label = title ? `[Context ${idx + 1}${title}]` : `[Context ${idx + 1}]`
return `${label}\n${doc.text}`
})
.join('\n\n')
}
/**
* Request a context window big enough to hold the system messages, but only
* once they are large enough to be at risk.
*
* Ollama respects num_ctx per request; LM Studio ignores it gracefully. Below
* the trigger we send nothing and inherit the server default — which for Ollama
* is a silent 2048.
*/
export function deriveNumCtx(messages: BudgetableMessage[]): number | undefined {
const systemChars = messages
.filter((m) => m.role === 'system')
.reduce((sum, m) => sum + m.content.length, 0)
const estimatedSystemTokens = Math.ceil(systemChars / PROMPT_CHARS_PER_TOKEN)
if (estimatedSystemTokens <= NUM_CTX_TRIGGER_TOKENS) return undefined
const needed = estimatedSystemTokens + NUM_CTX_RESPONSE_HEADROOM
return NUM_CTX_LADDER.find((n) => n >= needed) ?? NUM_CTX_LADDER[NUM_CTX_LADDER.length - 1]
}