Inky9/LFM2-WebGPU
0
1import React, {2 useState,3 useEffect,4 useCallback,5 useRef,6 useMemo,7} from "react";8import { openDB, type IDBPDatabase } from "idb";9import { Play, Plus, Zap, RotateCcw, Settings, X } from "lucide-react";10import { useLLM } from "./hooks/useLLM";11 12import type { Tool } from "./components/ToolItem";13 14import {15 parsePythonicCalls,16 extractPythonicCalls,17 extractFunctionAndRenderer,18 generateSchemaFromCode,19 extractToolCallContent,20 mapArgsToNamedParams,21 getErrorMessage,22 isMobileOrTablet,23} from "./utils";24 25import { DEFAULT_SYSTEM_PROMPT } from "./constants/systemPrompt";26import { DB_NAME, STORE_NAME, SETTINGS_STORE_NAME } from "./constants/db";27 28import { DEFAULT_TOOLS, TEMPLATE } from "./tools";29import ToolResultRenderer from "./components/ToolResultRenderer";30import ToolCallIndicator from "./components/ToolCallIndicator";31import ToolItem from "./components/ToolItem";32import ResultBlock from "./components/ResultBlock";33import ExamplePrompts from "./components/ExamplePrompts";34 35import { LoadingScreen } from "./components/LoadingScreen";36 37interface RenderInfo {38 call: string;39 result?: any;40 renderer?: string;41 input?: Record<string, any>;42 error?: string;43}44 45interface BaseMessage {46 role: "system" | "user" | "assistant";47 content: string;48}49interface ToolMessage {50 role: "tool";51 content: string;52 renderInfo: RenderInfo[]; // Rich data for the UI53}54type Message = BaseMessage | ToolMessage;55 56async function getDB(): Promise<IDBPDatabase> {57 return openDB(DB_NAME, 1, {58 upgrade(db) {59 if (!db.objectStoreNames.contains(STORE_NAME)) {60 db.createObjectStore(STORE_NAME, {61 keyPath: "id",62 autoIncrement: true,63 });64 }65 if (!db.objectStoreNames.contains(SETTINGS_STORE_NAME)) {66 db.createObjectStore(SETTINGS_STORE_NAME, { keyPath: "key" });67 }68 },69 });70}71 72const App: React.FC = () => {73 const [systemPrompt, setSystemPrompt] = useState<string>(74 DEFAULT_SYSTEM_PROMPT,75 );76 const [isSystemPromptModalOpen, setIsSystemPromptModalOpen] =77 useState<boolean>(false);78 const [tempSystemPrompt, setTempSystemPrompt] = useState<string>("");79 const [messages, setMessages] = useState<Message[]>([]);80 const [tools, setTools] = useState<Tool[]>([]);81 const [input, setInput] = useState<string>("");82 const [isGenerating, setIsGenerating] = useState<boolean>(false);83 const isMobile = useMemo(isMobileOrTablet, []);84 const [selectedModelId, setSelectedModelId] = useState<string>(85 isMobile ? "350M" : "1.2B",86 );87 const [isModelDropdownOpen, setIsModelDropdownOpen] =88 useState<boolean>(false);89 const chatContainerRef = useRef<HTMLDivElement>(null);90 const debounceTimers = useRef<Record<number, NodeJS.Timeout>>({});91 const toolsContainerRef = useRef<HTMLDivElement>(null);92 const inputRef = useRef<HTMLInputElement>(null);93 const {94 isLoading,95 isReady,96 error,97 progress,98 loadModel,99 generateResponse,100 clearPastKeyValues,101 } = useLLM(selectedModelId);102 103 const loadTools = useCallback(async (): Promise<void> => {104 const db = await getDB();105 const allTools: Tool[] = await db.getAll(STORE_NAME);106 if (allTools.length === 0) {107 const defaultTools: Tool[] = Object.entries(DEFAULT_TOOLS).map(108 ([name, code], id) => ({109 id,110 name,111 code,112 enabled: true,113 isCollapsed: false,114 }),115 );116 const tx = db.transaction(STORE_NAME, "readwrite");117 await Promise.all(defaultTools.map((tool) => tx.store.put(tool)));118 await tx.done;119 setTools(defaultTools);120 } else {121 setTools(allTools.map((t) => ({ ...t, isCollapsed: false })));122 }123 }, []);124 125 useEffect(() => {126 loadTools();127 }, [loadTools]);128 129 useEffect(() => {130 if (chatContainerRef.current) {131 chatContainerRef.current.scrollTop =132 chatContainerRef.current.scrollHeight;133 }134 }, [messages]);135 136 const updateToolInDB = async (tool: Tool): Promise<void> => {137 const db = await getDB();138 await db.put(STORE_NAME, tool);139 };140 141 const saveToolDebounced = (tool: Tool): void => {142 if (tool.id !== undefined && debounceTimers.current[tool.id]) {143 clearTimeout(debounceTimers.current[tool.id]);144 }145 if (tool.id !== undefined) {146 debounceTimers.current[tool.id] = setTimeout(() => {147 updateToolInDB(tool);148 }, 300);149 }150 };151 152 const clearChat = useCallback(() => {153 setMessages([]);154 clearPastKeyValues();155 }, [clearPastKeyValues]);156 157 const addTool = async (): Promise<void> => {158 const newTool: Omit<Tool, "id"> = {159 name: "new_tool",160 code: TEMPLATE,161 enabled: true,162 isCollapsed: false,163 };164 const db = await getDB();165 const id = await db.add(STORE_NAME, newTool);166 setTools((prev) => {167 const updated = [...prev, { ...newTool, id: id as number }];168 setTimeout(() => {169 if (toolsContainerRef.current) {170 toolsContainerRef.current.scrollTop =171 toolsContainerRef.current.scrollHeight;172 }173 }, 0);174 return updated;175 });176 clearChat();177 };178 179 const deleteTool = async (id: number): Promise<void> => {180 if (debounceTimers.current[id]) {181 clearTimeout(debounceTimers.current[id]);182 }183 const db = await getDB();184 await db.delete(STORE_NAME, id);185 setTools(tools.filter((tool) => tool.id !== id));186 clearChat();187 };188 189 const toggleToolEnabled = (id: number): void => {190 let changedTool: Tool | undefined;191 const newTools = tools.map((tool) => {192 if (tool.id === id) {193 changedTool = { ...tool, enabled: !tool.enabled };194 return changedTool;195 }196 return tool;197 });198 setTools(newTools);199 if (changedTool) saveToolDebounced(changedTool);200 };201 202 const toggleToolCollapsed = (id: number): void => {203 setTools(204 tools.map((tool) =>205 tool.id === id ? { ...tool, isCollapsed: !tool.isCollapsed } : tool,206 ),207 );208 };209 210 const expandTool = (id: number): void => {211 setTools(212 tools.map((tool) =>213 tool.id === id ? { ...tool, isCollapsed: false } : tool,214 ),215 );216 };217 218 const handleToolCodeChange = (id: number, newCode: string): void => {219 let changedTool: Tool | undefined;220 const newTools = tools.map((tool) => {221 if (tool.id === id) {222 const { functionCode } = extractFunctionAndRenderer(newCode);223 const schema = generateSchemaFromCode(functionCode);224 changedTool = { ...tool, code: newCode, name: schema.name };225 return changedTool;226 }227 return tool;228 });229 setTools(newTools);230 if (changedTool) saveToolDebounced(changedTool);231 };232 233 const executeToolCall = async (callString: string): Promise<string> => {234 const parsedCall = parsePythonicCalls(callString);235 if (!parsedCall) throw new Error(`Invalid tool call format: ${callString}`);236 237 const { name, positionalArgs, keywordArgs } = parsedCall;238 const toolToUse = tools.find((t) => t.name === name && t.enabled);239 if (!toolToUse) throw new Error(`Tool '${name}' not found or is disabled.`);240 241 const { functionCode } = extractFunctionAndRenderer(toolToUse.code);242 const schema = generateSchemaFromCode(functionCode);243 const paramNames = Object.keys(schema.parameters.properties);244 245 const finalArgs: any[] = [];246 const requiredParams = schema.parameters.required || [];247 248 for (let i = 0; i < paramNames.length; ++i) {249 const paramName = paramNames[i];250 if (i < positionalArgs.length) {251 finalArgs.push(positionalArgs[i]);252 } else if (keywordArgs.hasOwnProperty(paramName)) {253 finalArgs.push(keywordArgs[paramName]);254 } else if (255 schema.parameters.properties[paramName].hasOwnProperty("default")256 ) {257 finalArgs.push(schema.parameters.properties[paramName].default);258 } else if (!requiredParams.includes(paramName)) {259 finalArgs.push(undefined);260 } else {261 throw new Error(`Missing required argument: ${paramName}`);262 }263 }264 265 const bodyMatch = functionCode.match(/function[^{]+\{([\s\S]*)\}/);266 if (!bodyMatch) {267 throw new Error(268 "Could not parse function body. Ensure it's a standard `function` declaration.",269 );270 }271 const body = bodyMatch[1];272 const AsyncFunction = Object.getPrototypeOf(273 async function () {},274 ).constructor;275 const func = new AsyncFunction(...paramNames, body);276 const result = await func(...finalArgs);277 return JSON.stringify(result);278 };279 280 const executeToolCalls = async (281 toolCallContent: string,282 ): Promise<RenderInfo[]> => {283 const toolCalls = extractPythonicCalls(toolCallContent);284 if (toolCalls.length === 0)285 return [{ call: "", error: "No valid tool calls found." }];286 287 const results: RenderInfo[] = [];288 for (const call of toolCalls) {289 try {290 const result = await executeToolCall(call);291 const parsedCall = parsePythonicCalls(call);292 const toolUsed = parsedCall293 ? tools.find((t) => t.name === parsedCall.name && t.enabled)294 : null;295 const { rendererCode } = toolUsed296 ? extractFunctionAndRenderer(toolUsed.code)297 : { rendererCode: undefined };298 299 let parsedResult;300 try {301 parsedResult = JSON.parse(result);302 } catch {303 parsedResult = result;304 }305 306 let namedParams: Record<string, any> = Object.create(null);307 if (parsedCall && toolUsed) {308 const schema = generateSchemaFromCode(309 extractFunctionAndRenderer(toolUsed.code).functionCode,310 );311 const paramNames = Object.keys(schema.parameters.properties);312 namedParams = mapArgsToNamedParams(313 paramNames,314 parsedCall.positionalArgs,315 parsedCall.keywordArgs,316 );317 }318 319 results.push({320 call,321 result: parsedResult,322 renderer: rendererCode,323 input: namedParams,324 });325 } catch (error) {326 const errorMessage = getErrorMessage(error);327 results.push({ call, error: errorMessage });328 }329 }330 return results;331 };332 333 const handleSendMessage = async (): Promise<void> => {334 if (!input.trim() || !isReady) return;335 336 const userMessage: Message = { role: "user", content: input };337 let currentMessages: Message[] = [...messages, userMessage];338 setMessages(currentMessages);339 setInput("");340 setIsGenerating(true);341 342 try {343 const toolSchemas = tools344 .filter((tool) => tool.enabled)345 .map((tool) => generateSchemaFromCode(tool.code));346 347 while (true) {348 const messagesForGeneration = [349 { role: "system" as const, content: systemPrompt },350 ...currentMessages,351 ];352 353 setMessages([...currentMessages, { role: "assistant", content: "" }]);354 355 let accumulatedContent = "";356 const response = await generateResponse(357 messagesForGeneration,358 toolSchemas,359 (token: string) => {360 accumulatedContent += token;361 setMessages((current) => {362 const updated = [...current];363 updated[updated.length - 1] = {364 role: "assistant",365 content: accumulatedContent,366 };367 return updated;368 });369 },370 );371 372 currentMessages.push({ role: "assistant", content: response });373 const toolCallContent = extractToolCallContent(response);374 375 if (toolCallContent) {376 const toolResults = await executeToolCalls(toolCallContent);377 378 const toolMessage: ToolMessage = {379 role: "tool",380 content: JSON.stringify(toolResults.map((r) => r.result ?? null)),381 renderInfo: toolResults,382 };383 currentMessages.push(toolMessage);384 setMessages([...currentMessages]);385 continue;386 } else {387 setMessages(currentMessages);388 break;389 }390 }391 } catch (error) {392 const errorMessage = getErrorMessage(error);393 setMessages([394 ...currentMessages,395 {396 role: "assistant",397 content: `Error generating response: ${errorMessage}`,398 },399 ]);400 } finally {401 setIsGenerating(false);402 setTimeout(() => inputRef.current?.focus(), 0);403 }404 };405 406 const loadSystemPrompt = useCallback(async (): Promise<void> => {407 try {408 const db = await getDB();409 const stored = await db.get(SETTINGS_STORE_NAME, "systemPrompt");410 if (stored && stored.value) setSystemPrompt(stored.value);411 } catch (error) {412 console.error("Failed to load system prompt:", error);413 }414 }, []);415 416 const saveSystemPrompt = useCallback(417 async (prompt: string): Promise<void> => {418 try {419 const db = await getDB();420 await db.put(SETTINGS_STORE_NAME, {421 key: "systemPrompt",422 value: prompt,423 });424 } catch (error) {425 console.error("Failed to save system prompt:", error);426 }427 },428 [],429 );430 431 const loadSelectedModel = useCallback(async (): Promise<void> => {432 try {433 await loadModel();434 } catch (error) {435 console.error("Failed to load model:", error);436 }437 }, [selectedModelId, loadModel]);438 439 const loadSelectedModelId = useCallback(async (): Promise<void> => {440 try {441 const db = await getDB();442 const stored = await db.get(SETTINGS_STORE_NAME, "selectedModelId");443 if (stored && stored.value) {444 setSelectedModelId(stored.value);445 }446 } catch (error) {447 console.error("Failed to load selected model ID:", error);448 }449 }, []);450 451 useEffect(() => {452 loadSystemPrompt();453 }, [loadSystemPrompt]);454 455 const handleOpenSystemPromptModal = (): void => {456 setTempSystemPrompt(systemPrompt);457 setIsSystemPromptModalOpen(true);458 };459 460 const handleSaveSystemPrompt = (): void => {461 setSystemPrompt(tempSystemPrompt);462 saveSystemPrompt(tempSystemPrompt);463 setIsSystemPromptModalOpen(false);464 };465 466 const handleCancelSystemPrompt = (): void => {467 setTempSystemPrompt("");468 setIsSystemPromptModalOpen(false);469 };470 471 const handleResetSystemPrompt = (): void => {472 setTempSystemPrompt(DEFAULT_SYSTEM_PROMPT);473 };474 475 const saveSelectedModel = useCallback(476 async (modelId: string): Promise<void> => {477 try {478 const db = await getDB();479 await db.put(SETTINGS_STORE_NAME, {480 key: "selectedModelId",481 value: modelId,482 });483 } catch (error) {484 console.error("Failed to save selected model ID:", error);485 }486 },487 [],488 );489 490 useEffect(() => {491 loadSystemPrompt();492 loadSelectedModelId();493 }, [loadSystemPrompt, loadSelectedModelId]);494 495 const handleModelSelect = async (modelId: string) => {496 setSelectedModelId(modelId);497 setIsModelDropdownOpen(false);498 await saveSelectedModel(modelId);499 };500 501 const handleExampleClick = async (messageText: string): Promise<void> => {502 if (!isReady || isGenerating) return;503 setInput(messageText);504 505 const userMessage: Message = { role: "user", content: messageText };506 const currentMessages: Message[] = [...messages, userMessage];507 setMessages(currentMessages);508 setInput("");509 setIsGenerating(true);510 511 try {512 const toolSchemas = tools513 .filter((tool) => tool.enabled)514 .map((tool) => generateSchemaFromCode(tool.code));515 516 while (true) {517 const messagesForGeneration = [518 { role: "system" as const, content: systemPrompt },519 ...currentMessages,520 ];521 522 setMessages([...currentMessages, { role: "assistant", content: "" }]);523 524 let accumulatedContent = "";525 const response = await generateResponse(526 messagesForGeneration,527 toolSchemas,528 (token: string) => {529 accumulatedContent += token;530 setMessages((current) => {531 const updated = [...current];532 updated[updated.length - 1] = {533 role: "assistant",534 content: accumulatedContent,535 };536 return updated;537 });538 },539 );540 541 currentMessages.push({ role: "assistant", content: response });542 const toolCallContent = extractToolCallContent(response);543 544 if (toolCallContent) {545 const toolResults = await executeToolCalls(toolCallContent);546 547 const toolMessage: ToolMessage = {548 role: "tool",549 content: JSON.stringify(toolResults.map((r) => r.result ?? null)),550 renderInfo: toolResults,551 };552 currentMessages.push(toolMessage);553 setMessages([...currentMessages]);554 continue;555 } else {556 setMessages(currentMessages);557 break;558 }559 }560 } catch (error) {561 const errorMessage = getErrorMessage(error);562 setMessages([563 ...currentMessages,564 {565 role: "assistant",566 content: `Error generating response: ${errorMessage}`,567 },568 ]);569 } finally {570 setIsGenerating(false);571 setTimeout(() => inputRef.current?.focus(), 0);572 }573 };574 575 return (576 <div className="font-sans bg-gray-900">577 {!isReady ? (578 <LoadingScreen579 isLoading={isLoading}580 progress={progress}581 error={error}582 loadSelectedModel={loadSelectedModel}583 selectedModelId={selectedModelId}584 isModelDropdownOpen={isModelDropdownOpen}585 setIsModelDropdownOpen={setIsModelDropdownOpen}586 handleModelSelect={handleModelSelect}587 />588 ) : (589 <div className="flex h-screen text-white">590 <div className="w-1/2 flex flex-col p-4">591 <div className="flex items-center justify-between mb-4">592 <div className="flex items-center gap-3">593 <h1 className="text-3xl font-bold text-gray-200">594 LFM2 WebGPU595 </h1>596 </div>597 <div className="flex items-center gap-3">598 <div className="flex items-center text-green-400">599 <Zap size={16} className="mr-2" />600 Ready601 </div>602 <button603 disabled={isGenerating}604 onClick={clearChat}605 className={`h-10 flex items-center px-3 py-2 rounded-lg font-bold transition-colors text-sm ${606 isGenerating607 ? "bg-gray-600 cursor-not-allowed opacity-50"608 : "bg-gray-600 hover:bg-gray-700"609 }`}610 title="Clear chat"611 >612 <RotateCcw size={14} className="mr-2" /> Clear613 </button>614 <button615 onClick={handleOpenSystemPromptModal}616 className="h-10 flex items-center px-3 py-2 rounded-lg font-bold transition-colors bg-gray-600 hover:bg-gray-700 text-sm"617 title="Edit system prompt"618 >619 <Settings size={16} />620 </button>621 </div>622 </div>623 624 <div625 ref={chatContainerRef}626 className="flex-grow bg-gray-800 rounded-lg p-4 overflow-y-auto mb-4 space-y-4"627 >628 {messages.length === 0 && isReady ? (629 <ExamplePrompts onExampleClick={handleExampleClick} />630 ) : (631 messages.map((msg, index) => {632 const key = `${msg.role}-${index}`;633 634 if (msg.role === "user") {635 return (636 <div key={key} className="flex justify-end">637 <div className="p-3 rounded-lg max-w-md bg-indigo-600">638 <p className="text-sm whitespace-pre-wrap">639 {msg.content}640 </p>641 </div>642 </div>643 );644 } else if (msg.role === "assistant") {645 const isToolCall = msg.content.includes(646 "<|tool_call_start|>",647 );648 649 if (isToolCall) {650 const nextMessage = messages[index + 1];651 const isCompleted = nextMessage?.role === "tool";652 const hasError =653 isCompleted &&654 (nextMessage as ToolMessage).renderInfo.some(655 (info) => !!info.error,656 );657 658 return (659 <div key={key} className="flex justify-start">660 <div className="p-3 rounded-lg bg-gray-700">661 <ToolCallIndicator662 content={msg.content}663 isRunning={!isCompleted}664 hasError={hasError}665 />666 </div>667 </div>668 );669 }670 671 return (672 <div key={key} className="flex justify-start">673 <div className="p-3 rounded-lg max-w-md bg-gray-700">674 <p className="text-sm whitespace-pre-wrap">675 {msg.content}676 </p>677 </div>678 </div>679 );680 } else if (msg.role === "tool") {681 const visibleToolResults = msg.renderInfo.filter(682 (info) =>683 info.error || (info.result != null && info.renderer),684 );685 686 if (visibleToolResults.length === 0) return null;687 688 return (689 <div key={key} className="flex justify-start">690 <div className="p-3 rounded-lg bg-gray-700 max-w-lg">691 <div className="space-y-3">692 {visibleToolResults.map((info, idx) => (693 <div className="flex flex-col gap-2" key={idx}>694 <div className="text-xs text-gray-400 font-mono">695 {info.call}696 </div>697 {info.error ? (698 <ResultBlock error={info.error} />699 ) : (700 <ToolResultRenderer701 result={info.result}702 rendererCode={info.renderer}703 input={info.input}704 />705 )}706 </div>707 ))}708 </div>709 </div>710 </div>711 );712 }713 return null;714 })715 )}716 </div>717 718 <div className="flex">719 <input720 ref={inputRef}721 type="text"722 value={input}723 onChange={(e) => setInput(e.target.value)}724 onKeyDown={(e) =>725 e.key === "Enter" &&726 !isGenerating &&727 isReady &&728 handleSendMessage()729 }730 disabled={isGenerating || !isReady}731 className="flex-grow bg-gray-700 rounded-l-lg p-3 focus:outline-none focus:ring-2 focus:ring-indigo-500 disabled:opacity-50"732 placeholder={733 isReady734 ? "Type your message here..."735 : "Load model first to enable chat"736 }737 />738 <button739 onClick={handleSendMessage}740 disabled={isGenerating || !isReady}741 className="bg-indigo-600 hover:bg-indigo-700 disabled:bg-gray-600 disabled:cursor-not-allowed text-white font-bold p-3 rounded-r-lg transition-colors"742 >743 <Play size={20} />744 </button>745 </div>746 </div>747 748 <div className="w-1/2 flex flex-col p-4 border-l border-gray-700">749 <div className="flex justify-between items-center mb-4">750 <h2 className="text-2xl font-bold text-teal-400">Tools</h2>751 <button752 onClick={addTool}753 className="flex items-center bg-teal-600 hover:bg-teal-700 text-white font-bold py-2 px-4 rounded-lg transition-colors"754 >755 <Plus size={16} className="mr-2" /> Add Tool756 </button>757 </div>758 <div759 ref={toolsContainerRef}760 className="flex-grow bg-gray-800 rounded-lg p-4 overflow-y-auto space-y-3"761 >762 {tools.map((tool) => (763 <ToolItem764 key={tool.id}765 tool={tool}766 onToggleEnabled={() => toggleToolEnabled(tool.id)}767 onToggleCollapsed={() => toggleToolCollapsed(tool.id)}768 onExpand={() => expandTool(tool.id)}769 onDelete={() => deleteTool(tool.id)}770 onCodeChange={(newCode) =>771 handleToolCodeChange(tool.id, newCode)772 }773 />774 ))}775 </div>776 </div>777 </div>778 )}779 780 {isSystemPromptModalOpen && (781 <div className="fixed inset-0 bg-black bg-opacity-50 flex items-center justify-center z-50">782 <div className="bg-gray-800 rounded-lg p-6 w-3/4 max-w-4xl max-h-3/4 flex flex-col text-gray-100">783 <div className="flex justify-between items-center mb-4">784 <h2 className="text-xl font-bold text-indigo-400">785 Edit System Prompt786 </h2>787 <button788 onClick={handleCancelSystemPrompt}789 className="text-gray-400 hover:text-white"790 >791 <X size={20} />792 </button>793 </div>794 <div className="flex-grow mb-4">795 <textarea796 value={tempSystemPrompt}797 onChange={(e) => setTempSystemPrompt(e.target.value)}798 className="w-full h-full bg-gray-700 text-white p-4 rounded-lg resize-none focus:outline-none focus:ring-2 focus:ring-indigo-500"799 placeholder="Enter your system prompt here..."800 style={{ minHeight: "300px" }}801 />802 </div>803 <div className="flex justify-between">804 <button805 onClick={handleResetSystemPrompt}806 className="px-4 py-2 bg-teal-600 hover:bg-teal-700 rounded-lg transition-colors"807 >808 Reset809 </button>810 <div className="flex gap-3">811 <button812 onClick={handleSaveSystemPrompt}813 className="px-4 py-2 bg-indigo-600 hover:bg-indigo-700 rounded-lg transition-colors"814 >815 Save816 </button>817 </div>818 </div>819 </div>820 </div>821 )}822 </div>823 );824};825 826export default App;827 