feat(AI): make RAG retreival optional per chat (#1247)
This commit is contained in:
Jake Turner
GitHub
parent
4b2b3b5ebe
commit
e6acf9938a
|
|
@ -36,6 +36,12 @@ server/temp
|
|||
.vscode
|
||||
.idea
|
||||
|
||||
# Agent Files
|
||||
.agents
|
||||
.claude
|
||||
.codegraph
|
||||
.mcp.json
|
||||
|
||||
# Frontend assets compiled code
|
||||
admin/public/assets
|
||||
|
||||
|
|
|
|||
|
|
@ -63,9 +63,14 @@ export default class OllamaController {
|
|||
// Everything from system-prompt assembly through query rewriting,
|
||||
// retrieval, context trimming and the num_ctx decision lives in
|
||||
// RagPipelineService so the eval harness exercises this exact code path.
|
||||
// Knowledge base retrieval is user-toggleable (chat header + AI Assistant
|
||||
// settings, both writing rag.enabled). Unset means on, preserving the
|
||||
// behaviour from before the toggle existed.
|
||||
const ragEnabled = (await KVStore.getValue('rag.enabled')) ?? true
|
||||
const collectionFilter: string | null = request.input('collection', null)
|
||||
const trace = await this.ragPipelineService.buildPrompt(reqData.messages, reqData.model, {
|
||||
collection: collectionFilter ?? undefined,
|
||||
skipRetrieval: !ragEnabled,
|
||||
})
|
||||
reqData.messages = trace.messages
|
||||
const numCtx = trace.numCtx
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ export default class SettingsController {
|
|||
const ollamaFlashAttention = await KVStore.getValue('ai.ollamaFlashAttention')
|
||||
const autoThinking = await KVStore.getValue('ai.autoThinking')
|
||||
const tasksModel = await KVStore.getValue('ai.tasksModel')
|
||||
const ragEnabled = await KVStore.getValue('rag.enabled')
|
||||
return inertia.render('settings/models', {
|
||||
models: {
|
||||
availableModels: availableModels?.models || [],
|
||||
|
|
@ -83,6 +84,7 @@ export default class SettingsController {
|
|||
ollamaFlashAttention: ollamaFlashAttention ?? true,
|
||||
autoThinking: autoThinking ?? false,
|
||||
tasksModel: tasksModel ?? '',
|
||||
ragEnabled: ragEnabled ?? true,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
|
|
|||
|
|
@ -88,6 +88,12 @@ export class RagPipelineService {
|
|||
if (opts.oracleContext) {
|
||||
relevantDocs = opts.oracleContext
|
||||
trace.retrieved = relevantDocs
|
||||
} else if (opts.skipRetrieval) {
|
||||
// Retrieval turned off by the user (rag.enabled). Bail before the
|
||||
// hasDocuments check, the rewrite LLM call and the vector search — the
|
||||
// whole point is to spend nothing here. relevantDocs stays empty, so no
|
||||
// context block is injected below.
|
||||
logger.debug('[RagPipeline] Retrieval disabled by setting, skipping')
|
||||
} else {
|
||||
const rewriteStart = Date.now()
|
||||
const { query, didRewrite } = await this.resolveRetrievalQuery(working, model, opts)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ export const SETTINGS_KEYS: KVStoreKey[] = [
|
|||
'ai.autoThinking',
|
||||
'ai.tasksModel',
|
||||
'rag.defaultIngestPolicy',
|
||||
'rag.enabled',
|
||||
'autoUpdate.enabled',
|
||||
'autoUpdate.windowStart',
|
||||
'autoUpdate.windowEnd',
|
||||
|
|
|
|||
|
|
@ -76,6 +76,33 @@ export default function Chat({
|
|||
const autoThinkingDefault =
|
||||
autoThinkingSetting?.value === true || autoThinkingSetting?.value === 'true'
|
||||
|
||||
// Knowledge base retrieval, shared with AI Assistant settings (same KV key).
|
||||
// Unset means on, so coerce off the negative — an absent value must not read
|
||||
// as false.
|
||||
const { data: ragEnabledSetting } = useSystemSetting({ key: 'rag.enabled', enabled })
|
||||
const ragEnabled = !(ragEnabledSetting?.value === false || ragEnabledSetting?.value === 'false')
|
||||
|
||||
const ragEnabledMutation = useMutation({
|
||||
mutationFn: async (value: boolean) => await api.updateSetting('rag.enabled', value),
|
||||
// Flip the switch immediately rather than after the round-trip, and roll
|
||||
// back if the write fails.
|
||||
onMutate: async (value: boolean) => {
|
||||
await queryClient.cancelQueries({ queryKey: ['system-setting', 'rag.enabled'] })
|
||||
const previous = queryClient.getQueryData(['system-setting', 'rag.enabled'])
|
||||
queryClient.setQueryData(['system-setting', 'rag.enabled'], {
|
||||
key: 'rag.enabled',
|
||||
value,
|
||||
})
|
||||
return { previous }
|
||||
},
|
||||
onError: (_err, _value, context) => {
|
||||
queryClient.setQueryData(['system-setting', 'rag.enabled'], context?.previous)
|
||||
},
|
||||
onSettled: () => {
|
||||
queryClient.invalidateQueries({ queryKey: ['system-setting', 'rag.enabled'] })
|
||||
},
|
||||
})
|
||||
|
||||
const { data: remoteStatus } = useQuery({
|
||||
queryKey: ['remoteOllamaStatus'],
|
||||
queryFn: () => api.getRemoteOllamaStatus(),
|
||||
|
|
@ -558,6 +585,7 @@ export default function Chat({
|
|||
{remoteStatus?.connected === false ? 'Remote Disconnected' : 'Remote Connected'}
|
||||
</span>
|
||||
)}
|
||||
{ragEnabled && (
|
||||
<div className="flex items-center gap-2">
|
||||
<label htmlFor="collection-select" className="text-sm text-text-secondary">
|
||||
Search in:
|
||||
|
|
@ -574,6 +602,7 @@ export default function Chat({
|
|||
))}
|
||||
</select>
|
||||
</div>
|
||||
)}
|
||||
<div className="flex items-center gap-2 min-w-0">
|
||||
<label htmlFor="model-select" className="text-sm text-text-secondary">
|
||||
Model:
|
||||
|
|
@ -598,6 +627,19 @@ export default function Chat({
|
|||
</select>
|
||||
)}
|
||||
</div>
|
||||
<div className="flex items-center">
|
||||
<span className="text-sm text-text-secondary select-none">Knowledge Base:</span>
|
||||
<InfoTooltip
|
||||
position="bottom"
|
||||
align="right"
|
||||
text="When on, the assistant searches your knowledge base for relevant documents before answering. Turning this off is faster and lighter on hardware, which helps when your knowledge base is small or empty. This is the same setting as in AI Assistant settings."
|
||||
/>
|
||||
<Switch
|
||||
id="chat-rag-toggle"
|
||||
checked={ragEnabled}
|
||||
onChange={(v) => ragEnabledMutation.mutate(v)}
|
||||
/>
|
||||
</div>
|
||||
{selectedModelSupportsThinking && (
|
||||
<div className="flex items-center">
|
||||
<span className="text-sm text-text-secondary select-none">Thinking:</span>
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import { SERVICE_NAMES } from '../../../constants/service_names'
|
|||
import Switch from '~/components/inputs/Switch'
|
||||
import Select from '~/components/inputs/Select'
|
||||
import StyledSectionHeader from '~/components/StyledSectionHeader'
|
||||
import { useMutation, useQuery } from '@tanstack/react-query'
|
||||
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
|
||||
import Input from '~/components/inputs/Input'
|
||||
import { IconSearch, IconRefresh } from '@tabler/icons-react'
|
||||
import { formatBytes } from '~/lib/util'
|
||||
|
|
@ -27,7 +27,7 @@ export default function ModelsPage(props: {
|
|||
models: {
|
||||
availableModels: NomadOllamaModel[]
|
||||
installedModels: NomadInstalledModel[]
|
||||
settings: { chatSuggestionsEnabled: boolean; aiAssistantCustomName: string; remoteOllamaUrl: string; ollamaFlashAttention: boolean; autoThinking: boolean; tasksModel: string }
|
||||
settings: { chatSuggestionsEnabled: boolean; aiAssistantCustomName: string; remoteOllamaUrl: string; ollamaFlashAttention: boolean; autoThinking: boolean; tasksModel: string; ragEnabled: boolean }
|
||||
}
|
||||
}) {
|
||||
const { aiAssistantName } = usePage<{ aiAssistantName: string }>().props
|
||||
|
|
@ -36,6 +36,7 @@ export default function ModelsPage(props: {
|
|||
const { openModal, closeAllModals } = useModals()
|
||||
const { debounce } = useDebounce()
|
||||
const { data: systemInfo } = useSystemInfo({})
|
||||
const queryClient = useQueryClient()
|
||||
|
||||
const [gpuBannerDismissed, setGpuBannerDismissed] = useState(() => {
|
||||
try {
|
||||
|
|
@ -100,6 +101,7 @@ export default function ModelsPage(props: {
|
|||
props.models.settings.ollamaFlashAttention
|
||||
)
|
||||
const [autoThinking, setAutoThinking] = useState(props.models.settings.autoThinking)
|
||||
const [ragEnabled, setRagEnabled] = useState(props.models.settings.ragEnabled)
|
||||
const [tasksModel, setTasksModel] = useState(props.models.settings.tasksModel)
|
||||
const [aiAssistantCustomName, setAiAssistantCustomName] = useState(
|
||||
props.models.settings.aiAssistantCustomName
|
||||
|
|
@ -258,7 +260,11 @@ export default function ModelsPage(props: {
|
|||
mutationFn: async ({ key, value }: { key: string; value: boolean | string }) => {
|
||||
return await api.updateSetting(key, value)
|
||||
},
|
||||
onSuccess: () => {
|
||||
onSuccess: (_data, { key }) => {
|
||||
// Anything reading this key through useSystemSetting (e.g. the chat
|
||||
// window's own copy of the retrieval toggle) should pick the change up
|
||||
// without a reload.
|
||||
queryClient.invalidateQueries({ queryKey: ['system-setting', key] })
|
||||
addNotification({
|
||||
message: 'Setting updated successfully.',
|
||||
type: 'success',
|
||||
|
|
@ -343,6 +349,15 @@ export default function ModelsPage(props: {
|
|||
label="Use thinking automatically when a model supports it"
|
||||
description="Sets the default for models that can think. You can still turn thinking on or off for an individual model in the chat window."
|
||||
/>
|
||||
<Switch
|
||||
checked={ragEnabled}
|
||||
onChange={(newVal) => {
|
||||
setRagEnabled(newVal)
|
||||
updateSettingMutation.mutate({ key: 'rag.enabled', value: newVal })
|
||||
}}
|
||||
label="Knowledge Base Retrieval"
|
||||
description="Search your knowledge base for relevant documents before answering. Turn this off to save memory and speed up replies when your knowledge base is small or empty. This is the same switch as the one in the chat window."
|
||||
/>
|
||||
<Input
|
||||
name="aiAssistantCustomName"
|
||||
label="Assistant Name"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,125 @@
|
|||
/**
|
||||
* The `rag.enabled` off-switch, at the level that actually matters: what
|
||||
* buildPrompt does and does not call.
|
||||
*
|
||||
* Turning retrieval off is a resource decision, so "no context injected" is
|
||||
* only half the contract — the point is that none of the three expensive steps
|
||||
* (hasDocuments, the query-rewrite LLM call, the Qdrant search) run at all.
|
||||
* These tests count calls, not just output.
|
||||
*
|
||||
* Japa suite (not the plain-node `test:unit` runner): RagPipelineService pulls
|
||||
* in the AdonisJS container, so it needs the ignited app that `node ace test`
|
||||
* provides. Run with:
|
||||
* node ace test --suites=unit --files=rag_retrieval_toggle
|
||||
*
|
||||
* Deliberately NOT named rag_pipeline_*: that glob is what `npm run test:eval`
|
||||
* feeds to the plain-node runner, which cannot boot Adonis.
|
||||
*/
|
||||
import { test } from '@japa/runner'
|
||||
import { RagPipelineService } from '#services/rag_pipeline_service'
|
||||
import { SYSTEM_PROMPTS } from '../../constants/ollama.js'
|
||||
import type { OllamaChatMessage } from '../../types/ollama.js'
|
||||
|
||||
/** Records every call so the tests can assert on what was *not* run. */
|
||||
function makeFakes(opts: { nomadMd?: string | null } = {}) {
|
||||
const calls = { hasDocuments: 0, search: 0, chat: 0 }
|
||||
|
||||
const ragService = {
|
||||
async hasDocuments() {
|
||||
calls.hasDocuments++
|
||||
return true
|
||||
},
|
||||
async searchSimilarDocuments() {
|
||||
calls.search++
|
||||
return [{ text: 'retrieved body', score: 0.9, metadata: { full_title: 'A Doc' } }]
|
||||
},
|
||||
}
|
||||
|
||||
const ollamaService = {
|
||||
async chat() {
|
||||
calls.chat++
|
||||
return { message: { content: 'rewritten query' } }
|
||||
},
|
||||
}
|
||||
|
||||
const nomadMdService = {
|
||||
async getSystemPrompt() {
|
||||
return opts.nomadMd ?? null
|
||||
},
|
||||
}
|
||||
|
||||
const service = new RagPipelineService(
|
||||
ollamaService as any,
|
||||
ragService as any,
|
||||
nomadMdService as any
|
||||
)
|
||||
|
||||
return { service, calls }
|
||||
}
|
||||
|
||||
const userTurn: OllamaChatMessage[] = [{ role: 'user', content: 'how do I purify water?' }]
|
||||
|
||||
const systemContents = (messages: OllamaChatMessage[]) =>
|
||||
messages.filter((m) => m.role === 'system').map((m) => m.content)
|
||||
|
||||
test.group('buildPrompt | skipRetrieval', () => {
|
||||
test('skips every expensive step when retrieval is disabled', async ({ assert }) => {
|
||||
const { service, calls } = makeFakes()
|
||||
|
||||
const trace = await service.buildPrompt(userTurn, 'llama3.1:8b', { skipRetrieval: true })
|
||||
|
||||
// The whole point of the toggle: nothing reaches Qdrant or the LLM.
|
||||
assert.equal(calls.hasDocuments, 0)
|
||||
assert.equal(calls.search, 0)
|
||||
assert.equal(calls.chat, 0)
|
||||
|
||||
// ...and no knowledge base context lands in the prompt.
|
||||
assert.deepEqual(trace.retrieved, [])
|
||||
assert.deepEqual(trace.injected, [])
|
||||
assert.isNull(trace.rewrittenQuery)
|
||||
assert.isFalse(trace.didRewrite)
|
||||
assert.isFalse(systemContents(trace.messages).some((c) => c.includes('[Context 1')))
|
||||
})
|
||||
|
||||
test('still injects the default system prompt and NOMAD.md when disabled', async ({ assert }) => {
|
||||
// Neither of these is RAG; turning retrieval off must not silently strip
|
||||
// the user's persistent instructions or the formatting prompt.
|
||||
const { service } = makeFakes({ nomadMd: 'Always answer in metric units.' })
|
||||
|
||||
const trace = await service.buildPrompt(userTurn, 'llama3.1:8b', { skipRetrieval: true })
|
||||
|
||||
const systems = systemContents(trace.messages)
|
||||
assert.include(systems, SYSTEM_PROMPTS.default)
|
||||
assert.include(systems, 'Always answer in metric units.')
|
||||
})
|
||||
|
||||
test('retrieves as normal when the option is omitted', async ({ assert }) => {
|
||||
// The eval harness never sets skipRetrieval. If this ever fails, every eval
|
||||
// run is silently scoring a no-retrieval pipeline.
|
||||
const { service, calls } = makeFakes()
|
||||
|
||||
const trace = await service.buildPrompt(userTurn, 'llama3.1:8b', {})
|
||||
|
||||
assert.equal(calls.hasDocuments, 1)
|
||||
assert.equal(calls.search, 1)
|
||||
assert.lengthOf(trace.retrieved, 1)
|
||||
assert.lengthOf(trace.injected, 1)
|
||||
assert.isTrue(systemContents(trace.messages).some((c) => c.includes('[Context 1')))
|
||||
})
|
||||
|
||||
test('oracleContext still wins over skipRetrieval', async ({ assert }) => {
|
||||
// Both are bypasses; oracle mode supplies its own context and is checked
|
||||
// first, so an eval running in oracle mode is unaffected by the setting.
|
||||
const { service, calls } = makeFakes()
|
||||
const oracle = [{ text: 'oracle body', score: 1, metadata: {} }]
|
||||
|
||||
const trace = await service.buildPrompt(userTurn, 'llama3.1:8b', {
|
||||
skipRetrieval: true,
|
||||
oracleContext: oracle,
|
||||
})
|
||||
|
||||
assert.equal(calls.search, 0)
|
||||
assert.deepEqual(trace.retrieved, oracle)
|
||||
assert.lengthOf(trace.injected, 1)
|
||||
})
|
||||
})
|
||||
|
|
@ -4,6 +4,11 @@ export const KV_STORE_SCHEMA = {
|
|||
'chat.lastModel': 'string',
|
||||
'rag.docsEmbedded': 'boolean',
|
||||
'rag.defaultIngestPolicy': 'string',
|
||||
// Master switch for chat-time knowledge base retrieval. Unset/null means ON —
|
||||
// the pre-existing behaviour. Turning it off skips the whole retrieval
|
||||
// pipeline (hasDocuments, the query-rewrite LLM call, and the Qdrant search),
|
||||
// which matters on small hardware and when the KB is small or empty.
|
||||
'rag.enabled': 'boolean',
|
||||
'system.updateAvailable': 'boolean',
|
||||
'system.latestVersion': 'string',
|
||||
'system.earlyAccess': 'boolean',
|
||||
|
|
|
|||
|
|
@ -89,6 +89,11 @@ export type PipelineOptions = {
|
|||
/** Ignore the user's NOMAD.md. Off in production; on in evals, where a
|
||||
* developer's personal instructions would silently skew every result. */
|
||||
skipNomadMd?: boolean
|
||||
/** Skip the entire retrieval pipeline — the hasDocuments check, the
|
||||
* query-rewrite LLM call and the Qdrant search — leaving the prompt with
|
||||
* system prompts only. Set from the `rag.enabled` KV setting. Opt-out by
|
||||
* design: the eval harness omits it and therefore always retrieves. */
|
||||
skipRetrieval?: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
Loading…
Reference in New Issue