diff --git a/surfsense_backend/app/agents/new_chat/llm_config.py b/surfsense_backend/app/agents/new_chat/llm_config.py index a99386df4..b6b406a38 100644 --- a/surfsense_backend/app/agents/new_chat/llm_config.py +++ b/surfsense_backend/app/agents/new_chat/llm_config.py @@ -87,10 +87,11 @@ def create_chat_litellm_from_config(llm_config: dict) -> ChatLiteLLM | None: provider_prefix = provider_map.get(provider, provider.lower()) model_string = f"{provider_prefix}/{llm_config['model_name']}" - # Create ChatLiteLLM instance + # Create ChatLiteLLM instance with streaming enabled litellm_kwargs = { "model": model_string, "api_key": llm_config.get("api_key"), + "streaming": True, # Enable streaming for real-time token streaming } # Add optional parameters diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 04f3f97c3..7f97643dc 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -12,10 +12,9 @@ from langchain_core.messages import HumanMessage from sqlalchemy.ext.asyncio import AsyncSession from app.agents.new_chat.chat_deepagent import ( - create_chat_litellm_from_config, create_surfsense_deep_agent, - load_llm_config_from_yaml, ) +from app.agents.new_chat.llm_config import create_chat_litellm_from_config, load_llm_config_from_yaml from app.services.connector_service import ConnectorService from app.services.new_streaming_service import VercelStreamingService diff --git a/surfsense_web/app/api/chat/route.ts b/surfsense_web/app/api/chat/route.ts deleted file mode 100644 index 1b8b602cb..000000000 --- a/surfsense_web/app/api/chat/route.ts +++ /dev/null @@ -1,74 +0,0 @@ -import type { UIMessage } from "ai"; - -export const maxDuration = 30; - -export async function POST(req: Request) { - try { - const body = await req.json(); - const { - messages, - chat_id, - search_space_id, - }: { - messages: UIMessage[]; - chat_id?: number; - search_space_id?: number; - } = body; - - // Get auth token from headers - const authHeader = req.headers.get("authorization"); - if (!authHeader) { - return new Response("Unauthorized", { status: 401 }); - } - - // Get the last user message - const lastUserMessage = messages.filter((m) => m.role === "user").pop(); - - if (!lastUserMessage) { - return new Response("No user message found", { status: 400 }); - } - - // Extract text content from the message - const userQuery = - typeof lastUserMessage.content === "string" - ? lastUserMessage.content - : lastUserMessage.content - .filter((c: any) => c.type === "text") - .map((c: any) => c.text) - .join(" "); - - // Call the DeepAgent backend - const backendUrl = process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"; - const response = await fetch(`${backendUrl}/api/v1/new_chat`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: authHeader, - }, - body: JSON.stringify({ - chat_id: chat_id || 0, - user_query: userQuery, - search_space_id: search_space_id || 0, - }), - }); - - if (!response.ok) { - return new Response(`Backend error: ${response.statusText}`, { - status: response.status, - }); - } - - // The backend returns SSE stream with Vercel AI SDK Data Stream Protocol - // We need to forward this stream to the client - return new Response(response.body, { - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - Connection: "keep-alive", - }, - }); - } catch (error) { - console.error("Error in deep-agent-chat route:", error); - return new Response("Internal Server Error", { status: 500 }); - } -} diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index bcaae2a10..86ed4974f 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -1,17 +1,44 @@ "use client"; -import { AssistantRuntimeProvider } from "@assistant-ui/react"; -import { useChatRuntime } from "@assistant-ui/react-ai-sdk"; +import { AssistantRuntimeProvider, useLocalRuntime } from "@assistant-ui/react"; +import { useParams } from "next/navigation"; +import { useMemo } from "react"; import { Thread } from "@/components/assistant-ui/thread"; +import { createNewChatAdapter } from "@/lib/chat/new-chat-transport"; export default function NewChatPage() { - // Using the official assistant-ui pattern - useChatRuntime with NO parameters - // It defaults to /api/chat endpoint - const runtime = useChatRuntime(); + const params = useParams(); + + // Extract search_space_id and chat_id from URL params + const searchSpaceId = useMemo(() => { + const id = params.search_space_id; + const parsed = typeof id === "string" ? Number.parseInt(id, 10) : 0; + return Number.isNaN(parsed) ? 0 : parsed; + }, [params.search_space_id]); + + const chatId = useMemo(() => { + const id = params.chat_id; + let parsed = 0; + if (Array.isArray(id) && id.length > 0) { + parsed = Number.parseInt(id[0], 10); + } else if (typeof id === "string") { + parsed = Number.parseInt(id, 10); + } + return Number.isNaN(parsed) ? 0 : parsed; + }, [params.chat_id]); + + // Create the adapter with the extracted params + const adapter = useMemo( + () => createNewChatAdapter({ searchSpaceId, chatId }), + [searchSpaceId, chatId] + ); + + // Use LocalRuntime with our custom adapter + const runtime = useLocalRuntime(adapter); return ( -
+
diff --git a/surfsense_web/components/assistant-ui/inline-citation.tsx b/surfsense_web/components/assistant-ui/inline-citation.tsx new file mode 100644 index 000000000..3f589a1b9 --- /dev/null +++ b/surfsense_web/components/assistant-ui/inline-citation.tsx @@ -0,0 +1,47 @@ +"use client"; + +import type { FC } from "react"; +import { useState } from "react"; +import { SheetTrigger } from "@/components/ui/sheet"; +import { SourceDetailSheet } from "@/components/chat/SourceDetailSheet"; + +interface InlineCitationProps { + chunkId: number; +} + +/** + * Inline citation component for the new chat. + * Renders a clickable badge that opens the SourceDetailSheet with document chunk details. + */ +export const InlineCitation: FC = ({ chunkId }) => { + const [isOpen, setIsOpen] = useState(false); + + return ( + + + + + + + + + + + ); +}; diff --git a/surfsense_web/components/assistant-ui/markdown-text.tsx b/surfsense_web/components/assistant-ui/markdown-text.tsx index 2513a0928..87043d5a4 100644 --- a/surfsense_web/components/assistant-ui/markdown-text.tsx +++ b/surfsense_web/components/assistant-ui/markdown-text.tsx @@ -9,12 +9,52 @@ import { useIsMarkdownCodeBlock, } from "@assistant-ui/react-markdown"; import { CheckIcon, CopyIcon } from "lucide-react"; -import { type FC, memo, useState } from "react"; +import { type FC, type ReactNode, memo, useState } from "react"; import remarkGfm from "remark-gfm"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { InlineCitation } from "@/components/assistant-ui/inline-citation"; import { cn } from "@/lib/utils"; +// Citation pattern: [citation:CHUNK_ID] +const CITATION_REGEX = /\[citation:(\d+)\]/g; + +/** + * Parses text and replaces [citation:XXX] patterns with InlineCitation components + */ +function parseTextWithCitations(text: string): ReactNode[] { + const parts: ReactNode[] = []; + let lastIndex = 0; + let match: RegExpExecArray | null; + let citationIndex = 0; + + // Reset regex state + CITATION_REGEX.lastIndex = 0; + + while ((match = CITATION_REGEX.exec(text)) !== null) { + // Add text before the citation + if (match.index > lastIndex) { + parts.push(text.substring(lastIndex, match.index)); + } + + // Add the citation component + const chunkId = Number.parseInt(match[1], 10); + parts.push( + + ); + + lastIndex = match.index + match[0].length; + citationIndex++; + } + + // Add any remaining text after the last citation + if (lastIndex < text.length) { + parts.push(text.substring(lastIndex)); + } + + return parts.length > 0 ? parts : [text]; +} + const MarkdownTextImpl = () => { return ( {parsed}; + } + + if (Array.isArray(children)) { + return children.map((child, index) => { + if (typeof child === "string") { + const parsed = parseTextWithCitations(child); + return parsed.length === 1 && typeof parsed[0] === "string" ? ( + child + ) : ( + {parsed} + ); + } + return child; + }); + } + + return children; +} + const defaultComponents = memoizeMarkdownComponents({ - h1: ({ className, ...props }) => ( + h1: ({ className, children, ...props }) => (

