From f81d6bde18e2425d66ee62bd99df5eca93d576a8 Mon Sep 17 00:00:00 2001 From: Kayshen-X Date: Wed, 18 Feb 2026 22:35:16 +0800 Subject: [PATCH] feat(ai): add model selection and fetching capabilities Enhance AI chat and generation functionalities by introducing model selection. Update the chat and generation endpoints to accept an optional model parameter. Implement a new API endpoint to fetch available AI models from the server, and integrate model selection into the AI chat panel for improved user experience. --- server/api/ai/chat.ts | 13 ++-- server/api/ai/generate.ts | 13 ++-- server/api/ai/models.ts | 47 +++++++++++++ src/components/panels/ai-chat-panel.tsx | 88 ++++++++++++++++++++++--- src/services/ai/ai-service.ts | 22 ++++++- src/stores/ai-store.ts | 20 +++++- 6 files changed, 179 insertions(+), 24 deletions(-) create mode 100644 server/api/ai/models.ts diff --git a/server/api/ai/chat.ts b/server/api/ai/chat.ts index 1fdd3fd70..77123e1c7 100644 --- a/server/api/ai/chat.ts +++ b/server/api/ai/chat.ts @@ -3,6 +3,7 @@ import { defineEventHandler, readBody, setResponseHeaders } from 'h3' interface ChatBody { system: string messages: Array<{ role: 'user' | 'assistant'; content: string }> + model?: string } /** @@ -27,16 +28,16 @@ export default defineEventHandler(async (event) => { const apiKey = process.env.ANTHROPIC_API_KEY if (apiKey) { try { - return await streamViaAnthropicSDK(apiKey, body) + return await streamViaAnthropicSDK(apiKey, body, body.model) } catch { // SDK not installed or failed — fall back to Agent SDK } } - return streamViaAgentSDK(body) + return streamViaAgentSDK(body, body.model) }) /** Stream via Anthropic SDK (when API key is available) */ -async function streamViaAnthropicSDK(apiKey: string, body: ChatBody) { +async function streamViaAnthropicSDK(apiKey: string, body: ChatBody, model?: string) { // @ts-expect-error — optional dependency, only used when ANTHROPIC_API_KEY is set const { default: Anthropic } = await import('@anthropic-ai/sdk') const client = new Anthropic({ apiKey }) @@ -46,7 +47,7 @@ async function streamViaAnthropicSDK(apiKey: string, body: ChatBody) { const encoder = new TextEncoder() try { const messageStream = client.messages.stream({ - model: 'claude-sonnet-4-5-20250929', + model: model || 'claude-sonnet-4-5-20250929', max_tokens: 16384, system: body.system, messages: body.messages, @@ -80,7 +81,7 @@ async function streamViaAnthropicSDK(apiKey: string, body: ChatBody) { } /** Stream via Claude Agent SDK (uses local Claude Code OAuth login, no API key needed) */ -function streamViaAgentSDK(body: ChatBody) { +function streamViaAgentSDK(body: ChatBody, model?: string) { const stream = new ReadableStream({ async start(controller) { const encoder = new TextEncoder() @@ -100,7 +101,7 @@ function streamViaAgentSDK(body: ChatBody) { prompt, options: { systemPrompt: body.system, - model: 'claude-sonnet-4-6', + model: model || 'claude-sonnet-4-6', maxTurns: 1, includePartialMessages: true, tools: [], diff --git a/server/api/ai/generate.ts b/server/api/ai/generate.ts index df97c97b3..7a502e423 100644 --- a/server/api/ai/generate.ts +++ b/server/api/ai/generate.ts @@ -3,6 +3,7 @@ import { defineEventHandler, readBody, setResponseHeaders } from 'h3' interface GenerateBody { system: string message: string + model?: string } /** @@ -21,22 +22,22 @@ export default defineEventHandler(async (event) => { const apiKey = process.env.ANTHROPIC_API_KEY if (apiKey) { try { - return await generateViaAnthropicSDK(apiKey, body) + return await generateViaAnthropicSDK(apiKey, body, body.model) } catch { // SDK not installed or failed — fall back to Agent SDK } } - return generateViaAgentSDK(body) + return generateViaAgentSDK(body, body.model) }) /** Generate via Anthropic SDK */ -async function generateViaAnthropicSDK(apiKey: string, body: GenerateBody) { +async function generateViaAnthropicSDK(apiKey: string, body: GenerateBody, model?: string) { try { // @ts-expect-error — optional dependency, only used when ANTHROPIC_API_KEY is set const { default: Anthropic } = await import('@anthropic-ai/sdk') const client = new Anthropic({ apiKey }) const response = await client.messages.create({ - model: 'claude-sonnet-4-5-20250929', + model: model || 'claude-sonnet-4-5-20250929', max_tokens: 4096, system: body.system, messages: [{ role: 'user', content: body.message }], @@ -51,7 +52,7 @@ async function generateViaAnthropicSDK(apiKey: string, body: GenerateBody) { } /** Generate via Claude Agent SDK (uses local Claude Code OAuth login, no API key needed) */ -async function generateViaAgentSDK(body: GenerateBody): Promise<{ text?: string; error?: string }> { +async function generateViaAgentSDK(body: GenerateBody, model?: string): Promise<{ text?: string; error?: string }> { try { const { query } = await import('@anthropic-ai/claude-agent-sdk') @@ -63,7 +64,7 @@ async function generateViaAgentSDK(body: GenerateBody): Promise<{ text?: string; prompt: body.message, options: { systemPrompt: body.system, - model: 'claude-sonnet-4-6', + model: model || 'claude-sonnet-4-6', maxTurns: 1, tools: [], permissionMode: 'plan', diff --git a/server/api/ai/models.ts b/server/api/ai/models.ts new file mode 100644 index 000000000..b6ff0bb3e --- /dev/null +++ b/server/api/ai/models.ts @@ -0,0 +1,47 @@ +import { defineEventHandler } from 'h3' + +interface ModelInfo { + value: string + displayName: string + description: string +} + +let cachedModels: ModelInfo[] | null = null + +/** + * Returns the list of available AI models. + * Fetches from Claude Agent SDK and caches the result. + */ +export default defineEventHandler(async () => { + if (cachedModels) { + return { models: cachedModels } + } + + try { + const { query } = await import('@anthropic-ai/claude-agent-sdk') + + const env = { ...process.env } as Record + delete env.CLAUDECODE + + const q = query({ + prompt: '', + options: { + model: 'claude-sonnet-4-6', + maxTurns: 1, + tools: [], + permissionMode: 'plan', + persistSession: false, + env, + }, + }) + + const models = await q.supportedModels() + cachedModels = models + q.close() + + return { models } + } catch (error) { + const message = error instanceof Error ? error.message : 'Unknown error' + return { models: [], error: message } + } +}) diff --git a/src/components/panels/ai-chat-panel.tsx b/src/components/panels/ai-chat-panel.tsx index 633e2dbf9..b06e73719 100644 --- a/src/components/panels/ai-chat-panel.tsx +++ b/src/components/panels/ai-chat-panel.tsx @@ -1,5 +1,5 @@ import { useState, useRef, useEffect, useCallback } from 'react' -import { Send, Trash2, Minus, Maximize2, Sparkles } from 'lucide-react' +import { Send, Trash2, Minus, Maximize2, Sparkles, ChevronDown } from 'lucide-react' import { nanoid } from 'nanoid' import { cn } from '@/lib/utils' import { Button } from '@/components/ui/button' @@ -7,7 +7,7 @@ import { useAIStore } from '@/stores/ai-store' import type { PanelCorner } from '@/stores/ai-store' import { useCanvasStore } from '@/stores/canvas-store' import { useDocumentStore } from '@/stores/document-store' -import { streamChat } from '@/services/ai/ai-service' +import { streamChat, fetchAvailableModels } from '@/services/ai/ai-service' import { CHAT_SYSTEM_PROMPT, DESIGN_GENERATOR_PROMPT, @@ -90,6 +90,7 @@ function useChatHandlers() { const [input, setInput] = useState('') const messages = useAIStore((s) => s.messages) const isStreaming = useAIStore((s) => s.isStreaming) + const model = useAIStore((s) => s.model) const addMessage = useAIStore((s) => s.addMessage) const updateLastMessage = useAIStore((s) => s.updateLastMessage) const setStreaming = useAIStore((s) => s.setStreaming) @@ -138,7 +139,7 @@ function useChatHandlers() { let accumulated = '' try { - for await (const chunk of streamChat(effectivePrompt, chatHistory)) { + for await (const chunk of streamChat(effectivePrompt, chatHistory, model)) { if (chunk.type === 'text') { accumulated += chunk.content updateLastMessage(accumulated) @@ -165,7 +166,7 @@ function useChatHandlers() { updateLastMessage(accumulated) try { - for await (const chunk of streamChat(effectivePrompt, retryHistory)) { + for await (const chunk of streamChat(effectivePrompt, retryHistory, model)) { if (chunk.type === 'text') { accumulated += chunk.content updateLastMessage(accumulated) @@ -192,7 +193,7 @@ function useChatHandlers() { return { messages: msgs } }) }, - [input, isStreaming, messages, addMessage, updateLastMessage, setStreaming], + [input, isStreaming, model, messages, addMessage, updateLastMessage, setStreaming], ) return { input, setInput, handleSend, isStreaming } @@ -239,12 +240,48 @@ export default function AIChatPanel() { const isMinimized = useAIStore((s) => s.isMinimized) const setPanelCorner = useAIStore((s) => s.setPanelCorner) const toggleMinimize = useAIStore((s) => s.toggleMinimize) + const model = useAIStore((s) => s.model) + const setModel = useAIStore((s) => s.setModel) + const availableModels = useAIStore((s) => s.availableModels) + const setAvailableModels = useAIStore((s) => s.setAvailableModels) + const isLoadingModels = useAIStore((s) => s.isLoadingModels) + const setLoadingModels = useAIStore((s) => s.setLoadingModels) + const [modelDropdownOpen, setModelDropdownOpen] = useState(false) const { input, setInput, handleSend } = useChatHandlers() useEffect(() => { messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }) }, [messages]) + // Fetch available models on mount + useEffect(() => { + if (availableModels.length > 0) return + setLoadingModels(true) + fetchAvailableModels().then((models) => { + if (models.length > 0) { + setAvailableModels(models) + // Set default model if current model is not in the list + if (!models.some((m) => m.value === model)) { + setModel(models[0].value) + } + } + setLoadingModels(false) + }) + }, []) // eslint-disable-line react-hooks/exhaustive-deps + + // Close model dropdown when clicking outside + useEffect(() => { + if (!modelDropdownOpen) return + const handler = (e: MouseEvent) => { + const panel = panelRef.current + if (panel && !panel.contains(e.target as Node)) { + setModelDropdownOpen(false) + } + } + document.addEventListener('pointerdown', handler) + return () => document.removeEventListener('pointerdown', handler) + }, [modelDropdownOpen]) + // Auto-expand when streaming starts while minimized useEffect(() => { if (isStreaming && isMinimized) { @@ -255,7 +292,7 @@ export default function AIChatPanel() { /* --- Drag-to-snap handlers --- */ const handleDragStart = useCallback((e: React.PointerEvent) => { - if ((e.target as HTMLElement).closest('button, input, textarea')) return + if ((e.target as HTMLElement).closest('button, input, textarea, select')) return const panel = panelRef.current if (!panel) return @@ -368,9 +405,42 @@ export default function AIChatPanel() { >
- - AI Assistant - + {/* Model selector */} +
+ + {modelDropdownOpen && availableModels.length > 0 && ( +
+ {availableModels.map((m) => ( + + ))} +
+ )} +
{messages.length > 0 && ( diff --git a/src/services/ai/ai-service.ts b/src/services/ai/ai-service.ts index ffa2f6466..01fa05eb4 100644 --- a/src/services/ai/ai-service.ts +++ b/src/services/ai/ai-service.ts @@ -1,4 +1,5 @@ import type { AIStreamChunk } from './ai-types' +import type { AIModelInfo } from '@/stores/ai-store' /** * Streams a chat response from the server-side AI endpoint. @@ -7,12 +8,13 @@ import type { AIStreamChunk } from './ai-types' export async function* streamChat( systemPrompt: string, messages: Array<{ role: 'user' | 'assistant'; content: string }>, + model?: string, ): AsyncGenerator { try { const response = await fetch('/api/ai/chat', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ system: systemPrompt, messages }), + body: JSON.stringify({ system: systemPrompt, messages, model }), }) if (!response.ok) { @@ -79,11 +81,12 @@ export async function* streamChat( export async function generateCompletion( systemPrompt: string, userMessage: string, + model?: string, ): Promise { const response = await fetch('/api/ai/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ system: systemPrompt, message: userMessage }), + body: JSON.stringify({ system: systemPrompt, message: userMessage, model }), }) if (!response.ok) { @@ -96,3 +99,18 @@ export async function generateCompletion( } return data.text ?? '' } + +/** + * Fetches available AI models from the server. + * The server queries Claude Agent SDK for the supported model list. + */ +export async function fetchAvailableModels(): Promise { + try { + const response = await fetch('/api/ai/models') + if (!response.ok) return [] + const data = await response.json() + return data.models ?? [] + } catch { + return [] + } +} diff --git a/src/stores/ai-store.ts b/src/stores/ai-store.ts index 87e2a7df9..58605df31 100644 --- a/src/stores/ai-store.ts +++ b/src/stores/ai-store.ts @@ -3,6 +3,12 @@ import type { ChatMessage } from '@/services/ai/ai-types' export type PanelCorner = 'top-left' | 'top-right' | 'bottom-left' | 'bottom-right' +export interface AIModelInfo { + value: string + displayName: string + description: string +} + interface AIState { messages: ChatMessage[] isStreaming: boolean @@ -10,9 +16,15 @@ interface AIState { activeTab: 'chat' | 'code' generatedCode: string codeFormat: 'react-tailwind' | 'html-css' | 'react-inline' + model: string + availableModels: AIModelInfo[] + isLoadingModels: boolean panelCorner: PanelCorner isMinimized: boolean + setModel: (model: string) => void + setAvailableModels: (models: AIModelInfo[]) => void + setLoadingModels: (v: boolean) => void addMessage: (msg: ChatMessage) => void updateLastMessage: (content: string) => void setStreaming: (v: boolean) => void @@ -33,7 +45,10 @@ export const useAIStore = create((set) => ({ activeTab: 'chat', generatedCode: '', codeFormat: 'react-tailwind', - panelCorner: 'bottom-right', + model: 'claude-sonnet-4-5-20250929', + availableModels: [], + isLoadingModels: false, + panelCorner: 'bottom-left', isMinimized: false, addMessage: (msg) => @@ -61,6 +76,9 @@ export const useAIStore = create((set) => ({ setCodeFormat: (codeFormat) => set({ codeFormat }), + setModel: (model) => set({ model }), + setAvailableModels: (availableModels) => set({ availableModels }), + setLoadingModels: (isLoadingModels) => set({ isLoadingModels }), clearMessages: () => set({ messages: [] }), setPanelCorner: (panelCorner) => set({ panelCorner }),