124 lines
4.8 KiB
TypeScript
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]
|
|
}
|