Skip to content

feat: integrate RAG environment selection in chat window #1227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jeffmaury wants to merge 1 commit into
base: main
Choose a base branch
from
+163 −1
Open

feat: integrate RAG environment selection in chat window #1227

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 25 additions & 1 deletion packages/renderer/src/lib/chat/components/chat-header.svelte
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ import ModelSelector from '/@/lib/chat/components/model-selector.svelte';
import { currentChatId } from '/@/lib/chat/state/current-chat-id.svelte';
import { cn } from '/@/lib/chat/utils/shadcn';
import { mcpRemoteServerInfos } from '/@/stores/mcp-remote-servers';
import type { MCPRemoteServerInfo } from '/@api/mcp/mcp-server-info';
import type { RagEnvironment } from '/@api/rag/rag-environment';

import Plus from './icons/plus.svelte';
import RagEnvironmentSelector from './rag-environment-selector.svelte';
import SidebarToggle from './sidebar-toggle.svelte';
import { Button } from './ui/button';
import { useSidebar } from './ui/sidebar';
Expand All @@ -20,11 +23,15 @@ let {
readonly,
models,
selectedModel = $bindable<ModelInfo | undefined>(),
onMCPServerAdd,
onMCPServerRemove,
selectedMCPToolsCount,
mcpSelectorOpen = $bindable(),
}: {
readonly: boolean;
selectedModel: ModelInfo | undefined;
onMCPServerAdd: (mcpServer: MCPRemoteServerInfo) => void;
onMCPServerRemove: (mcpServer: MCPRemoteServerInfo) => void;
models: Array<ModelInfo>;
/**
* Represent the number of tools selected
Expand All @@ -40,6 +47,18 @@ const noMcps = $derived($mcpRemoteServerInfos.length === 0);
function onToolSelection(): void {
mcpSelectorOpen = true;
}

let selectedRagEnvironment = $state<RagEnvironment | undefined>(undefined);

function onSelectRagEnvironment(ragEnvironment: RagEnvironment | undefined): void {
if (selectedRagEnvironment?.mcpServer !== undefined) {
onMCPServerRemove(selectedRagEnvironment?.mcpServer);
}
if (ragEnvironment?.mcpServer !== undefined) {
onMCPServerAdd(ragEnvironment.mcpServer);
}
selectedRagEnvironment = ragEnvironment;
}
</script>

<header class="bg-background sticky top-0 flex items-start gap-2 p-2">
Expand Down Expand Up @@ -77,7 +96,12 @@ function onToolSelection(): void {
models={models}
bind:value={selectedModel}
/>
<div class="flex flex-col gap-1">
<RagEnvironmentSelector
class="order-2 md:order-3"
onSelect={onSelectRagEnvironment}
/>

<div class="flex flex-col gap-1">
{#if noMcps}
<div class="flex items-center gap-1 px-1 text-xs text-muted-foreground">
<Button
Expand Down
14 changes: 14 additions & 0 deletions packages/renderer/src/lib/chat/components/chat.svelte
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import { mcpRemoteServerInfos } from '/@/stores/mcp-remote-servers';
import { providerInfos } from '/@/stores/providers';
import { MessageConfigSchema } from '/@api/chat/message-config';
import type { Chat as DbChat, Message as DbMessage } from '/@api/chat/schema.js';
import type { MCPRemoteServerInfo } from '/@api/mcp/mcp-server-info';

import ChatHeader from './chat-header.svelte';
import { IPCChatTransport } from './ipc-chat-transport';
Expand Down Expand Up @@ -82,6 +83,17 @@ const selectedMCPTools = new SvelteMap<string, SvelteSet<string>>(
}, new Map<string, SvelteSet<string>>()),
);

function onMCPServerAdd(mcpServer: MCPRemoteServerInfo): void {
const server = $mcpRemoteServerInfos.find(r => r.id === mcpServer.id);
if (server) {
selectedMCPTools.set(mcpServer.id, new SvelteSet(Object.keys(server.tools)));
}
Comment on lines +86 to +90
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Use the selected RAG MCP payload as fallback when syncing tools.

At Line 87-90, selection is applied only if the server is found in $mcpRemoteServerInfos. If that store is temporarily stale/empty, selecting a RAG environment silently does nothing and tool count won’t update.

Proposed fix 
 function onMCPServerAdd(mcpServer: MCPRemoteServerInfo): void {
-  const server = $mcpRemoteServerInfos.find(r => r.id === mcpServer.id);
-  if (server) {
-    selectedMCPTools.set(mcpServer.id, new SvelteSet(Object.keys(server.tools)));
-  }
+  const server = $mcpRemoteServerInfos.find(r => r.id === mcpServer.id) ?? mcpServer;
+  selectedMCPTools.set(mcpServer.id, new SvelteSet(Object.keys(server.tools)));
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
function onMCPServerAdd(mcpServer: MCPRemoteServerInfo): void {
const server = $mcpRemoteServerInfos.find(r => r.id === mcpServer.id);
if (server) {
selectedMCPTools.set(mcpServer.id, new SvelteSet(Object.keys(server.tools)));
}
function onMCPServerAdd(mcpServer: MCPRemoteServerInfo): void {
const server = $mcpRemoteServerInfos.find(r => r.id === mcpServer.id) ?? mcpServer;
selectedMCPTools.set(mcpServer.id, new SvelteSet(Object.keys(server.tools)));
}
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@packages/renderer/src/lib/chat/components/chat.svelte` around lines 86 - 90,
The onMCPServerAdd handler currently only updates selectedMCPTools when a
matching server exists in $mcpRemoteServerInfos, so selections get dropped if
that store is stale; change onMCPServerAdd to always call
selectedMCPTools.set(mcpServer.id, new SvelteSet(...)) using
Object.keys(server.tools) when server is found and falling back to
Object.keys(mcpServer.tools) (the incoming payload) when not found, keeping the
SvelteSet usage and preserving behavior when $mcpRemoteServerInfos is empty.

}

function onMCPServerRemove(mcpServer: MCPRemoteServerInfo): void {
selectedMCPTools.delete(mcpServer.id);
}

const selectedMCPToolsCount = $derived(
selectedMCPTools.entries().reduce((acc, [, tools]) => {
return acc + tools.size;
Expand Down Expand Up @@ -163,6 +175,8 @@ function onCheckMCPTool(mcpId: string, toolId: string, checked: boolean): void {
bind:mcpSelectorOpen={mcpSelectorOpen}
{readonly}
models={models}
onMCPServerAdd={onMCPServerAdd}
onMCPServerRemove={onMCPServerRemove}
selectedMCPToolsCount={selectedMCPToolsCount}
bind:selectedModel={selectedModel}
/>
Expand Down
124 changes: 124 additions & 0 deletions packages/renderer/src/lib/chat/components/rag-environment-selector.svelte
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
<script lang="ts">
import { faBook } from '@fortawesome/free-solid-svg-icons/faBook';
import Fa from 'svelte-fa';

import { cn } from '/@/lib/chat/utils/shadcn';
import { ragEnvironments } from '/@/stores/rag-environments';
import type { RagEnvironment } from '/@api/rag/rag-environment';

import CheckCircleFillIcon from './icons/check-circle-fill.svelte';
import ChevronDownIcon from './icons/chevron-down.svelte';
import { Button } from './ui/button';
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuGroup,
DropdownMenuItem,
DropdownMenuLabel,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from './ui/dropdown-menu';

interface Props {
class?: string;
disabled?: boolean;
open?: boolean;
onSelect: (env: RagEnvironment | undefined) => void;
}

let { class: className, disabled = false, open = $bindable(false), onSelect }: Props = $props();

// Filter RAG environments to only show those with an MCP server
const ragEnvironmentsWithMCP = $derived($ragEnvironments.filter(env => env.mcpServer !== undefined));

let selected = $state<RagEnvironment | undefined>(undefined);

const selectedText = $derived(selected?.name ?? 'No Knowledge Base');

function onSelectRagEnvironment(env: RagEnvironment | undefined, event: Event): void {
event.preventDefault(); // prevent dropdown from closing itself
onSelect(env);
selected = env;
open = false; // close the dropdown after selection
}
Comment on lines +34 to +43
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Reset invalid selection when RAG environments change.

Line 34-43 keeps selected even if the environment disappears/loses MCP after a store refresh, leaving stale label/state and potentially stale active MCP tools.

Proposed fix 
 let selected = $state<RagEnvironment | undefined>(undefined);

 const selectedText = $derived(selected?.name ?? 'No Knowledge Base');
+
+$effect(() => {
+  if (!selected) return;
+  const stillAvailable = ragEnvironmentsWithMCP.some(
+    env =>
+      env.name === selected.name &&
+      env.ragConnection.providerId === selected.ragConnection.providerId &&
+      env.ragConnection.name === selected.ragConnection.name,
+  );
+  if (!stillAvailable) {
+    selected = undefined;
+    onSelect(undefined);
+  }
+});
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
let selected = $state<RagEnvironment | undefined>(undefined);
const selectedText = $derived(selected?.name ?? 'No Knowledge Base');
function onSelectRagEnvironment(env: RagEnvironment | undefined, event: Event): void {
event.preventDefault(); // prevent dropdown from closing itself
onSelect(env);
selected = env;
open = false; // close the dropdown after selection
}
let selected = $state<RagEnvironment | undefined>(undefined);
const selectedText = $derived(selected?.name ?? 'No Knowledge Base');
$effect(() => {
if (!selected) return;
const stillAvailable = ragEnvironmentsWithMCP.some(
env =>
env.name === selected.name &&
env.ragConnection.providerId === selected.ragConnection.providerId &&
env.ragConnection.name === selected.ragConnection.name,
);
if (!stillAvailable) {
selected = undefined;
onSelect(undefined);
}
});
function onSelectRagEnvironment(env: RagEnvironment | undefined, event: Event): void {
event.preventDefault(); // prevent dropdown from closing itself
onSelect(env);
selected = env;
open = false; // close the dropdown after selection
}
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@packages/renderer/src/lib/chat/components/rag-environment-selector.svelte`
around lines 34 - 43, The selected RagEnvironment must be cleared when the
available RAG environments change to avoid stale labels/tools; add a
watcher/subscription to the environments store (the list used to populate the
dropdown) that checks whether the current selected (selected) still exists and
still has an active MCP/tool, and if not set selected = undefined and call
onSelect(undefined) to clear active MCPs; implement this as a reactive statement
or store subscription that compares selected.id or selected.name against the new
environments list and resets selected and the selection handler accordingly so
selectedText and active tools cannot remain stale.

</script>

<DropdownMenu {open} onOpenChange={(val): boolean => (open = val)}>
<DropdownMenuTrigger>
{#snippet child({ props })}
<Button
{...props}
aria-label="Select knowledge base"
variant="outline"
{disabled}
class={cn('data-[state=open]:bg-accent data-[state=open]:text-accent-foreground w-fit md:h-[34px] md:px-2', className)}
>
<Fa icon={faBook} />
{selectedText}
<ChevronDownIcon />
</Button>
{/snippet}
</DropdownMenuTrigger>
<DropdownMenuContent align="start" class="min-w-[300px]">
<DropdownMenuLabel>
<div class="flex flex-col gap-1">
<div class="font-semibold">Select Knowledge Base</div>
<div class="text-xs font-normal text-muted-foreground">Choose a knowledge environment to enhance your chat</div>
</div>
</DropdownMenuLabel>
<DropdownMenuSeparator />

{#if ragEnvironmentsWithMCP.length === 0}
<DropdownMenuItem
disabled
class="group/item flex flex-row items-center justify-between gap-4"
>
No RAG environments with MCP available
</DropdownMenuItem>
{:else}
<DropdownMenuGroup>
<!-- No Knowledge Base option -->
<DropdownMenuItem
onSelect={onSelectRagEnvironment.bind(undefined, undefined)}
class="group/item flex flex-row items-center justify-between gap-4"
data-active={selected === undefined}
>
<div class="flex flex-col items-start gap-1">
<div class="font-medium">No Knowledge Base</div>
<div class="text-xs text-muted-foreground">Chat without RAG enhancement</div>
</div>

<div
class="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100"
>
<CheckCircleFillIcon />
</div>
</DropdownMenuItem>

<DropdownMenuSeparator />

<!-- RAG Environment options -->
{#each ragEnvironmentsWithMCP as ragEnv (ragEnv.name)}
<DropdownMenuItem
onSelect={onSelectRagEnvironment.bind(undefined, ragEnv)}
class="group/item flex flex-row items-center justify-between gap-4"
data-active={selected?.name === ragEnv.name}
>
<div class="flex flex-col items-start gap-1">
<div class="font-medium">{ragEnv.name}</div>
<div class="text-xs text-muted-foreground">
{ragEnv.ragConnection.providerId} • {ragEnv.files.length} sources
</div>
</div>

<div
class="text-foreground dark:text-foreground opacity-0 group-data-[active=true]/item:opacity-100"
>
<CheckCircleFillIcon />
</div>
</DropdownMenuItem>
{/each}
</DropdownMenuGroup>
{/if}
</DropdownMenuContent>
</DropdownMenu>
Loading