Sources message section integration

This commit is contained in:
Utkarsh-Patel-13 2025-07-22 22:41:42 -07:00
parent 7c6437eef6
commit 2b647a9e7d
5 changed files with 195 additions and 8 deletions

View file

@ -60,7 +60,7 @@ class StreamingService:
self.message_annotations[1]["content"] = sources
# Return only the delta annotation
annotation = {"type": "SOURCES", "content": sources}
annotation = {"type": "SOURCES", "data": sources}
return f"8:[{json.dumps(annotation)}]\n"
def format_answer_delta(self, answer_chunk: str) -> str:

View file

@ -138,7 +138,7 @@ const ConnectorSelector = React.memo(
const [isOpen, setIsOpen] = useState(false);
const { connectorSourceItems, isLoading, isLoaded, fetchConnectors } =
useSearchSourceConnectors();
useSearchSourceConnectors(true);
const handleOpenChange = useCallback(
(open: boolean) => {

View file

@ -1,5 +1,6 @@
"use client";
import React from "react";
import {
ChatSection,
ChatHandler,
@ -12,8 +13,8 @@ import {
import { Document } from "@/hooks/use-documents";
import { CustomChatInput } from "@/components/chat_v2/ChatInputGroup";
import { ResearchMode } from "@/components/chat";
import React from "react";
import TerminalDisplay from "@/components/chat_v2/ChatTerminal";
import ChatSourcesDisplay from "@/components/chat_v2/ChatSources";
interface ChatInterfaceProps {
handler: ChatHandler;
@ -43,6 +44,7 @@ function ChatMessageDisplay({
{message.role === "assistant" ? (
<div className="flex-1 flex flex-col space-y-4">
<TerminalDisplay message={message} />
<ChatSourcesDisplay message={message} />
<ChatMessage.Content className="flex-1">
<ChatMessage.Content.Markdown />
</ChatMessage.Content>

View file

@ -0,0 +1,172 @@
"use client";
import { useState } from "react";
import { getAnnotationData, Message } from "@llamaindex/chat-ui";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogTrigger,
} from "@/components/ui/dialog";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import {
Card,
CardContent,
CardDescription,
CardHeader,
CardTitle,
} from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { ExternalLink, FileText, Github, Globe } from "lucide-react";
interface Source {
id: number;
title: string;
description: string;
url: string;
}
interface SourceGroup {
id: number;
name: string;
type: string;
sources: Source[];
}
function getSourceIcon(type: string) {
switch (type) {
case "GITHUB_CONNECTOR":
return <Github className="h-4 w-4" />;
case "NOTION_CONNECTOR":
return <FileText className="h-4 w-4" />;
case "FILE":
case "USER_SELECTED_FILE":
return <FileText className="h-4 w-4" />;
default:
return <Globe className="h-4 w-4" />;
}
}
function SourceCard({ source }: { source: Source }) {
const hasUrl = source.url && source.url.trim() !== "";
return (
<Card className="mb-3">
<CardHeader className="pb-2">
<div className="flex items-center justify-between">
<CardTitle className="text-sm font-medium truncate">
{source.title}
</CardTitle>
{hasUrl && (
<Button
variant="ghost"
size="sm"
className="h-6 w-6 p-0"
onClick={() => window.open(source.url, "_blank")}
>
<ExternalLink className="h-3 w-3" />
</Button>
)}
</div>
</CardHeader>
<CardContent className="pt-0">
<CardDescription className="text-xs line-clamp-3">
{source.description}
</CardDescription>
</CardContent>
</Card>
);
}
export default function ChatSourcesDisplay({ message }: { message: Message }) {
const [open, setOpen] = useState(false);
const annotations = getAnnotationData(message, "SOURCES");
// Flatten the nested array structure and ensure we have source groups
const sourceGroups: SourceGroup[] =
Array.isArray(annotations) && annotations.length > 0
? annotations
.flat()
.filter(
(group): group is SourceGroup =>
group !== null &&
group !== undefined &&
typeof group === "object" &&
"sources" in group &&
Array.isArray(group.sources)
)
: [];
if (sourceGroups.length === 0) {
return null;
}
const totalSources = sourceGroups.reduce(
(acc, group) => acc + group.sources.length,
0
);
return (
<Dialog open={open} onOpenChange={setOpen}>
<DialogTrigger asChild>
<Button variant="outline" size="sm" className="w-fit">
<FileText className="h-4 w-4 mr-2" />
View Sources ({totalSources})
</Button>
</DialogTrigger>
<DialogContent className="max-w-4xl h-[80vh] flex flex-col">
<DialogHeader className="flex-shrink-0">
<DialogTitle>Sources</DialogTitle>
</DialogHeader>
<Tabs
defaultValue={sourceGroups[0]?.type}
className="flex-1 flex flex-col min-h-0"
>
<TabsList
className="grid w-full flex-shrink-0"
style={{
gridTemplateColumns: `repeat(${sourceGroups.length}, 1fr)`,
}}
>
{sourceGroups.map((group) => (
<TabsTrigger
key={group.type}
value={group.type}
className="flex items-center gap-2"
>
{getSourceIcon(group.type)}
<span className="truncate">{group.name}</span>
<Badge
variant="secondary"
className="ml-1 h-5 text-xs"
>
{group.sources.length}
</Badge>
</TabsTrigger>
))}
</TabsList>
{sourceGroups.map((group) => (
<TabsContent
key={group.type}
value={group.type}
className="flex-1 min-h-0 mt-4"
>
<div className="h-full overflow-y-auto pr-2">
<div className="space-y-3">
{group.sources.map((source) => (
<SourceCard
key={source.id}
source={source}
/>
))}
</div>
</div>
</TabsContent>
))}
</Tabs>
</DialogContent>
</Dialog>
);
}

View file

@ -21,10 +21,10 @@ export interface ConnectorSourceItem {
/**
* Hook to fetch search source connectors from the API
*/
export const useSearchSourceConnectors = () => {
export const useSearchSourceConnectors = (lazy: boolean = false) => {
const [connectors, setConnectors] = useState<SearchSourceConnector[]>([]);
const [isLoading, setIsLoading] = useState(false);
const [isLoaded, setIsLoaded] = useState(false);
const [isLoading, setIsLoading] = useState(!lazy); // Don't show loading initially for lazy mode
const [isLoaded, setIsLoaded] = useState(false); // Memoization flag
const [error, setError] = useState<Error | null>(null);
const [connectorSourceItems, setConnectorSourceItems] = useState<
ConnectorSourceItem[]
@ -56,7 +56,7 @@ export const useSearchSourceConnectors = () => {
]);
const fetchConnectors = useCallback(async () => {
if (isLoaded) return; // Don't fetch if already loaded
if (isLoaded && lazy) return; // Avoid redundant calls in lazy mode
try {
setIsLoading(true);
@ -100,7 +100,19 @@ export const useSearchSourceConnectors = () => {
} finally {
setIsLoading(false);
}
}, [isLoaded]);
}, [isLoaded, lazy]);
useEffect(() => {
if (!lazy) {
fetchConnectors();
}
}, [lazy, fetchConnectors]);
// Function to refresh the connectors list
const refreshConnectors = useCallback(async () => {
setIsLoaded(false); // Reset memoization flag to allow refetch
await fetchConnectors();
}, [fetchConnectors]);
// Update connector source items when connectors change
const updateConnectorSourceItems = (
@ -363,5 +375,6 @@ export const useSearchSourceConnectors = () => {
indexConnector,
getConnectorSourceItems,
connectorSourceItems,
refreshConnectors,
};
};