mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-05 05:42:39 +02:00
Merge pull request #619 from MODSetter/dev
feat: in chat document mentions
This commit is contained in:
Rohan Verma
GitHub
commit
ae971e0bbd
7 changed files with 747 additions and 55 deletions
|
|
@ -705,6 +705,7 @@ async def handle_new_chat(
|
|||
session=session,
|
||||
llm_config_id=llm_config_id,
|
||||
attachments=request.attachments,
|
||||
mentioned_document_ids=request.mentioned_document_ids,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
|
|
|
|||
|
|
@ -160,3 +160,6 @@ class NewChatRequest(BaseModel):
|
|||
attachments: list[ChatAttachment] | None = (
|
||||
None # Optional attachments with extracted content
|
||||
)
|
||||
mentioned_document_ids: list[int] | None = (
|
||||
None # Optional document IDs mentioned with @ in the chat
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from collections.abc import AsyncGenerator
|
|||
|
||||
from langchain_core.messages import HumanMessage
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent
|
||||
from app.agents.new_chat.checkpointer import get_checkpointer
|
||||
|
|
@ -24,6 +25,7 @@ from app.agents.new_chat.llm_config import (
|
|||
load_agent_config,
|
||||
load_llm_config_from_yaml,
|
||||
)
|
||||
from app.db import Document
|
||||
from app.schemas.new_chat import ChatAttachment
|
||||
from app.services.connector_service import ConnectorService
|
||||
from app.services.new_streaming_service import VercelStreamingService
|
||||
|
|
@ -46,6 +48,27 @@ def format_attachments_as_context(attachments: list[ChatAttachment]) -> str:
|
|||
return "\n".join(context_parts)
|
||||
|
||||
|
||||
def format_mentioned_documents_as_context(documents: list[Document]) -> str:
|
||||
"""Format mentioned documents as context for the agent."""
|
||||
if not documents:
|
||||
return ""
|
||||
|
||||
context_parts = ["<mentioned_documents>"]
|
||||
context_parts.append(
|
||||
"The user has explicitly mentioned the following documents from their knowledge base. "
|
||||
"These documents are directly relevant to the query and should be prioritized as primary sources."
|
||||
)
|
||||
for i, doc in enumerate(documents, 1):
|
||||
context_parts.append(
|
||||
f"<document index='{i}' id='{doc.id}' title='{doc.title}' type='{doc.document_type.value}'>"
|
||||
)
|
||||
context_parts.append(f"<![CDATA[{doc.content}]]>")
|
||||
context_parts.append("</document>")
|
||||
context_parts.append("</mentioned_documents>")
|
||||
|
||||
return "\n".join(context_parts)
|
||||
|
||||
|
||||
async def stream_new_chat(
|
||||
user_query: str,
|
||||
search_space_id: int,
|
||||
|
|
@ -53,6 +76,7 @@ async def stream_new_chat(
|
|||
session: AsyncSession,
|
||||
llm_config_id: int = -1,
|
||||
attachments: list[ChatAttachment] | None = None,
|
||||
mentioned_document_ids: list[int] | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Stream chat responses from the new SurfSense deep agent.
|
||||
|
|
@ -68,6 +92,8 @@ async def stream_new_chat(
|
|||
session: The database session
|
||||
llm_config_id: The LLM configuration ID (default: -1 for first global config)
|
||||
messages: Optional chat history from frontend (list of ChatMessage)
|
||||
attachments: Optional attachments with extracted content
|
||||
mentioned_document_ids: Optional list of document IDs mentioned with @ in the chat
|
||||
|
||||
Yields:
|
||||
str: SSE formatted response strings
|
||||
|
|
@ -136,13 +162,32 @@ async def stream_new_chat(
|
|||
# Build input with message history from frontend
|
||||
langchain_messages = []
|
||||
|
||||
# Format the user query with attachment context if any
|
||||
final_query = user_query
|
||||
if attachments:
|
||||
attachment_context = format_attachments_as_context(attachments)
|
||||
final_query = (
|
||||
f"{attachment_context}\n\n<user_query>{user_query}</user_query>"
|
||||
# Fetch mentioned documents if any
|
||||
mentioned_documents: list[Document] = []
|
||||
if mentioned_document_ids:
|
||||
result = await session.execute(
|
||||
select(Document).filter(
|
||||
Document.id.in_(mentioned_document_ids),
|
||||
Document.search_space_id == search_space_id,
|
||||
)
|
||||
)
|
||||
mentioned_documents = list(result.scalars().all())
|
||||
|
||||
# Format the user query with context (attachments + mentioned documents)
|
||||
final_query = user_query
|
||||
context_parts = []
|
||||
|
||||
if attachments:
|
||||
context_parts.append(format_attachments_as_context(attachments))
|
||||
|
||||
if mentioned_documents:
|
||||
context_parts.append(
|
||||
format_mentioned_documents_as_context(mentioned_documents)
|
||||
)
|
||||
|
||||
if context_parts:
|
||||
context = "\n\n".join(context_parts)
|
||||
final_query = f"{context}\n\n<user_query>{user_query}</user_query>"
|
||||
|
||||
# if messages:
|
||||
# # Convert frontend messages to LangChain format
|
||||
|
|
|
|||
|
|
@ -6,9 +6,16 @@ import {
|
|||
type ThreadMessageLike,
|
||||
useExternalStoreRuntime,
|
||||
} from "@assistant-ui/react";
|
||||
import { useAtomValue, useSetAtom } from "jotai";
|
||||
import { useParams, useRouter } from "next/navigation";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import {
|
||||
type MentionedDocumentInfo,
|
||||
mentionedDocumentIdsAtom,
|
||||
mentionedDocumentsAtom,
|
||||
messageDocumentsMapAtom,
|
||||
} from "@/atoms/chat/mentioned-documents.atom";
|
||||
import { Thread } from "@/components/assistant-ui/thread";
|
||||
import { ChatHeader } from "@/components/new-chat/chat-header";
|
||||
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
|
||||
|
|
@ -47,6 +54,23 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] {
|
|||
return thinkingPart?.steps || [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract mentioned documents from message content
|
||||
*/
|
||||
function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
|
||||
if (!Array.isArray(content)) return [];
|
||||
|
||||
const docsPart = content.find(
|
||||
(part: unknown) =>
|
||||
typeof part === "object" &&
|
||||
part !== null &&
|
||||
"type" in part &&
|
||||
(part as { type: string }).type === "mentioned-documents"
|
||||
) as { type: "mentioned-documents"; documents: MentionedDocumentInfo[] } | undefined;
|
||||
|
||||
return docsPart?.documents || [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert backend message to assistant-ui ThreadMessageLike format
|
||||
* Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps
|
||||
|
|
@ -57,16 +81,13 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
|
|||
if (typeof msg.content === "string") {
|
||||
content = [{ type: "text", text: msg.content }];
|
||||
} else if (Array.isArray(msg.content)) {
|
||||
// Filter out thinking-steps part - it's handled separately via messageThinkingSteps
|
||||
const filteredContent = msg.content.filter(
|
||||
(part: unknown) =>
|
||||
!(
|
||||
typeof part === "object" &&
|
||||
part !== null &&
|
||||
"type" in part &&
|
||||
(part as { type: string }).type === "thinking-steps"
|
||||
)
|
||||
);
|
||||
// Filter out custom metadata parts - they're handled separately
|
||||
const filteredContent = msg.content.filter((part: unknown) => {
|
||||
if (typeof part !== "object" || part === null || !("type" in part)) return true;
|
||||
const partType = (part as { type: string }).type;
|
||||
// Filter out thinking-steps and mentioned-documents
|
||||
return partType !== "thinking-steps" && partType !== "mentioned-documents";
|
||||
});
|
||||
content =
|
||||
filteredContent.length > 0
|
||||
? (filteredContent as ThreadMessageLike["content"])
|
||||
|
|
@ -117,6 +138,13 @@ export default function NewChatPage() {
|
|||
);
|
||||
const abortControllerRef = useRef<AbortController | null>(null);
|
||||
|
||||
// Get mentioned document IDs from the composer
|
||||
const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
|
||||
const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
|
||||
const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
|
||||
const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
|
||||
const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
|
||||
|
||||
// Create the attachment adapter for file processing
|
||||
const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []);
|
||||
|
||||
|
|
@ -154,6 +182,9 @@ export default function NewChatPage() {
|
|||
|
||||
// Extract and restore thinking steps from persisted messages
|
||||
const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
|
||||
// Extract and restore mentioned documents from persisted messages
|
||||
const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
|
||||
|
||||
for (const msg of response.messages) {
|
||||
if (msg.role === "assistant") {
|
||||
const steps = extractThinkingSteps(msg.content);
|
||||
|
|
@ -161,10 +192,19 @@ export default function NewChatPage() {
|
|||
restoredThinkingSteps.set(`msg-${msg.id}`, steps);
|
||||
}
|
||||
}
|
||||
if (msg.role === "user") {
|
||||
const docs = extractMentionedDocuments(msg.content);
|
||||
if (docs.length > 0) {
|
||||
restoredDocsMap[`msg-${msg.id}`] = docs;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (restoredThinkingSteps.size > 0) {
|
||||
setMessageThinkingSteps(restoredThinkingSteps);
|
||||
}
|
||||
if (Object.keys(restoredDocsMap).length > 0) {
|
||||
setMessageDocumentsMap(restoredDocsMap);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Create new thread
|
||||
|
|
@ -181,7 +221,7 @@ export default function NewChatPage() {
|
|||
} finally {
|
||||
setIsInitializing(false);
|
||||
}
|
||||
}, [urlChatId, searchSpaceId, router]);
|
||||
}, [urlChatId, searchSpaceId, router, setMessageDocumentsMap]);
|
||||
|
||||
// Initialize on mount
|
||||
useEffect(() => {
|
||||
|
|
@ -243,10 +283,37 @@ export default function NewChatPage() {
|
|||
};
|
||||
setMessages((prev) => [...prev, userMessage]);
|
||||
|
||||
// Persist user message (don't await, fire and forget)
|
||||
// Store mentioned documents with this message for display
|
||||
if (mentionedDocuments.length > 0) {
|
||||
const docsInfo: MentionedDocumentInfo[] = mentionedDocuments.map((doc) => ({
|
||||
id: doc.id,
|
||||
title: doc.title,
|
||||
document_type: doc.document_type,
|
||||
}));
|
||||
setMessageDocumentsMap((prev) => ({
|
||||
...prev,
|
||||
[userMsgId]: docsInfo,
|
||||
}));
|
||||
}
|
||||
|
||||
// Persist user message with mentioned documents (don't await, fire and forget)
|
||||
const persistContent =
|
||||
mentionedDocuments.length > 0
|
||||
? [
|
||||
...message.content,
|
||||
{
|
||||
type: "mentioned-documents",
|
||||
documents: mentionedDocuments.map((doc) => ({
|
||||
id: doc.id,
|
||||
title: doc.title,
|
||||
document_type: doc.document_type,
|
||||
})),
|
||||
},
|
||||
]
|
||||
: message.content;
|
||||
appendMessage(threadId, {
|
||||
role: "user",
|
||||
content: message.content,
|
||||
content: persistContent,
|
||||
}).catch((err) => console.error("Failed to persist user message:", err));
|
||||
|
||||
// Start streaming response
|
||||
|
|
@ -385,6 +452,15 @@ export default function NewChatPage() {
|
|||
// Extract attachment content to send with the request
|
||||
const attachments = extractAttachmentContent(messageAttachments);
|
||||
|
||||
// Get mentioned document IDs for context
|
||||
const documentIds = mentionedDocumentIds.length > 0 ? [...mentionedDocumentIds] : undefined;
|
||||
|
||||
// Clear mentioned documents after capturing them
|
||||
if (mentionedDocumentIds.length > 0) {
|
||||
setMentionedDocumentIds([]);
|
||||
setMentionedDocuments([]);
|
||||
}
|
||||
|
||||
const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
|
|
@ -397,6 +473,7 @@ export default function NewChatPage() {
|
|||
search_space_id: searchSpaceId,
|
||||
messages: messageHistory,
|
||||
attachments: attachments.length > 0 ? attachments : undefined,
|
||||
mentioned_document_ids: documentIds,
|
||||
}),
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
|
@ -558,7 +635,16 @@ export default function NewChatPage() {
|
|||
// Note: We no longer clear thinking steps - they persist with the message
|
||||
}
|
||||
},
|
||||
[threadId, searchSpaceId, messages]
|
||||
[
|
||||
threadId,
|
||||
searchSpaceId,
|
||||
messages,
|
||||
mentionedDocumentIds,
|
||||
mentionedDocuments,
|
||||
setMentionedDocumentIds,
|
||||
setMentionedDocuments,
|
||||
setMessageDocumentsMap,
|
||||
]
|
||||
);
|
||||
|
||||
// Convert message (pass through since already in correct format)
|
||||
|
|
|
|||
31
surfsense_web/atoms/chat/mentioned-documents.atom.ts
Normal file
31
surfsense_web/atoms/chat/mentioned-documents.atom.ts
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
"use client";
|
||||
|
||||
import { atom } from "jotai";
|
||||
import type { Document } from "@/contracts/types/document.types";
|
||||
|
||||
/**
|
||||
* Atom to store the IDs of documents mentioned in the current chat composer.
|
||||
* This is used to pass document context to the backend when sending a message.
|
||||
*/
|
||||
export const mentionedDocumentIdsAtom = atom<number[]>([]);
|
||||
|
||||
/**
|
||||
* Atom to store the full document objects mentioned in the current chat composer.
|
||||
* This persists across component remounts.
|
||||
*/
|
||||
export const mentionedDocumentsAtom = atom<Document[]>([]);
|
||||
|
||||
/**
|
||||
* Simplified document info for display purposes
|
||||
*/
|
||||
export interface MentionedDocumentInfo {
|
||||
id: number;
|
||||
title: string;
|
||||
document_type: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Atom to store mentioned documents per message ID.
|
||||
* This allows displaying which documents were mentioned with each user message.
|
||||
*/
|
||||
export const messageDocumentsMapAtom = atom<Record<string, MentionedDocumentInfo[]>>({});
|
||||
|
|
@ -8,8 +8,9 @@ import {
|
|||
ThreadPrimitive,
|
||||
useAssistantState,
|
||||
useMessage,
|
||||
useThreadViewport,
|
||||
} from "@assistant-ui/react";
|
||||
import { useAtomValue } from "jotai";
|
||||
import { useAtom, useAtomValue, useSetAtom } from "jotai";
|
||||
import {
|
||||
AlertCircle,
|
||||
ArrowDownIcon,
|
||||
|
|
@ -21,6 +22,7 @@ import {
|
|||
ChevronRightIcon,
|
||||
CopyIcon,
|
||||
DownloadIcon,
|
||||
FileText,
|
||||
Loader2,
|
||||
PencilIcon,
|
||||
Plug2,
|
||||
|
|
@ -29,10 +31,27 @@ import {
|
|||
Search,
|
||||
Sparkles,
|
||||
SquareIcon,
|
||||
X,
|
||||
} from "lucide-react";
|
||||
import Link from "next/link";
|
||||
import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useParams } from "next/navigation";
|
||||
import {
|
||||
createContext,
|
||||
type FC,
|
||||
useCallback,
|
||||
useContext,
|
||||
useEffect,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
import { createPortal } from "react-dom";
|
||||
import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon";
|
||||
import {
|
||||
mentionedDocumentIdsAtom,
|
||||
mentionedDocumentsAtom,
|
||||
messageDocumentsMapAtom,
|
||||
} from "@/atoms/chat/mentioned-documents.atom";
|
||||
import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
|
||||
import {
|
||||
globalNewLLMConfigsAtom,
|
||||
|
|
@ -49,6 +68,10 @@ import {
|
|||
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
||||
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
|
||||
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
|
||||
import {
|
||||
DocumentsDataTable,
|
||||
type DocumentsDataTableRef,
|
||||
} from "@/components/new-chat/DocumentsDataTable";
|
||||
import {
|
||||
ChainOfThought,
|
||||
ChainOfThoughtContent,
|
||||
|
|
@ -60,6 +83,7 @@ import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
|
|||
import { Button } from "@/components/ui/button";
|
||||
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
|
||||
import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
|
||||
import type { Document } from "@/contracts/types/document.types";
|
||||
import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
|
|
@ -73,8 +97,6 @@ interface ThreadProps {
|
|||
}
|
||||
|
||||
// Context to pass thinking steps to AssistantMessage
|
||||
import { createContext, useContext } from "react";
|
||||
|
||||
const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
|
||||
|
||||
/**
|
||||
|
|
@ -213,6 +235,56 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
|
|||
);
|
||||
};
|
||||
|
||||
/**
|
||||
* Component that handles auto-scroll when thinking steps update.
|
||||
* Uses useThreadViewport to scroll to bottom when thinking steps change,
|
||||
* ensuring the user always sees the latest content during streaming.
|
||||
*/
|
||||
const ThinkingStepsScrollHandler: FC = () => {
|
||||
const thinkingStepsMap = useContext(ThinkingStepsContext);
|
||||
const viewport = useThreadViewport();
|
||||
const isRunning = useAssistantState(({ thread }) => thread.isRunning);
|
||||
// Track the serialized state to detect any changes
|
||||
const prevStateRef = useRef<string>("");
|
||||
|
||||
useEffect(() => {
|
||||
// Only act during streaming
|
||||
if (!isRunning) {
|
||||
prevStateRef.current = "";
|
||||
return;
|
||||
}
|
||||
|
||||
// Serialize the thinking steps state to detect any changes
|
||||
// This catches new steps, status changes, and item additions
|
||||
let stateString = "";
|
||||
thinkingStepsMap.forEach((steps, msgId) => {
|
||||
steps.forEach((step) => {
|
||||
stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`;
|
||||
});
|
||||
});
|
||||
|
||||
// If state changed at all during streaming, scroll
|
||||
if (stateString !== prevStateRef.current && stateString !== "") {
|
||||
prevStateRef.current = stateString;
|
||||
|
||||
// Multiple attempts to ensure scroll happens after DOM updates
|
||||
const scrollAttempt = () => {
|
||||
try {
|
||||
viewport.scrollToBottom();
|
||||
} catch {
|
||||
// Ignore errors - viewport might not be ready
|
||||
}
|
||||
};
|
||||
|
||||
// Delayed attempts to handle async DOM updates
|
||||
requestAnimationFrame(scrollAttempt);
|
||||
setTimeout(scrollAttempt, 100);
|
||||
}
|
||||
}, [thinkingStepsMap, viewport, isRunning]);
|
||||
|
||||
return null; // This component doesn't render anything
|
||||
};
|
||||
|
||||
export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map(), header }) => {
|
||||
return (
|
||||
<ThinkingStepsContext.Provider value={messageThinkingSteps}>
|
||||
|
|
@ -316,12 +388,15 @@ const getTimeBasedGreeting = (userEmail?: string): string => {
|
|||
const ThreadWelcome: FC = () => {
|
||||
const { data: user } = useAtomValue(currentUserAtom);
|
||||
|
||||
// Memoize greeting so it doesn't change on re-renders (only on user change)
|
||||
const greeting = useMemo(() => getTimeBasedGreeting(user?.email), [user?.email]);
|
||||
|
||||
return (
|
||||
<div className="aui-thread-welcome-root mx-auto flex w-full max-w-(--thread-max-width) grow flex-col items-center px-4 relative">
|
||||
{/* Greeting positioned above the composer - fixed position */}
|
||||
<div className="aui-thread-welcome-message absolute bottom-[calc(50%+5rem)] left-0 right-0 flex flex-col items-center text-center">
|
||||
<h1 className="aui-thread-welcome-message-inner fade-in slide-in-from-bottom-2 animate-in text-5xl delay-100 duration-500 ease-out fill-mode-both">
|
||||
{getTimeBasedGreeting(user?.email)}
|
||||
{greeting}
|
||||
</h1>
|
||||
</div>
|
||||
{/* Composer - top edge fixed, expands downward only */}
|
||||
|
|
@ -333,42 +408,226 @@ const ThreadWelcome: FC = () => {
|
|||
};
|
||||
|
||||
const Composer: FC = () => {
|
||||
// Check if a model is configured - needed to disable input
|
||||
const { data: userConfigs } = useAtomValue(newLLMConfigsAtom);
|
||||
const { data: globalConfigs } = useAtomValue(globalNewLLMConfigsAtom);
|
||||
const { data: preferences } = useAtomValue(llmPreferencesAtom);
|
||||
// ---- State for document mentions (using atoms to persist across remounts) ----
|
||||
const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
|
||||
const [showDocumentPopover, setShowDocumentPopover] = useState(false);
|
||||
const [mentionQuery, setMentionQuery] = useState("");
|
||||
const inputRef = useRef<HTMLTextAreaElement | null>(null);
|
||||
const documentPickerRef = useRef<DocumentsDataTableRef>(null);
|
||||
const { search_space_id } = useParams();
|
||||
const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
|
||||
|
||||
const hasModelConfigured = useMemo(() => {
|
||||
if (!preferences) return false;
|
||||
const agentLlmId = preferences.agent_llm_id;
|
||||
if (agentLlmId === null || agentLlmId === undefined) return false;
|
||||
// Sync mentioned document IDs to atom for use in chat request
|
||||
useEffect(() => {
|
||||
setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
|
||||
}, [mentionedDocuments, setMentionedDocumentIds]);
|
||||
|
||||
// Check if the configured model actually exists
|
||||
if (agentLlmId < 0) {
|
||||
return globalConfigs?.some((c) => c.id === agentLlmId) ?? false;
|
||||
// Extract mention query (text after @)
|
||||
const extractMentionQuery = useCallback((value: string): string => {
|
||||
const atIndex = value.lastIndexOf("@");
|
||||
if (atIndex === -1) return "";
|
||||
return value.slice(atIndex + 1);
|
||||
}, []);
|
||||
|
||||
const handleKeyUp = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
|
||||
const textarea = e.currentTarget;
|
||||
const value = textarea.value;
|
||||
|
||||
// Open document picker when user types '@'
|
||||
if (e.key === "@" || (e.key === "2" && e.shiftKey)) {
|
||||
setShowDocumentPopover(true);
|
||||
setMentionQuery("");
|
||||
return;
|
||||
}
|
||||
return userConfigs?.some((c) => c.id === agentLlmId) ?? false;
|
||||
}, [preferences, globalConfigs, userConfigs]);
|
||||
|
||||
// Check if value contains @ and extract query
|
||||
if (value.includes("@")) {
|
||||
const query = extractMentionQuery(value);
|
||||
|
||||
// Close popup if query starts with space (user typed "@ ")
|
||||
if (query.startsWith(" ")) {
|
||||
setShowDocumentPopover(false);
|
||||
setMentionQuery("");
|
||||
return;
|
||||
}
|
||||
|
||||
// Reopen popup if @ is present and query doesn't start with space
|
||||
// (handles case where user deleted the space after @)
|
||||
if (!showDocumentPopover) {
|
||||
setShowDocumentPopover(true);
|
||||
}
|
||||
setMentionQuery(query);
|
||||
} else {
|
||||
// Close popover if '@' is no longer in the input (user deleted it)
|
||||
if (showDocumentPopover) {
|
||||
setShowDocumentPopover(false);
|
||||
setMentionQuery("");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
|
||||
// When popup is open, handle navigation keys
|
||||
if (showDocumentPopover) {
|
||||
if (e.key === "ArrowDown") {
|
||||
e.preventDefault();
|
||||
documentPickerRef.current?.moveDown();
|
||||
return;
|
||||
}
|
||||
if (e.key === "ArrowUp") {
|
||||
e.preventDefault();
|
||||
documentPickerRef.current?.moveUp();
|
||||
return;
|
||||
}
|
||||
if (e.key === "Enter") {
|
||||
e.preventDefault();
|
||||
documentPickerRef.current?.selectHighlighted();
|
||||
return;
|
||||
}
|
||||
if (e.key === "Escape") {
|
||||
e.preventDefault();
|
||||
setShowDocumentPopover(false);
|
||||
setMentionQuery("");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Remove last document chip when pressing backspace at the beginning of input
|
||||
if (e.key === "Backspace" && mentionedDocuments.length > 0) {
|
||||
const textarea = e.currentTarget;
|
||||
const selectionStart = textarea.selectionStart;
|
||||
const selectionEnd = textarea.selectionEnd;
|
||||
|
||||
// Only remove chip if cursor is at position 0 and nothing is selected
|
||||
if (selectionStart === 0 && selectionEnd === 0) {
|
||||
e.preventDefault();
|
||||
// Remove the last document chip
|
||||
setMentionedDocuments((prev) => prev.slice(0, -1));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const handleDocumentsMention = (documents: Document[]) => {
|
||||
// Update mentioned documents (merge with existing, avoid duplicates)
|
||||
setMentionedDocuments((prev) => {
|
||||
const existingIds = new Set(prev.map((d) => d.id));
|
||||
const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
|
||||
return [...prev, ...newDocs];
|
||||
});
|
||||
|
||||
// Clean up the '@...' mention text from input
|
||||
if (inputRef.current) {
|
||||
const input = inputRef.current;
|
||||
const currentValue = input.value;
|
||||
const atIndex = currentValue.lastIndexOf("@");
|
||||
|
||||
if (atIndex !== -1) {
|
||||
// Remove @ and everything after it
|
||||
const newValue = currentValue.slice(0, atIndex);
|
||||
const nativeInputValueSetter = Object.getOwnPropertyDescriptor(
|
||||
window.HTMLTextAreaElement.prototype,
|
||||
"value"
|
||||
)?.set;
|
||||
if (nativeInputValueSetter) {
|
||||
nativeInputValueSetter.call(input, newValue);
|
||||
input.dispatchEvent(new Event("input", { bubbles: true }));
|
||||
}
|
||||
}
|
||||
// Focus the input so user can continue typing
|
||||
input.focus();
|
||||
}
|
||||
|
||||
// Reset mention query
|
||||
setMentionQuery("");
|
||||
};
|
||||
|
||||
const handleRemoveDocument = (docId: number) => {
|
||||
setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId));
|
||||
};
|
||||
|
||||
return (
|
||||
<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
|
||||
<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
|
||||
<ComposerAttachments />
|
||||
<ComposerPrimitive.Input
|
||||
placeholder={
|
||||
hasModelConfigured
|
||||
? "Ask SurfSense"
|
||||
: "Select a model from the header to start chatting..."
|
||||
}
|
||||
className={cn(
|
||||
"aui-composer-input mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0",
|
||||
!hasModelConfigured && "cursor-not-allowed opacity-60"
|
||||
{/* -------- Input field with inline document chips -------- */}
|
||||
<div className="aui-composer-input-wrapper flex flex-wrap items-center gap-1.5 px-3 pt-2 pb-6">
|
||||
{/* Inline document chips */}
|
||||
{mentionedDocuments.map((doc) => (
|
||||
<span
|
||||
key={doc.id}
|
||||
className="inline-flex items-center gap-1 pl-2 pr-1 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20 shrink-0"
|
||||
title={doc.title}
|
||||
>
|
||||
<span className="max-w-[120px] truncate">{doc.title}</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => handleRemoveDocument(doc.id)}
|
||||
className="size-4 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors"
|
||||
aria-label={`Remove ${doc.title}`}
|
||||
>
|
||||
<X className="size-3" />
|
||||
</button>
|
||||
</span>
|
||||
))}
|
||||
{/* Text input */}
|
||||
<ComposerPrimitive.Input
|
||||
ref={inputRef}
|
||||
onKeyUp={handleKeyUp}
|
||||
onKeyDown={handleKeyDown}
|
||||
placeholder={
|
||||
mentionedDocuments.length > 0
|
||||
? "Ask about these documents..."
|
||||
: "Ask SurfSense (type @ to mention docs)"
|
||||
}
|
||||
className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1"
|
||||
rows={1}
|
||||
autoFocus
|
||||
aria-label="Message input"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* -------- Document mention popover (rendered via portal) -------- */}
|
||||
{showDocumentPopover &&
|
||||
typeof document !== "undefined" &&
|
||||
createPortal(
|
||||
<>
|
||||
{/* Backdrop */}
|
||||
<button
|
||||
type="button"
|
||||
className="fixed inset-0 cursor-default"
|
||||
style={{ zIndex: 9998 }}
|
||||
onClick={() => setShowDocumentPopover(false)}
|
||||
aria-label="Close document picker"
|
||||
/>
|
||||
{/* Popover positioned above input */}
|
||||
<div
|
||||
className="fixed shadow-2xl rounded-lg border border-border overflow-hidden"
|
||||
style={{
|
||||
zIndex: 9999,
|
||||
backgroundColor: "#18181b",
|
||||
bottom: inputRef.current
|
||||
? `${window.innerHeight - inputRef.current.getBoundingClientRect().top + 8}px`
|
||||
: "200px",
|
||||
left: inputRef.current
|
||||
? `${inputRef.current.getBoundingClientRect().left}px`
|
||||
: "50%",
|
||||
}}
|
||||
>
|
||||
<DocumentsDataTable
|
||||
ref={documentPickerRef}
|
||||
searchSpaceId={Number(search_space_id)}
|
||||
onSelectionChange={handleDocumentsMention}
|
||||
onDone={() => {
|
||||
setShowDocumentPopover(false);
|
||||
setMentionQuery("");
|
||||
}}
|
||||
initialSelectedDocuments={mentionedDocuments}
|
||||
externalSearch={mentionQuery}
|
||||
/>
|
||||
</div>
|
||||
</>,
|
||||
document.body
|
||||
)}
|
||||
rows={1}
|
||||
autoFocus={hasModelConfigured}
|
||||
aria-label="Message input"
|
||||
disabled={!hasModelConfigured}
|
||||
/>
|
||||
<ComposerAction />
|
||||
</ComposerPrimitive.AttachmentDropzone>
|
||||
</ComposerPrimitive.Root>
|
||||
|
|
@ -469,7 +728,7 @@ const ConnectorIndicator: FC = () => {
|
|||
</div>
|
||||
<div className="flex flex-wrap gap-2">
|
||||
{/* Document types from the search space */}
|
||||
{activeDocumentTypes.map(([docType, count]) => (
|
||||
{activeDocumentTypes.map(([docType]) => (
|
||||
<div
|
||||
key={docType}
|
||||
className="flex items-center gap-1.5 rounded-md bg-muted/80 px-2.5 py-1.5 text-xs border border-border/50"
|
||||
|
|
@ -642,7 +901,7 @@ const ThinkingStepsPart: FC = () => {
|
|||
const thinkingStepsMap = useContext(ThinkingStepsContext);
|
||||
|
||||
// Get the current message ID to look up thinking steps
|
||||
const messageId = useMessage((m) => m.id);
|
||||
const messageId = useAssistantState(({ message }) => message?.id);
|
||||
const thinkingSteps = thinkingStepsMap.get(messageId) || [];
|
||||
|
||||
// Check if thread is still running (for stopping the spinner when cancelled)
|
||||
|
|
@ -725,6 +984,10 @@ const AssistantActionBar: FC = () => {
|
|||
};
|
||||
|
||||
const UserMessage: FC = () => {
|
||||
const messageId = useAssistantState(({ message }) => message?.id);
|
||||
const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
|
||||
const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
|
||||
|
||||
return (
|
||||
<MessagePrimitive.Root
|
||||
className="aui-user-message-root fade-in slide-in-from-bottom-1 mx-auto grid w-full max-w-(--thread-max-width) animate-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] content-start gap-y-2 px-2 py-3 duration-150 [&:where(>*)]:col-start-2"
|
||||
|
|
@ -733,10 +996,25 @@ const UserMessage: FC = () => {
|
|||
<UserMessageAttachments />
|
||||
|
||||
<div className="aui-user-message-content-wrapper relative col-start-2 min-w-0">
|
||||
{/* Display mentioned documents as chips */}
|
||||
{mentionedDocs && mentionedDocs.length > 0 && (
|
||||
<div className="flex flex-wrap gap-1.5 mb-2 justify-end">
|
||||
{mentionedDocs.map((doc) => (
|
||||
<span
|
||||
key={doc.id}
|
||||
className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
|
||||
title={doc.title}
|
||||
>
|
||||
<FileText className="size-3" />
|
||||
<span className="max-w-[150px] truncate">{doc.title}</span>
|
||||
</span>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
|
||||
<MessagePrimitive.Parts />
|
||||
</div>
|
||||
<div className="aui-user-action-bar-wrapper -translate-x-full -translate-y-1/2 absolute top-1/2 left-0 pr-2">
|
||||
<div className="aui-user-action-bar-wrapper -translate-x-full -translate-y-full absolute top-full left-0 pr-2">
|
||||
<UserActionBar />
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
248
surfsense_web/components/new-chat/DocumentsDataTable.tsx
Normal file
248
surfsense_web/components/new-chat/DocumentsDataTable.tsx
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
"use client";
|
||||
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { FileText } from "lucide-react";
|
||||
import {
|
||||
forwardRef,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useImperativeHandle,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
|
||||
import type { Document } from "@/contracts/types/document.types";
|
||||
import { documentsApiService } from "@/lib/apis/documents-api.service";
|
||||
import { cacheKeys } from "@/lib/query-client/cache-keys";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
export interface DocumentsDataTableRef {
|
||||
selectHighlighted: () => void;
|
||||
moveUp: () => void;
|
||||
moveDown: () => void;
|
||||
}
|
||||
|
||||
interface DocumentsDataTableProps {
|
||||
searchSpaceId: number;
|
||||
onSelectionChange: (documents: Document[]) => void;
|
||||
onDone: () => void;
|
||||
initialSelectedDocuments?: Document[];
|
||||
externalSearch?: string;
|
||||
}
|
||||
|
||||
function useDebounced<T>(value: T, delay = 300) {
|
||||
const [debounced, setDebounced] = useState(value);
|
||||
useEffect(() => {
|
||||
const t = setTimeout(() => setDebounced(value), delay);
|
||||
return () => clearTimeout(t);
|
||||
}, [value, delay]);
|
||||
return debounced;
|
||||
}
|
||||
|
||||
export const DocumentsDataTable = forwardRef<DocumentsDataTableRef, DocumentsDataTableProps>(
|
||||
function DocumentsDataTable(
|
||||
{
|
||||
searchSpaceId,
|
||||
onSelectionChange,
|
||||
onDone,
|
||||
initialSelectedDocuments = [],
|
||||
externalSearch = "",
|
||||
},
|
||||
ref
|
||||
) {
|
||||
// Use external search
|
||||
const search = externalSearch;
|
||||
const debouncedSearch = useDebounced(search, 150);
|
||||
const [highlightedIndex, setHighlightedIndex] = useState(0);
|
||||
const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
|
||||
|
||||
const fetchQueryParams = useMemo(
|
||||
() => ({
|
||||
search_space_id: searchSpaceId,
|
||||
page: 0,
|
||||
page_size: 20,
|
||||
}),
|
||||
[searchSpaceId]
|
||||
);
|
||||
|
||||
const searchQueryParams = useMemo(() => {
|
||||
return {
|
||||
search_space_id: searchSpaceId,
|
||||
page: 0,
|
||||
page_size: 20,
|
||||
title: debouncedSearch,
|
||||
};
|
||||
}, [debouncedSearch, searchSpaceId]);
|
||||
|
||||
// Use query for fetching documents
|
||||
const { data: documents, isLoading: isDocumentsLoading } = useQuery({
|
||||
queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
|
||||
queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
|
||||
staleTime: 3 * 60 * 1000,
|
||||
enabled: !!searchSpaceId && !debouncedSearch.trim(),
|
||||
});
|
||||
|
||||
// Searching
|
||||
const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
|
||||
queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
|
||||
queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
|
||||
staleTime: 3 * 60 * 1000,
|
||||
enabled: !!searchSpaceId && !!debouncedSearch.trim(),
|
||||
});
|
||||
|
||||
const actualDocuments = debouncedSearch.trim()
|
||||
? searchedDocuments?.items || []
|
||||
: documents?.items || [];
|
||||
const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
|
||||
|
||||
// Track already selected document IDs
|
||||
const selectedIds = useMemo(
|
||||
() => new Set(initialSelectedDocuments.map((d) => d.id)),
|
||||
[initialSelectedDocuments]
|
||||
);
|
||||
|
||||
// Filter out already selected documents for navigation
|
||||
const selectableDocuments = useMemo(
|
||||
() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
|
||||
[actualDocuments, selectedIds]
|
||||
);
|
||||
|
||||
const handleSelectDocument = useCallback(
|
||||
(doc: Document) => {
|
||||
onSelectionChange([...initialSelectedDocuments, doc]);
|
||||
onDone();
|
||||
},
|
||||
[initialSelectedDocuments, onSelectionChange, onDone]
|
||||
);
|
||||
|
||||
// Scroll highlighted item into view
|
||||
useEffect(() => {
|
||||
const item = itemRefs.current.get(highlightedIndex);
|
||||
if (item) {
|
||||
item.scrollIntoView({ block: "nearest", behavior: "smooth" });
|
||||
}
|
||||
}, [highlightedIndex]);
|
||||
|
||||
// Reset highlighted index when external search changes
|
||||
const prevSearchRef = useRef(search);
|
||||
if (prevSearchRef.current !== search) {
|
||||
prevSearchRef.current = search;
|
||||
if (highlightedIndex !== 0) {
|
||||
setHighlightedIndex(0);
|
||||
}
|
||||
}
|
||||
|
||||
// Expose methods to parent via ref
|
||||
useImperativeHandle(
|
||||
ref,
|
||||
() => ({
|
||||
selectHighlighted: () => {
|
||||
if (selectableDocuments[highlightedIndex]) {
|
||||
handleSelectDocument(selectableDocuments[highlightedIndex]);
|
||||
}
|
||||
},
|
||||
moveUp: () => {
|
||||
setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
|
||||
},
|
||||
moveDown: () => {
|
||||
setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
|
||||
},
|
||||
}),
|
||||
[selectableDocuments, highlightedIndex, handleSelectDocument]
|
||||
);
|
||||
|
||||
// Handle keyboard navigation
|
||||
const handleKeyDown = useCallback(
|
||||
(e: React.KeyboardEvent) => {
|
||||
if (selectableDocuments.length === 0) return;
|
||||
|
||||
switch (e.key) {
|
||||
case "ArrowDown":
|
||||
e.preventDefault();
|
||||
setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
|
||||
break;
|
||||
case "ArrowUp":
|
||||
e.preventDefault();
|
||||
setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
|
||||
break;
|
||||
case "Enter":
|
||||
e.preventDefault();
|
||||
if (selectableDocuments[highlightedIndex]) {
|
||||
handleSelectDocument(selectableDocuments[highlightedIndex]);
|
||||
}
|
||||
break;
|
||||
case "Escape":
|
||||
e.preventDefault();
|
||||
onDone();
|
||||
break;
|
||||
}
|
||||
},
|
||||
[selectableDocuments, highlightedIndex, handleSelectDocument, onDone]
|
||||
);
|
||||
|
||||
return (
|
||||
<div
|
||||
className="flex flex-col w-[280px] sm:w-[320px] bg-zinc-900 rounded-lg"
|
||||
onKeyDown={handleKeyDown}
|
||||
role="listbox"
|
||||
tabIndex={-1}
|
||||
>
|
||||
{/* Document List */}
|
||||
<div className="max-h-[280px] overflow-y-auto">
|
||||
{actualLoading ? (
|
||||
<div className="flex items-center justify-center py-4">
|
||||
<div className="animate-spin h-5 w-5 border-2 border-primary border-t-transparent rounded-full" />
|
||||
</div>
|
||||
) : actualDocuments.length === 0 ? (
|
||||
<div className="flex flex-col items-center justify-center py-4 text-center px-4">
|
||||
<FileText className="h-5 w-5 text-muted-foreground/50 mb-1" />
|
||||
<p className="text-sm text-muted-foreground">No documents found</p>
|
||||
</div>
|
||||
) : (
|
||||
<div className="py-1">
|
||||
{actualDocuments.map((doc) => {
|
||||
const isAlreadySelected = selectedIds.has(doc.id);
|
||||
const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id);
|
||||
const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex;
|
||||
|
||||
return (
|
||||
<button
|
||||
key={doc.id}
|
||||
ref={(el) => {
|
||||
if (el && selectableIndex >= 0) {
|
||||
itemRefs.current.set(selectableIndex, el);
|
||||
}
|
||||
}}
|
||||
type="button"
|
||||
onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
|
||||
onMouseEnter={() => {
|
||||
if (!isAlreadySelected && selectableIndex >= 0) {
|
||||
setHighlightedIndex(selectableIndex);
|
||||
}
|
||||
}}
|
||||
disabled={isAlreadySelected}
|
||||
className={cn(
|
||||
"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
|
||||
isAlreadySelected ? "opacity-50 cursor-not-allowed" : "cursor-pointer",
|
||||
isHighlighted && "bg-accent"
|
||||
)}
|
||||
>
|
||||
{/* Type icon */}
|
||||
<span className="flex-shrink-0 text-muted-foreground text-sm">
|
||||
{getConnectorIcon(doc.document_type)}
|
||||
</span>
|
||||
{/* Title */}
|
||||
<span className="flex-1 text-sm truncate" title={doc.title}>
|
||||
{doc.title}
|
||||
</span>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
);
|
||||
Loading…
Add table
Add a link
Reference in a new issue