feat: multiple models & ui improvements (#455)

8 files changed, 253 insertions, 29 deletions

diff --git a/apps/web/app/(navigation)/settings/support/page.tsx b/apps/web/app/(navigation)/settings/support/page.tsx
index de4b3c78..7dc2a207 100644
--- a/apps/web/app/(navigation)/settings/support/page.tsx
+++ b/apps/web/app/(navigation)/settings/support/page.tsx
@@ -69,7 +69,7 @@ export default function SupportPage() {
 								What's included in the Pro plan?
 							</h4>
 							<p className="text-muted-foreground text-sm">
-								Pro includes 5,000 memories (vs 200 in free), 10 connections to
+								Pro includes unlimited memories (vs 200 in free), 10 connections to
 								external services like Google Drive and Notion, advanced search
 								features, and priority support.
 							</p>
diff --git a/apps/web/components/chat-input.tsx b/apps/web/components/chat-input.tsx
index 8a6edbf8..38f89815 100644
--- a/apps/web/components/chat-input.tsx
+++ b/apps/web/components/chat-input.tsx
@@ -1,20 +1,44 @@
 "use client"
 
-import { useState } from "react"
+import { useState, useEffect } from "react"
 import { useRouter } from "next/navigation"
 import { generateId } from "@lib/generate-id"
 import { usePersistentChat } from "@/stores/chat"
 import { ArrowUp } from "lucide-react"
 import { Button } from "@ui/components/button"
 import { ProjectSelector } from "./project-selector"
+import { ModelSelector } from "./model-selector"
 import { useAuth } from "@lib/auth-context"
 
 export function ChatInput() {
 	const [message, setMessage] = useState("")
+	const [selectedModel, setSelectedModel] = useState<
+		"gpt-5" | "claude-sonnet-4.5" | "gemini-2.5-pro"
+	>("gemini-2.5-pro")
 	const router = useRouter()
 	const { setCurrentChatId } = usePersistentChat()
 	const { user } = useAuth()
 
+	useEffect(() => {
+		const savedModel = localStorage.getItem("selectedModel") as
+			| "gpt-5"
+			| "claude-sonnet-4.5"
+			| "gemini-2.5-pro"
+		if (
+			savedModel &&
+			["gpt-5", "claude-sonnet-4.5", "gemini-2.5-pro"].includes(savedModel)
+		) {
+			setSelectedModel(savedModel)
+		}
+	}, [])
+
+	const handleModelChange = (
+		modelId: "gpt-5" | "claude-sonnet-4.5" | "gemini-2.5-pro",
+	) => {
+		setSelectedModel(modelId)
+		localStorage.setItem("selectedModel", modelId)
+	}
+
 	const handleSend = () => {
 		if (!message.trim()) return
 
@@ -22,8 +46,8 @@ export function ChatInput() {
 
 		setCurrentChatId(newChatId)
 
-		// Store the initial message in sessionStorage for the chat page to pick up
 		sessionStorage.setItem(`chat-initial-${newChatId}`, message.trim())
+		sessionStorage.setItem(`chat-model-${newChatId}`, selectedModel)
 
 		router.push(`/chat/${newChatId}`)
 
@@ -64,15 +88,21 @@ export function ChatInput() {
 						/>
 						<div className="flex items-center gap-2 w-full justify-between py-2 px-3 rounded-b-[14px]">
 							<ProjectSelector />
-							<Button
-								onClick={handleSend}
-								disabled={!message.trim()}
-								className="text-primary-foreground border-0 rounded-xl transition-colors disabled:opacity-50 disabled:cursor-not-allowed !bg-primary h-8 w-8"
-								variant="outline"
-								size="icon"
-							>
-								<ArrowUp className="size-3.5" />
-							</Button>
+							<div className="flex items-center gap-2">
+								<ModelSelector
+									selectedModel={selectedModel}
+									onModelChange={handleModelChange}
+								/>
+								<Button
+									onClick={handleSend}
+									disabled={!message.trim()}
+									className="text-primary-foreground border-0 rounded-xl transition-colors disabled:opacity-50 disabled:cursor-not-allowed !bg-primary h-8 w-8"
+									variant="outline"
+									size="icon"
+								>
+									<ArrowUp className="size-3.5" />
+								</Button>
+							</div>
 						</div>
 					</form>
 				</div>
diff --git a/apps/web/components/chrome-extension-button.tsx b/apps/web/components/chrome-extension-button.tsx
index 8f2b74c8..3f7510e3 100644
--- a/apps/web/components/chrome-extension-button.tsx
+++ b/apps/web/components/chrome-extension-button.tsx
@@ -26,7 +26,11 @@ export function ChromeExtensionButton() {
 	useEffect(() => {
 		const dismissed =
 			localStorage.getItem("chrome-extension-dismissed") === "true"
+		const minimized =
+			localStorage.getItem("chrome-extension-minimized") === "true"
+
 		setIsDismissed(dismissed)
+		setIsMinimized(minimized)
 
 		const checkExtension = () => {
 			const message = { action: "check-extension" }
@@ -35,9 +39,10 @@ export function ChromeExtensionButton() {
 				setIsExtensionInstalled(false)
 				setIsChecking(false)
 				// Auto-minimize after 3 seconds if extension is not installed and not dismissed
-				if (!dismissed) {
+				if (!dismissed && !minimized) {
 					setTimeout(() => {
 						setIsMinimized(true)
+						localStorage.setItem("chrome-extension-minimized", "true")
 					}, 3000)
 				}
 			}, 1000)
@@ -79,6 +84,7 @@ export function ChromeExtensionButton() {
 
 	const handleDismiss = () => {
 		localStorage.setItem("chrome-extension-dismissed", "true")
+		localStorage.removeItem("chrome-extension-minimized")
 		setIsDismissed(true)
 	}
 
diff --git a/apps/web/components/header.tsx b/apps/web/components/header.tsx
index 0d3682c0..7e6c800a 100644
--- a/apps/web/components/header.tsx
+++ b/apps/web/components/header.tsx
@@ -22,6 +22,7 @@ import {
 import { DropdownMenuItem } from "@ui/components/dropdown-menu"
 import { DropdownMenu } from "@ui/components/dropdown-menu"
 import { Avatar, AvatarFallback, AvatarImage } from "@ui/components/avatar"
+import { Tooltip, TooltipContent, TooltipTrigger } from "@ui/components/tooltip"
 import { useAuth } from "@lib/auth-context"
 import { ConnectAIModal } from "./connect-ai-modal"
 import { useTheme } from "next-themes"
@@ -86,19 +87,31 @@ export function Header({ onAddMemory }: { onAddMemory?: () => void }) {
 							c
 						</span>
 					</Button>
-					<Button
-						variant="ghost"
-						size="sm"
-						onClick={() => setGraphModalOpen(true)}
-					>
-						<WaypointsIcon className="h-5 w-5" />
-						{/*<span className="hidden md:inline">Graph View</span>*/}
-					</Button>
+					<Tooltip>
+						<TooltipTrigger asChild>
+							<Button
+								variant="ghost"
+								size="sm"
+								onClick={() => setGraphModalOpen(true)}
+							>
+								<WaypointsIcon className="h-5 w-5" />
+							</Button>
+						</TooltipTrigger>
+						<TooltipContent>
+							<p>Graph View</p>
+						</TooltipContent>
+					</Tooltip>
 					<ConnectAIModal>
-						<Button variant="ghost" size="sm" className="gap-1.5">
-							<MCPIcon className="h-4 w-4" />
-							{/*<span className="hidden lg:inline">Connect to AI (MCP)</span>*/}
-						</Button>
+						<Tooltip>
+							<TooltipTrigger asChild>
+								<Button variant="ghost" size="sm" className="gap-1.5">
+									<MCPIcon className="h-4 w-4" />
+								</Button>
+							</TooltipTrigger>
+							<TooltipContent>
+								<p>Connect to AI (MCP)</p>
+							</TooltipContent>
+						</Tooltip>
 					</ConnectAIModal>
 					<DropdownMenu>
 						<DropdownMenuTrigger>
diff --git a/apps/web/components/model-selector.tsx b/apps/web/components/model-selector.tsx
new file mode 100644
index 00000000..38b68f9f
--- /dev/null
+++ b/apps/web/components/model-selector.tsx
@@ -0,0 +1,146 @@
+"use client"
+
+import { useState } from "react"
+import { Button } from "@repo/ui/components/button"
+import { ChevronDown } from "lucide-react"
+import { motion } from "motion/react"
+
+const models = [
+	{
+		id: "gpt-5",
+		name: "GPT 5",
+		description: "OpenAI's latest model",
+	},
+	{
+		id: "claude-sonnet-4.5",
+		name: "Claude Sonnet 4.5",
+		description: "Anthropic's advanced model",
+	},
+	{
+		id: "gemini-2.5-pro",
+		name: "Gemini 2.5 Pro",
+		description: "Google's most capable model",
+	},
+] as const
+
+type ModelId = (typeof models)[number]["id"]
+
+interface ModelSelectorProps {
+	selectedModel?: ModelId
+	onModelChange?: (modelId: ModelId) => void
+	disabled?: boolean
+}
+
+export function ModelSelector({
+	selectedModel = "gemini-2.5-pro",
+	onModelChange,
+	disabled = false,
+}: ModelSelectorProps) {
+	const [isOpen, setIsOpen] = useState(false)
+
+	const currentModel = models.find((m) => m.id === selectedModel) || models[0]
+
+	const handleModelSelect = (modelId: ModelId) => {
+		onModelChange?.(modelId)
+		setIsOpen(false)
+	}
+
+	return (
+		<div className="relative">
+			<Button
+				variant="ghost"
+				className="flex items-center gap-1.5 px-2 py-1.5 rounded-md transition-colors"
+				onClick={() => !disabled && setIsOpen(!isOpen)}
+				disabled={disabled}
+			>
+				<svg
+					xmlns="http://www.w3.org/2000/svg"
+					width="24"
+					height="24"
+					viewBox="0 0 24 24"
+					fill="none"
+				>
+					<title>Model Selector</title>
+					<g clip-path="url(#clip0_4418_9868)">
+						<path
+							d="M12.92 2.25984L19.43 5.76984C20.19 6.17984 20.19 7.34984 19.43 7.75984L12.92 11.2698C12.34 11.5798 11.66 11.5798 11.08 11.2698L4.57 7.75984C3.81 7.34984 3.81 6.17984 4.57 5.76984L11.08 2.25984C11.66 1.94984 12.34 1.94984 12.92 2.25984Z"
+							stroke="currentColor"
+							stroke-width="1.5"
+							stroke-linecap="round"
+							stroke-linejoin="round"
+						/>
+						<path
+							d="M3.61 10.1297L9.66 13.1597C10.41 13.5397 10.89 14.3097 10.89 15.1497V20.8697C10.89 21.6997 10.02 22.2297 9.28 21.8597L3.23 18.8297C2.48 18.4497 2 17.6797 2 16.8397V11.1197C2 10.2897 2.87 9.75968 3.61 10.1297Z"
+							stroke="currentColor"
+							stroke-width="1.5"
+							stroke-linecap="round"
+							stroke-linejoin="round"
+						/>
+						<path
+							d="M20.39 10.1297L14.34 13.1597C13.59 13.5397 13.11 14.3097 13.11 15.1497V20.8697C13.11 21.6997 13.98 22.2297 14.72 21.8597L20.77 18.8297C21.52 18.4497 22 17.6797 22 16.8397V11.1197C22 10.2897 21.13 9.75968 20.39 10.1297Z"
+							stroke="currentColor"
+							stroke-width="1.5"
+							stroke-linecap="round"
+							stroke-linejoin="round"
+						/>
+					</g>
+					<defs>
+						<clipPath id="clip0_4418_9868">
+							<rect width="24" height="24" fill="white" />
+						</clipPath>
+					</defs>
+				</svg>
+				<span className="text-xs font-medium max-w-32 truncate">
+					{currentModel.name}
+				</span>
+				<motion.div
+					animate={{ rotate: isOpen ? 180 : 0 }}
+					transition={{ duration: 0.25 }}
+				>
+					<ChevronDown className="h-3 w-3" />
+				</motion.div>
+			</Button>
+
+			{isOpen && (
+				<>
+					<button
+						type="button"
+						className="fixed inset-0 z-40"
+						onClick={() => setIsOpen(false)}
+						onKeyDown={(e) => e.key === "Escape" && setIsOpen(false)}
+						aria-label="Close model selector"
+					/>
+
+					<div className="absolute top-full left-0 mt-1 w-64 bg-background/95 backdrop-blur-xl border border-border rounded-md shadow-xl z-50 overflow-hidden space-y-1">
+						<div className="p-1.5 space-y-1">
+							{models.map((model) => (
+								<button
+									key={model.id}
+									type="button"
+									className={`flex items-center p-1 px-2 rounded-md transition-colors cursor-pointer w-full text-left ${
+										selectedModel === model.id
+											? "bg-accent"
+											: "hover:bg-accent/50"
+									}`}
+									onClick={() => handleModelSelect(model.id)}
+									onKeyDown={(e) =>
+										e.key === "Enter" && handleModelSelect(model.id)
+									}
+								>
+									<div className="flex-1 min-w-0">
+										<div className="text-sm font-medium text-foreground">
+											{model.name}
+										</div>
+										<div className="text-xs text-muted-foreground truncate">
+											{model.description}
+										</div>
+									</div>
+								</button>
+							))}
+						</div>
+					</div>
+				</>
+			)}
+		</div>
+	)
+}
diff --git a/apps/web/components/views/chat/chat-messages.tsx b/apps/web/components/views/chat/chat-messages.tsx
index db0ae209..2224005e 100644
--- a/apps/web/components/views/chat/chat-messages.tsx
+++ b/apps/web/components/views/chat/chat-messages.tsx
@@ -216,20 +216,32 @@ export function ChatMessages() {
 		getCurrentChat,
 	} = usePersistentChat()
 
+	const storageKey = `chat-model-${currentChatId}`
+
 	const [input, setInput] = useState("")
+	const [selectedModel, setSelectedModel] = useState<
+		"gpt-5" | "claude-sonnet-4.5" | "gemini-2.5-pro"
+	>((sessionStorage.getItem(storageKey) as "gpt-5" | "claude-sonnet-4.5" | "gemini-2.5-pro" || "gemini-2.5-pro") || "gemini-2.5-pro")
 	const activeChatIdRef = useRef<string | null>(null)
 	const shouldGenerateTitleRef = useRef<boolean>(false)
 	const hasRunInitialMessageRef = useRef<boolean>(false)
 
 	const { setDocumentIds } = useGraphHighlights()
 
+	console.log("selectedModel", selectedModel)
+
 	const { messages, sendMessage, status, stop, setMessages, id, regenerate } =
 		useChat({
 			id: currentChatId ?? undefined,
 			transport: new DefaultChatTransport({
 				api: `${process.env.NEXT_PUBLIC_BACKEND_URL}/chat`,
 				credentials: "include",
-				body: { metadata: { projectId: selectedProject } },
+				body: {
+					metadata: {
+						projectId: selectedProject,
+						model: selectedModel,
+					},
+				},
 			}),
 			maxSteps: 2,
 			onFinish: (result) => {
@@ -255,6 +267,23 @@ export function ChatMessages() {
 	}, [currentChatId, id])
 
 	useEffect(() => {
+		if (currentChatId) {
+			const savedModel = sessionStorage.getItem(storageKey) as
+				| "gpt-5"
+				| "claude-sonnet-4.5"
+				| "gemini-2.5-pro"
+
+			if (
+				savedModel &&
+				["gpt-5", "claude-sonnet-4.5", "gemini-2.5-pro"].includes(savedModel)
+			) {
+				console.log("savedModel", savedModel)
+				setSelectedModel(savedModel)
+			}
+		}
+	}, [currentChatId])
+
+	useEffect(() => {
 		if (currentChatId && !hasRunInitialMessageRef.current) {
 			// Check if there's an initial message from the home page in sessionStorage
 			const storageKey = `chat-initial-${currentChatId}`
diff --git a/apps/web/components/views/profile.tsx b/apps/web/components/views/profile.tsx
index 5aa21231..cd4e6085 100644
--- a/apps/web/components/views/profile.tsx
+++ b/apps/web/components/views/profile.tsx
@@ -301,7 +301,7 @@ export function ProfileView() {
 										</div>
 									</h4>
 									<ul className="space-y-2">
-										<li className="flex items-center gap-2 text-sm text-black/90">
+										<li className="flex items-center gap-2 text-xs sm:text-sm text-foreground">
 											<CheckCircle className="h-4 w-4 text-green-400" />
 											Unlimited memories
 										</li>
diff --git a/packages/ui/components/tooltip.tsx b/packages/ui/components/tooltip.tsx
index 32a78f06..daea283d 100644
--- a/packages/ui/components/tooltip.tsx
+++ b/packages/ui/components/tooltip.tsx
@@ -43,7 +43,7 @@ function TooltipContent({
 		<TooltipPrimitive.Portal>
 			<TooltipPrimitive.Content
 				className={cn(
-					"bg-primary text-primary-foreground animate-in fade-in-0 zoom-in-95 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 z-50 w-fit origin-(--radix-tooltip-content-transform-origin) rounded-md px-3 py-1.5 text-xs text-balance",
+					"bg-secondary text-secondary-foreground animate-in fade-in-0 zoom-in-95 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 z-50 w-fit origin-(--radix-tooltip-content-transform-origin) rounded-md px-3 py-1.5 text-xs text-balance",
 					className,
 				)}
 				data-slot="tooltip-content"
@@ -51,7 +51,7 @@ function TooltipContent({
 				{...props}
 			>
 				{children}
-				<TooltipPrimitive.Arrow className="bg-primary fill-primary z-50 size-2.5 translate-y-[calc(-50%_-_2px)] rotate-45 rounded-sm" />
+				<TooltipPrimitive.Arrow className="bg-secondary fill-secondary z-50 size-2.5 translate-y-[calc(-50%_-_2px)] rotate-45 rounded-sm" />
 			</TooltipPrimitive.Content>
 		</TooltipPrimitive.Portal>
 	);