+ > + {processChildrenWithCitations(children)} +

), - h2: ({ className, ...props }) => ( + h2: ({ className, children, ...props }) => (

+ > + {processChildrenWithCitations(children)} +

), - h3: ({ className, ...props }) => ( + h3: ({ className, children, ...props }) => (

+ > + {processChildrenWithCitations(children)} +

), - h4: ({ className, ...props }) => ( + h4: ({ className, children, ...props }) => (

+ > + {processChildrenWithCitations(children)} +

), - h5: ({ className, ...props }) => ( + h5: ({ className, children, ...props }) => (
+ > + {processChildrenWithCitations(children)} +
), - h6: ({ className, ...props }) => ( -
+ h6: ({ className, children, ...props }) => ( +
+ {processChildrenWithCitations(children)} +
), - p: ({ className, ...props }) => ( -

+ p: ({ className, children, ...props }) => ( +

+ {processChildrenWithCitations(children)} +

), - a: ({ className, ...props }) => ( + a: ({ className, children, ...props }) => ( + > + {processChildrenWithCitations(children)} + ), - blockquote: ({ className, ...props }) => ( -
+ blockquote: ({ className, children, ...props }) => ( +
+ {processChildrenWithCitations(children)} +
), ul: ({ className, ...props }) => (
    li]:mt-2", className)} {...props} /> @@ -124,6 +217,11 @@ const defaultComponents = memoizeMarkdownComponents({ ol: ({ className, ...props }) => (
      li]:mt-2", className)} {...props} /> ), + li: ({ className, children, ...props }) => ( +
    1. + {processChildrenWithCitations(children)} +
      2. + ), hr: ({ className, ...props }) => (
      ), @@ -136,23 +234,27 @@ const defaultComponents = memoizeMarkdownComponents({ {...props} /> ), - th: ({ className, ...props }) => ( + th: ({ className, children, ...props }) => ( + > + {processChildrenWithCitations(children)} + ), - td: ({ className, ...props }) => ( + td: ({ className, children, ...props }) => ( + > + {processChildrenWithCitations(children)} + ), tr: ({ className, ...props }) => ( ); }, + strong: ({ className, children, ...props }) => ( + + {processChildrenWithCitations(children)} + + ), + em: ({ className, children, ...props }) => ( + + {processChildrenWithCitations(children)} + + ), CodeHeader, }); diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 26ee44872..c2dab7b3c 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -118,7 +118,7 @@ const ThreadSuggestions: FC = () => { className="aui-thread-welcome-suggestion-display fade-in slide-in-from-bottom-2 @md:nth-[n+3]:block nth-[n+3]:hidden animate-in fill-mode-both duration-200" style={{ animationDelay: `${100 + index * 50}ms` }} > - +